diffusers/v0.43.0/ru/optimization/fp16.md

Ускорение вывода

Модели диффузии медленны на этапе вывода, поскольку генерация — это итеративный процесс, в ходе которого шум постепенно уточняется до изображения или видео за определённое количество «шагов». Чтобы ускорить этот процесс, можно поэкспериментировать с различными планировщиками, снизить точность весов модели для более быстрых вычислений, использовать более эффективные по памяти механизмы внимания и так далее.

Сочетайте и применяйте эти техники вместе, чтобы ускорить инференс по сравнению с использованием любой из них по отдельности.

В этом руководстве описано, как ускорить инференс.

Тип данных модели

Точность и тип данных весов модели влияют на скорость инференса, поскольку более высокая точность требует больше памяти для загрузки и больше времени для выполнения вычислений. PyTorch по умолчанию загружает веса модели в float32 или с полной точностью, поэтому смена типа данных — простой способ быстро ускорить инференс.

bfloat16 похож на float16, но более устойчив к численным ошибкам. Поддержка bfloat16 на уровне аппаратного обеспечения различается, однако большинство современных GPU способны работать с bfloat16.

import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

float16 похож на bfloat16, но более подвержен численным ошибкам.

import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

Режим TensorFloat-32 (tf32) поддерживается на GPU NVIDIA Ampere и выполняет операции свёртки и матричного умножения в tf32. Хранение данных и другие операции остаются в float32. Это позволяет значительно ускорить вычисления при использовании в сочетании с bfloat16 или float16.

PyTorch по умолчанию включает режим tf32 только для свёрток, поэтому для матричных умножений его необходимо включить явно.

import torch
from diffusers import StableDiffusionXLPipeline

torch.backends.cuda.matmul.allow_tf32 = True

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

См. документацию по обучению с смешанной точностью для получения дополнительных сведений.

Масштабированное внимание на основе скалярного произведения

Memory-efficient attention оптимизирует скорость и потребление памяти при инференсе!

Scaled dot product attention (SDPA) реализует несколько бэкендов внимания: FlashAttention, xFormers и нативную реализацию на C++. Он автоматически выбирает наиболее оптимальный бэкенд для вашего оборудования.

SDPA включён по умолчанию, если вы используете PyTorch >= 2.0, и никаких дополнительных изменений в вашем коде не требуется. Однако вы можете поэкспериментировать с другими бэкендами внимания, если хотите выбрать свой. В примере ниже используется контекстный менеджер torch.nn.attention.sdpa_kernel для включения эффективного внимания.

from torch.nn.attention import SDPBackend, sdpa_kernel
import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"

with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
  image = pipeline(prompt, num_inference_steps=30).images[0]

torch.compile

torch.compile ускоряет инференс, компилируя код и операции PyTorch в оптимизированные ядра. Diffusers, как правило, компилирует модели с высокой вычислительной нагрузкой, такие как UNet, transformer или VAE.

Включите следующие настройки компилятора для максимальной скорости (см. полный список для других вариантов).

import torch
from diffusers import StableDiffusionXLPipeline

torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True

Загрузите и скомпилируйте UNet и VAE. Доступно несколько режимов, но "max-autotune" оптимизирует работу для максимальной скорости, компилируя код в CUDA graph. CUDA graph эффективно снижает накладные расходы, запуская несколько операций на GPU через одну операцию на CPU.

В PyTorch 2.3.1 можно управлять поведением кэширования torch.compile. Это особенно полезно для режимов компиляции, таких как "max-autotune", который выполняет поиск по сетке нескольких флагов компиляции, чтобы найти оптимальную конфигурацию. Подробнее в учебнике Compile Time Caching in torch.compile.

Изменение раскладки памяти на channels_last также оптимизирует использование памяти и скорость вывода.

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet.to(memory_format=torch.channels_last)
pipeline.vae.to(memory_format=torch.channels_last)
pipeline.unet = torch.compile(
    pipeline.unet, mode="max-autotune", fullgraph=True
)
pipeline.vae.decode = torch.compile(
    pipeline.vae.decode,
    mode="max-autotune",
    fullgraph=True
)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

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

Компиляция динамических форм

Всегда используйте ночную версию PyTorch для лучшей поддержки.

torch.compile отслеживает формы входных данных и условия, и если они отличаются, пересобирает модель. Например, если модель была скомпилирована для изображения с разрешением 1024x1024, а затем используется с изображением другого разрешения, это вызывает повторную компиляцию.

Чтобы избежать повторной компиляции, добавьте dynamic=True, чтобы попытаться сгенерировать более динамичное ядро и избежать повторной компиляции при изменении условий.

+ torch.fx.experimental._config.use_duck_shape = False
+ pipeline.unet = torch.compile(
    pipeline.unet, fullgraph=True, dynamic=True
)

Указание use_duck_shape=False сообщает компилятору, следует ли использовать одну и ту же символьную переменную для представления входных размеров, которые совпадают. Подробнее см. в этом комментарии.

Не все модели могут получать преимущества от динамической компиляции «из коробки» и могут требовать изменений. См. этот PR, в котором была улучшена реализация AuraFlowPipeline для использования преимуществ динамической компиляции.

Если динамическая компиляция не работает как ожидается для модели Diffusers, не стесняйтесь создать issue.

Региональная компиляция

Региональная компиляция сокращает задержку холодного запуска, компилируя только небольшие и часто повторяющиеся блоки модели — как правило, слой трансформера — и позволяет переиспользовать скомпилированные артефакты для каждого последующего вхождения. Для многих архитектур диффузии это обеспечивает такое же ускорение времени выполнения, как и компиляция полного графа, и сокращает время компиляции в 8–10 раз.

Используйте метод compile_repeated_blocks() — вспомогательный, который оборачивает torch.compile — для любого компонента, например для трансформерной модели, как показано ниже.

# pip install -U diffusers
import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"

# compile only the repeated transformer layers inside the UNet
pipeline.unet.compile_repeated_blocks(fullgraph=True)

Чтобы включить региональную компиляцию для новой модели, добавьте атрибут _repeated_blocks в класс модели, содержащий имена классов (в виде строк) блоков, которые вы хотите скомпилировать.

class MyUNet(ModelMixin):
    _repeated_blocks = ("Transformer2DModel",)  # ← compiled by default

Для примеров региональной компиляции см. ссылку PR.

В Accelerate также есть метод compile_regions, который автоматически выбирает блоки модели для компиляции. Остальная часть графа компилируется отдельно. Это удобно для быстрых экспериментов, так как вам не нужно настраивать, какие именно блоки компилировать, или подбирать флаги компиляции.

# pip install -U accelerate
import torch
from diffusers import StableDiffusionXLPipeline
from accelerate.utils import compile_regions

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet = compile_regions(pipeline.unet, mode="reduce-overhead", fullgraph=True)

compile_repeated_blocks() намеренно явен. Перечислите блоки, подлежащие повторению, в _repeated_blocks, и вспомогательная функция скомпилирует только эти блоки. Это обеспечивает предсказуемое поведение и позволяет легко рассуждать о повторном использовании кэша в одной строке кода.

Разрывы графа

Важно указать fullgraph=True в torch.compile, чтобы убедиться, что в базовой модели нет разрывов графа. Это позволяет использовать torch.compile без потери производительности. Для UNet и VAE это меняет способ доступа к возвращаемым переменным.

- latents = unet(
-   latents, timestep=timestep, encoder_hidden_states=prompt_embeds
-).sample

+ latents = unet(
+   latents, timestep=timestep, encoder_hidden_states=prompt_embeds, return_dict=False
+)[0]

Синхронизация GPU

Функция step() вызывается на планировщике каждый раз после того, как денуизер делает предсказание, а переменная sigmas индексируется. При размещении на GPU это вносит задержку из-за синхронизации обмена данными между CPU и GPU. Это становится более заметным, когда денуизер уже скомпилирован.

В целом, sigmas следует оставлять на CPU, чтобы избежать синхронизации и задержек при обмене данными.

Refer to the torch.compile and Diffusers: A Hands-On Guide to Peak Performance blog post for maximizing performance with torch.compile for diffusion models.

Бенчмарки

См. датасет diffusers/benchmarks, чтобы ознакомиться с данными о задержке инференса и использовании памяти для скомпилированных пайплайнов.

Репозиторий diffusers-torchao также содержит результаты бенчмарков для скомпилированных версий Flux и CogVideoX.

Ядра

Kernels — это библиотека для создания, распространения и загрузки оптимизированных вычислительных ядер на Hub. Она поддерживает ядра attention и пользовательские CUDA-ядра для операций, таких как RMSNorm, GEGLU, RoPE и AdaLN.

Руководство по интеграции с пайплайнами Diffusers показывает, как интегрировать ядро с навыком add cuda-kernels. Этот навык позволяет агенту, такому как Claude или Codex, писать кастомные ядра, нацеленные на конкретную модель и ваше оборудование.

Install the add cuda-kernels skill to teach an agent how to write a kernel. The Custom kernels for all from Codex and Claude blog post covers this in more detail.

Например, пользовательское ядро RMSNorm (сгенерированное навыком add cuda-kernels) в сочетании с torch.compile ускоряет генерацию LTX-Video в 1,43 раза на H100.