Acelerar a inferência

Os modelos de difusão são lentos na inferência porque a geração é um processo iterativo no qual o ruído é gradualmente refinado em uma imagem ou vídeo ao longo de um determinado número de “etapas”. Para acelerar esse processo, você pode tentar experimentar diferentes schedulers, reduzir a precisão dos pesos do modelo para cálculos mais rápidos, usar mecanismos de atenção mais eficientes em termos de memória, entre outras coisas.

Combine e utilize essas técnicas em conjunto para tornar a inferência mais rápida do que ao usar qualquer técnica isoladamente.

Este guia explicará como acelerar a inferência.

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 PyTorch carrega os pesos do modelo em float32 ou precisão completa por padrão, 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

A atenção eficiente em memória otimiza a velocidade de inferência e o uso de memória!

Scaled dot product attention (SDPA) implementa vários backends de atenção, FlashAttention, xFormers e uma implementação nativa em C++. Ele seleciona automaticamente o backend mais otimizado para o seu hardware.

O SDPA é ativado por padrão se você estiver usando PyTorch >= 2.0 e nenhuma alteração adicional é necessária no seu código. No entanto, você pode experimentar outros backends de atenção, caso queira escolher o seu próprio. O exemplo abaixo usa o gerenciador de contexto torch.nn.attention.sdpa_kernel para habilitar a atenção eficiente.

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 acelera a inferência compilando o código e as operações do PyTorch em kernels otimizados. O Diffusers normalmente compila os modelos mais intensivos em computação, como o UNet, o transformer ou 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 UNet e o VAE. Existem vários modos diferentes que você pode escolher, mas "max-autotune" otimiza para a velocidade mais rápida compilando para um grafo CUDA. Os grafos CUDA reduzem efetivamente o overhead ao lançar múltiplas operações de GPU por meio de uma única operação de CPU.

Com o PyTorch 2.3.1, você pode controlar o comportamento de cache do torch.compile. Isso é particularmente benéfico para modos de compilação como "max-autotune", que executa uma busca em grade sobre várias flags de compilação para encontrar a configuração ideal. Saiba mais no tutorial Compile Time Caching in torch.compile.

Alterar o layout de memória para channels_last também otimiza a memória e a velocidade de inferência.

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]

A compilação é lenta na primeira vez, mas, uma vez compilado, o processo é significativamente mais rápido. Tente usar o pipeline compilado apenas para o mesmo tipo de operações de inferência. Chamar o pipeline compilado com um tamanho de imagem diferente reativa a compilação, o que é lento e ineficiente.

Compilação de formas dinâmicas

Certifique-se de sempre usar a versão nightly do PyTorch para um melhor suporte.

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 entrega as mesmas acelerações de tempo de execução que a compilação de grafo completo e reduz o tempo de compilação em 8 a 10 vezes.

Use o método compile_repeated_blocks(), um auxiliar que encapsula o torch.compile, em qualquer componente, como o modelo transformer, conforme mostrado abaixo.

# 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)

Para habilitar a compilação regional para um novo modelo, adicione um atributo _repeated_blocks a uma classe de modelo contendo os nomes das classes (como strings) dos blocos que você deseja compilar.

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

Para mais exemplos de compilação regional, consulte a referência PR.

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

É importante especificar fullgraph=True no torch.compile para garantir que não haja quebras de grafo no modelo subjacente. Isso permite aproveitar o torch.compile sem qualquer degradação de desempenho. Para o UNet e o VAE, isso altera a forma como você acessa as variáveis de retorno.

- 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]

Sincronização de GPU

A função step() é chamada no scheduler a cada vez que o denoiser faz uma previsão, e a variável sigmas é indexada. Quando colocada na GPU, ela introduz latência devido à sincronização de comunicação entre a CPU e a GPU. Isso se torna mais evidente quando o denoiser já foi compilado.

Em geral, os sigmas devem permanecer na CPU para evitar a sincronização de comunicação e a latência.

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.

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 com o 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 para o seu hardware.

Instale a habilidade add cuda-kernels para ensinar um agente a escrever um kernel. O post no blog Custom kernels for all from Codex and Claude aborda isso em 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

Quantização dinâmica melhora a velocidade de inferência ao reduzir a precisão para permitir operações matemáticas mais rápidas. Esse tipo específico de quantização determina como escalar as ativações com base nos dados em tempo de execução, em vez de usar um fator de escala fixo. Como resultado, o fator de escala está mais precisamente alinhado com os dados.

O exemplo abaixo aplica quantização dinâmica int8 ao UNet e ao VAE com a biblioteca torchao.

Consulte nossa documentação de torchao para saber mais sobre como usar a integração do torchao do Diffusers.

Configure as tags do compilador para a velocidade máxima.

import torch
from torchao import apply_dynamic_quant
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
torch._inductor.config.force_fuse_int_mm_with_mul = True
torch._inductor.config.use_mixed_mm = True

Filtre alguns camadas lineares na UNet e na VAE que não se beneficiam da quantização dinâmica com o dynamic_quant_filter_fn.

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

apply_dynamic_quant(pipeline.unet, dynamic_quant_filter_fn)
apply_dynamic_quant(pipeline.vae, dynamic_quant_filter_fn)

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

Matrizes de projeção fundidas

O método fuse_qkv_projections é experimental e o suporte é limitado principalmente aos pipelines do Stable Diffusion. Dê uma olhada neste PR para saber mais sobre como ativá-lo para outros pipelines

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()

Recursos

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.
Leia o post no blog torch.compile e Diffusers: um guia prático para desempenho máximo para maximizar o desempenho ao usar torch.compile.

Update on GitHub