Ускорение работы моделей Stable Diffusion на процессорах Intel

Недавно мы рассказывали о последнем поколении процессоров Intel Xeon (кодовое название Sapphire Rapids). Мы говорили об их новых аппаратных возможностях, ориентированных на ускорение задач глубокого обучения,  разбирались с тем, как использовать их для ускорения распределённого дообучения трансформеров, занимающихся обработкой естественного языка, как применять их для ускорения работы таких моделей.

В этом материале мы собираемся остановиться на различных подходах к ускорению моделей Stable Diffusion на процессорах Sapphire Rapids. В следующем похожем посте речь пойдёт о распределённом дообучении.

Во время написания этой статьи легче всего попробовать в деле сервер с процессором Sapphire Rapids можно, воспользовавшись инстансом Amazon EC2 семейства R7iz. Так как эти системы всё ещё находятся в статусе Preview, вам, чтобы получить к ним доступ, нужно зарегистрироваться. Я, как и в предыдущих материалах, пользуюсь инстансом r7iz.metal-16xl (64 vCPU, 512GB RAM) с установленной Ubuntu 20.04 AMI (ami-07cd3e6c4915b2d18).

Код к статье можно найти на GitLab.

Библиотека Diffusers

Библиотека Diffusers очень сильно упрощает генерирование изображений с помощью моделей Stable Diffusion. Если вы не знакомы с этими моделями — вот иллюстрированное объяснение принципов их работы.

Для начала создадим виртуальное окружение, содержащее необходимые библиотеки: Transformers, Diffusers, Accelerate и PyTorch.

virtualenv sd_inference
source sd_inference/bin/activate
pip install pip --upgrade
pip install transformers diffusers accelerate torch==1.13.1

Теперь напишем простую функцию для измерения производительности, которая по много раз запускает процесс генерирования изображения и возвращает среднее время выполнения одного действия (среднее время ожидания результата).

import time

def elapsed_time(pipeline, prompt, nb_pass=10, num_inference_steps=20):
    # разогрев
    images = pipeline(prompt, num_inference_steps=10).images
    start = time.time()
    for _ in range(nb_pass):
        _ = pipeline(prompt, num_inference_steps=num_inference_steps, output_type="np")
    end = time.time()
    return (end - start) / nb_pass

Далее — соберём StableDiffusionPipeline с типом данных float32, используемым по умолчанию, и измерим скорость генерирования изображений.

from diffusers import StableDiffusionPipeline

model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id)
prompt = "sailing ship in storm by Rembrandt"
latency = elapsed_time(pipe, prompt)
print(latency)

Среднее время ожидания результата составило 32,3 секунды. Как показано в разделе Intel на Hugging Face, тот же код, запускаемый на процессорах Intel Xeon предыдущего поколения (кодовое название Ice Lake), выдаёт результат примерно через 45 секунд.

Без применения каких-то особых приёмов процессоры Sapphire Rapids, на точно таком же коде, оказываются значительно быстрее процессоров Ice Lake.

А теперь давайте их ускорим!

Optimum Intel и OpenVINO

Пакет Optimum Intel ускоряет все этапы ИИ-конвейеров на архитектурах Intel. Его API очень похож на обычный API библиотеки Diffusers. Это до крайности облегчает адаптацию существующего кода под него.

Optimum Intel поддерживает OpenVINO — опенсорсный набор инструментов Intel, направленный на организацию высокопроизводительной работы моделей глубокого обучения.

Optimum Intel и OpenVINO можно установить так:

pip install optimum[openvino]

Если переделывать под Optimum Intel предыдущий фрагмент кода — достаточно будет заменить StableDiffusionPipeline на OVStableDiffusionPipeline. Чтобы, что называется, «на лету», загрузить модель PyTorch и конвертировать её в формат OpenVINO, можно, при загрузке модели, установить флаг export=True.

from optimum.intel.openvino import OVStableDiffusionPipeline
...
ov_pipe = OVStableDiffusionPipeline.from_pretrained(model_id, export=True)
latency = elapsed_time(ov_pipe, prompt)
print(latency)

# Не забудьте сохранить экспортированную модель
ov_pipe.save_pretrained("./openvino")

OpenVINO автоматически оптимизирует модель в расчёте на формат bfloat16. Благодаря этому среднее время генерирования одного изображения теперь составляет 16,7 секунд. Мы получили очень приятное двукратное ускорение кода.

Вышеприведённый конвейер поддерживает динамические входные данные, не накладывает ограничений ни на количество изображений, ни на их разрешение. При применении Stable Diffusion приложение обычно ограничено одним (или несколькими) различными выходными разрешениями, такими, как 512x512 или 256x256. Получается, что имеет смысл попытка значительного ускорения кода путём переделывания конвейера под фиксированное разрешение. Если вы нуждаетесь в более чем одном выходном разрешении, вы можете просто поддерживать несколько экземпляров конвейера — по одному для каждого необходимого разрешения.

ov_pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
latency = elapsed_time(ov_pipe, prompt)

Применение статического разрешения ведёт к резкому сокращению среднего времени генерирования изображения — до 4,7 секунд. А это — дополнительное 3,5-кратное ускорение.

Как видите, OpenVINO — это простой и эффективный инструмент для ускорения работы моделей Stable Diffusion. Если применить этот инструмент в коде, запускаемом на процессорах Sapphire Rapids, то, в сравнении с работой обычного кода на процессорах Xeon семейства Ice Lake, получается почти 10-кратное ускорение.

Если вы не можете или не хотите пользоваться инструментами OpenVINO, тогда вам могут пригодиться другие техники оптимизации, о которых мы поговорим ниже. Пристегните ремни!

Оптимизация системного уровня

Модели Diffusers — это большие структуры, размеры которых исчисляются многими гигабайтами. Генерирование изображений — это операция, предусматривающая интенсивное использование памяти. Установка высокопроизводительной библиотеки для выделения памяти должна способствовать ускорению операций по работе с памятью, должна помогать параллельному выполнению таких операций несколькими ядрами процессора Xeon. Прошу обратить внимание на то, что подобная операция изменит стандартную библиотеку выделения памяти в системе. Конечно, можно вернуться к библиотеке, используемой по умолчанию, удалив новую библиотеку.

Одинаковый интерес для нас представляют библиотеки jemalloc и tcmalloc. Тут я показываю установку библиотеки jemalloc, так как она способствует несколько большему улучшению производительности моих тестов. Эту библиотеку, кроме того, можно подстроить в расчёте на определённую рабочую нагрузку. Например — для того, чтобы максимизировать использование процессора. Подробности можно найти в этом руководстве по настройке jemalloc.

sudo apt-get install -y libjemalloc-dev
export LD_PRELOAD=$LD_PRELOAD:/usr/lib/x86_64-linux-gnu/libjemalloc.so
export MALLOC_CONF="oversize_threshold:1,background_thread:true,metadata_thp:auto,dirty_decay_ms: 60000,muzzy_decay_ms:60000"

Далее — установим библиотеку libiomp для оптимизации параллельной обработки данных. Она является частью библиотеки Intel OpenMP* Runtime.

sudo apt-get install intel-mkl
export LD_PRELOAD=$LD_PRELOAD:/usr/lib/x86_64-linux-gnu/libiomp5.so
export OMP_NUM_THREADS=32

И наконец — устанавливаем инструмент командной строки numactl. Он позволяет прикреплять Python-процессы к определённым ядрам и избегать некоторой части избыточной траты ресурсов, связанной с переключением контекста.

numactl -C 0-31 python sd_blog_1.py

Благодаря этим оптимизациям изначальный код, использующий Diffusers, теперь показывает результат в 11,8 секунд. Это — почти в 3 раза быстрее, и достигается это без каких-либо изменений кода. Эти инструменты, определённо, отлично работают на нашем 32-ядерном Xeon.

Но мы ещё далеко не закончили. Добавим в наш арсенал пакет Intel Extension for PyTorch.

IPEX и BF16

Пакет Intel Extension for Pytorch (IPEX) расширяет PyTorch и пользуется возможностями аппаратного ускорения, которые имеются в процессорах Intel. Например — это AVX-512 Vector Neural Network Instructions (AVX512 VNNI) и Advanced Matrix Extensions (AMX).

Установим пакет:

pip install intel_extension_for_pytorch==1.13.100

Затем отредактируем код, оптимизировав все элементы конвейера с использованием IPEX (получить их список можно, выведя на печать объект pipe). Для этого потребуется конвертировать их в формат torch.channels_last.

import torch
import intel_extension_for_pytorch as ipex
...
pipe = StableDiffusionPipeline.from_pretrained(model_id)

# преобразование в формат torch.channels_last
pipe.unet = pipe.unet.to(memory_format=torch.channels_last)
pipe.vae = pipe.vae.to(memory_format=torch.channels_last)
pipe.text_encoder = pipe.text_encoder.to(memory_format=torch.channels_last)
pipe.safety_checker = pipe.safety_checker.to(memory_format=torch.channels_last)

# Создание случайных входных данных для включения JIT-компиляции
sample = torch.randn(2,4,64,64)
timestep = torch.rand(1)*999
encoder_hidden_status = torch.randn(2,77,768)
input_example = (sample, timestep, encoder_hidden_status)

# Оптимизация с использованием IPEX
pipe.unet = ipex.optimize(pipe.unet.eval(), dtype=torch.bfloat16, inplace=True, sample_input=input_example)
pipe.vae = ipex.optimize(pipe.vae.eval(), dtype=torch.bfloat16, inplace=True)
pipe.text_encoder = ipex.optimize(pipe.text_encoder.eval(), dtype=torch.bfloat16, inplace=True)
pipe.safety_checker = ipex.optimize(pipe.safety_checker.eval(), dtype=torch.bfloat16, inplace=True)

Мы, кроме того, включили использование формата данных bloat16 для того чтобы задействовать ускоритель AMX Tile Matrix Multiply Unit (TMMU) в процессорах Sapphire Rapids.

with torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16):
    latency = elapsed_time(pipe, prompt)
    print(latency)

Эта обновлённая версия кода демонстрирует дальнейшее снижение среднего времени ожидания — с 11,9 секунд до 5,4 секунд. Это значит, что IPEX и AMX дали нам двукратное ускорение.

Можно ли выжать из процессоров Sapphire Rapids ещё немного скорости? Да — можно. В этом нам помогут планировщики.

Планировщики

Библиотека Diffusers позволяет прикрепить планировщик к конвейеру Stable Diffusion. Планировщики пытаются найти наилучший компромисс между скоростью и качеством устранения шума.

Вот что написано об этом в документации: «На момент написания этого документа DPMSolverMultistepScheduler даёт, вероятно, наилучшее соотношение скорость/качество, работа может быть выполнена всего за 20 шагов».

Попробуем эту штуку:

from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
...
dpm = DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=dpm)

Эта последняя версия кода демонстрирует дальнейшее снижение среднего времени ожидания результата. Теперь это — 5,05 секунды. Если сравнить это с нашим исходным экспериментом на Sapphire Rapids (32,3 секунды), то мы получим почти 6,5-кратное ускорение!

Итоги

Возможность создания высококачественных изображений за время, исчисляемое секундами, должна прийтись очень кстати во многих ситуациях. Например — в клиентских приложениях, в задачах генерирования контента для маркетинга и сетевых изданий, при формировании синтетических картинок, используемых для расширения наборов данных.

Вот несколько полезных ресурсов на эту тему:

• Документация по Diffusers.

• Документация по Optimum Intel.

• Страница Intel IPEX на GitHub.

• Материалы для разработчиков от Intel и Hugging Face.

Если у вас есть вопросы или пожелания — мы ждём их на форуме Hugging Face.

Спасибо всем, кто дочитал до этого места!