Точность и компиляция

Снижение точности и компиляция — два способа ускорить инференс в Diffusers. Загрузите веса в bfloat16 или float16, затем скомпилируйте денуайзер (UNet или трансформер) с помощью torch.compile или региональной компиляции.

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

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

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]

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]

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.float32
).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]

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

Scaled dot product attention (SDPA) является механизмом внимания по умолчанию в PyTorch 2.0 и более поздних версиях, через AttnProcessor2_0 или native бэкенд диспетчера внимания. PyTorch выбирает ядро SDPA для вашего оборудования. О FlashAttention, SageAttention, xFormers, Hub-ядрах и других бэкендах см. Attention backends.

torch.compile

torch.compile ускоряет инференс, компилируя код и операции PyTorch в оптимизированные ядра. Обычно компилируют денуизер, который определяет время выполнения, — pipeline.unet или pipeline.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

Загрузите и скомпилируйте denoiser и VAE. Для UNet-пайплайнов, таких как Stable Diffusion XL, используйте pipeline.unet, а для пайплайнов Flux и DiT — pipeline.transformer. Доступно несколько режимов, но "max-autotune" ищет самые быстрые ядра и использует CUDA graphs. CUDA graphs снижают накладные расходы, запуская несколько операций на GPU через одну операцию на CPU.

Изменение раскладки памяти на channels_last также может ускорить модели с преобладанием свёрток, такие как UNet и VAE. Сначала проведите бенчмаркинг, так как некоторые модели работают медленнее.

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.unet = torch.compile(
    pipeline.unet, mode="max-autotune", fullgraph=True
)
pipeline.vae.to(memory_format=torch.channels_last)
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]

Компиляция выполняется медленно при первом запуске, особенно с "max-autotune", поскольку компилятор ищет ядра, а затем продолжает использовать тот же объект скомпилированного конвейера (см. Кэширование времени компиляции и Конфигурация кэширования, чтобы настроить поведение кэширования). Вызов скомпилированного конвейера с другим размером изображения запускает повторную компиляцию.

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

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 — для денойзера. Вызывайте его на pipeline.unet или pipeline.transformer в зависимости от пайплайна.

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"
pipeline.unet.compile_repeated_blocks(fullgraph=True)

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

_repeated_blocks = ["BasicTransformerBlock"]

В 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 выдавал ошибку при разрыве графа, а не молча разбивал его, что снижает ускорение.

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

После каждого предсказания денойзера пайплайн вызывает метод step планировщика, который обращается к sigmas. Если хранить sigmas на GPU, это может вызывать синхронизацию CPU↔GPU на каждом шаге. Эта стоимость легко упускается из виду, пока денойзер не скомпилирован. Рекомендуется оставлять sigmas на CPU (планировщики Diffusers, такие как Euler, уже делают это).

Бенчмарки

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

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

Ядра

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

The Diffusers Pipeline Integration guide shows how to integrate a kernel with the add cuda-kernels skill. This skill enables an agent, like Claude or Codex, to write custom kernels targeted towards a specific model and your hardware. The Custom kernels for all from Codex and Claude post has more detail.

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

Динамическое квантование

Динамическая квантизация активаций int8 доступна в torchao. Используйте TorchAoConfig с конфигурацией, например Int8DynamicActivationInt8WeightConfig, или API quantize_ из torchao. Не используйте устаревший помощник apply_dynamic_quant.

Сквозные матрицы проекций

~DiffusionPipeline.fuse_qkv_projections является экспериментальным.

Входные данные проецируются в три подпространства, представленные матрицами проекций Q, K и V, в блоке внимания. Обычно эти проекции вычисляются отдельно, но их можно объединить по горизонтали в единую матрицу и выполнить проекцию за один шаг. Это увеличивает размер матричных умножений для входных проекций и также улучшает влияние квантования.

pipeline.fuse_qkv_projections()

Дальнейшие шаги

• Ознакомьтесь с статьей в блоге «Presenting Flux Fast: Making Flux go brrr on H100s», чтобы узнать, как можно объединить все эти оптимизации с TorchInductor и AOTInductor для ускорения примерно в 2,5 раза, используя рецепты из flux-fast.

  Эти рецепты поддерживают оборудование AMD и Flux.1 Kontext Dev.

• 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.