Precisão e compilação

Precisão reduzida e compilação são duas formas de acelerar a inferência do Diffusers. Carregue os pesos em bfloat16 ou float16 e, em seguida, compile o denoiser (o UNet ou o transformer) com torch.compile ou compilação regional.

Tipo de dados do modelo

A precisão e o tipo de dados dos pesos do modelo afetam a velocidade de inferência, pois uma precisão maior exige mais memória para carregar e mais tempo para executar os cálculos. O Diffusers carrega os pesos do modelo em float32 quando você omite dtype, portanto, alterar o tipo de dados é uma maneira simples de obter inferência mais rápida rapidamente.

bfloat16

float16

TensorFloat-32

Atenção de produto escalar

Scaled dot product attention (SDPA) é a atenção padrão no PyTorch 2.0 e versões posteriores, por meio de AttnProcessor2_0 ou do backend native do despachador de atenção. O PyTorch seleciona um kernel de SDPA para o seu hardware. Para FlashAttention, SageAttention, xFormers, kernels do Hub e outros backends, consulte Attention backends.

torch.compile

torch.compile acelera a inferência compilando o código e as operações do PyTorch em kernels otimizados. Normalmente, você compila o denoiser, que domina o tempo de execução, seja pipeline.unet ou pipeline.transformer, e às vezes também o VAE.

Ative as seguintes configurações do compilador para máxima velocidade (consulte a lista completa para mais opções).

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

Carregue e compile o denoiser e o VAE. Use pipeline.unet em pipelines UNet, como o Stable Diffusion XL, ou pipeline.transformer em pipelines Flux e DiT. Existem vários modos diferentes para você escolher, mas "max-autotune" busca os kernels mais rápidos e usa CUDA graphs. Os CUDA graphs reduzem o overhead ao lançar múltiplas operações de GPU por meio de uma única operação de CPU.

Alterar o layout de memória para channels_last também pode acelerar modelos com muitas convoluções, como UNets e VAEs. Faça um benchmark primeiro, pois alguns modelos ficam mais lentos.

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]

A compilação é lenta na primeira vez, especialmente com "max-autotune", porque o compilador busca kernels e continua usando o mesmo objeto de pipeline compilado em seguida (veja Cache de tempo de compilação e Configuração de cache para saber como configurar o comportamento de cache). Chamar o pipeline compilado com um tamanho de imagem diferente dispara a recompilação.

Compilação de formas dinâmicas

torch.compile acompanha as formas e condições de entrada e, se estas forem diferentes, recompila o modelo. Por exemplo, se um modelo for compilado em uma imagem de resolução 1024x1024 e usado em uma imagem com resolução diferente, isso aciona a recompilação.

Para evitar recompilação, adicione dynamic=True para tentar gerar um kernel mais dinâmico e evitar recompilação quando as condições mudarem.

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

Especificar use_duck_shape=False instrui o compilador se ele deve usar a mesma variável simbólica para representar tamanhos de entrada que são iguais. Para mais detalhes, confira este comentário.

Nem todos os modelos podem se beneficiar da compilação dinâmica imediatamente e podem exigir alterações. Consulte este PR que melhorou a implementação do AuraFlowPipeline para se beneficiar da compilação dinâmica.

Sinta-se à vontade para abrir uma issue se a compilação dinâmica não funcionar como esperado para um modelo do Diffusers.

Compilação regional

Compilação regional reduz a latência de inicialização ao compilar apenas o bloco pequeno e repetido com frequência de um modelo — tipicamente uma camada de transformer — e permite reutilizar os artefatos compilados para cada ocorrência subsequente. Para muitas arquiteturas de difusão, isso proporciona acelerações de tempo de execução comparáveis às da compilação de grafo completo e pode reduzir o tempo de compilação em até 8–10x.

Use o método compile_repeated_blocks(), um auxiliar que encapsula torch.compile, no denoiser. Chame-o em pipeline.unet ou pipeline.transformer, dependendo do pipeline.

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)

Para habilitar a compilação regional para um novo modelo, defina _repeated_blocks com os nomes das classes dos blocos repetidos (strings). Para o UNet do SDXL, use BasicTransformerBlock. Os denoisers Transformer listam os nomes das classes de seus blocos repetidos da mesma forma.

_repeated_blocks = ["BasicTransformerBlock"]

Há também um método compile_regions no Accelerate que seleciona automaticamente os blocos candidatos em um modelo para compilação. O restante do grafo é compilado separadamente. Isso é útil para experimentos rápidos, pois não há tantas opções para você definir quais blocos compilar ou ajustar as flags de compilação.

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() é intencionalmente explícito. Liste os blocos a repetir em _repeated_blocks e o helper compila apenas esses blocos. Ele oferece um comportamento previsível e facilita o raciocínio sobre a reutilização do cache em uma linha de código.

Quebras de gráfico

Defina fullgraph=True para que o torch.compile gere um erro em caso de quebra de grafo, em vez de dividir o grafo silenciosamente, o que reduz a aceleração.

Sincronização de GPU

Após cada previsão do denoiser, o pipeline chama o método step do scheduler, que indexa sigmas. Manter sigmas na GPU pode forçar uma sincronização CPU↔GPU a cada passo. Esse custo é fácil de ignorar até que o denoiser seja compilado. Prefira deixar sigmas na CPU (schedulers do Diffusers, como o Euler, já fazem isso).

Benchmarks

Consulte o conjunto de dados diffusers/benchmarks para ver dados de latência de inferência e uso de memória para pipelines compilados.

O repositório diffusers-torchao também contém resultados de benchmarking para as versões compiladas do Flux e do CogVideoX.

Kernels

Kernels é uma biblioteca para construir, distribuir e carregar kernels de computação otimizados no Hub. Ela suporta kernels de atenção e kernels CUDA personalizados para operações como RMSNorm, GEGLU, RoPE e AdaLN.

O guia de Integração do Pipeline do Diffusers mostra como integrar um kernel com a habilidade add cuda-kernels. Essa habilidade permite que um agente, como Claude ou Codex, escreva kernels personalizados voltados para um modelo específico e seu hardware. O post Custom kernels for all from Codex and Claude traz mais detalhes.

Por exemplo, um kernel RMSNorm personalizado (gerado pela skill add cuda-kernels) com torch.compile acelera a geração do LTX-Video em 1,43x em um H100.

Quantização dinâmica

A quantização dinâmica de ativação int8 está disponível no torchao. Use TorchAoConfig com uma configuração como Int8DynamicActivationInt8WeightConfig, ou a API quantize_ do torchao. Não use o helper mais antigo apply_dynamic_quant.

Matrizes de projeção fundidas

~DiffusionPipeline.fuse_qkv_projections é experimental.

Uma entrada é projetada em três subespaços, representados pelas matrizes de projeção Q, K e V, em um bloco de atenção. Essas projeções são tipicamente calculadas separadamente, mas é possível combiná-las horizontalmente em uma única matriz e realizar a projeção em uma única etapa. Isso aumenta o tamanho das multiplicações de matrizes das projeções de entrada e também melhora o impacto da quantização.

pipeline.fuse_qkv_projections()

Próximas etapas

Leia o post no blog Presenting Flux Fast: Making Flux go brrr on H100s para saber mais sobre como você pode combinar todas essas otimizações com TorchInductor e AOTInductor para uma aceleração de ~2,5x usando as receitas de flux-fast.

Essas receitas suportam hardware AMD e Flux.1 Kontext Dev.
Consulte o post no blog torch.compile and Diffusers: A Hands-On Guide to Peak Performance para maximizar o desempenho com torch.compile em modelos de difusão.

Update on GitHub