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

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 é semelhante a float16, mas é mais robusto a erros numéricos. O suporte de hardware para bfloat16 varia, mas a maioria das GPUs modernas é capaz de suportar 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 é semelhante a bfloat16, mas pode ser mais suscetível a erros numéricos.

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]

O modo TensorFloat-32 (tf32) é suportado em GPUs NVIDIA Ampere e executa as operações de convolução e multiplicação de matrizes em tf32. O armazenamento e as demais operações permanecem em float32. Isso permite computações significativamente mais rápidas quando combinado com bfloat16 ou float16.

O PyTorch habilita o modo tf32 apenas para convoluções por padrão e você precisará habilitá-lo explicitamente para multiplicações de matrizes.

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]

Consulte a documentação de treinamento com precisão mista para mais detalhes.

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 há menos 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.