torchao

torchao fornece dtypes de alto desempenho e otimizações baseadas em quantização e esparsidade para inferência e treinamento de modelos PyTorch. É suportado para qualquer modelo em qualquer modalidade, desde que suporte carregamento com Accelerate e contenha camadas torch.nn.Linear.

Certifique-se de que o Pytorch 2.5+ e o torchao estão instalados com o comando abaixo.

uv pip install -U torch torchao

Cada dtype de quantização está disponível como uma instância separada de uma classe AOBaseConfig. Isso fornece opções de configuração mais flexíveis ao expor mais argumentos disponíveis.

Passe o AOBaseConfig de um dtype de quantização, como Int4WeightOnlyConfig, para TorchAoConfig em from_pretrained().

import torch
from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConfig
from torchao.quantization import Int8WeightOnlyConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": TorchAoConfig(Int8WeightOnlyConfig(group_size=128, version=2))}
)
pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
    device_map="cuda"  # or "mps", "xpu", "cpu"
)

device_map="cuda" quantiza cada camada na GPU durante o carregamento. Isso é rápido, mas requer temporariamente memória de GPU adicional para os pesos originais e quantizados. Se o modelo já usa a maior parte da memória da sua GPU, o carregamento pode falhar com um erro de falta de memória. Nesse caso, remova device_map e mova o pipeline para a GPU após o carregamento:

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    torch_dtype=torch.bfloat16,
)
pipeline.to("cuda")  # or "mps", "xpu", "cpu"

Sem device_map, o Diffusers quantiza as camadas na CPU. Isso é mais lento, mas evita o pico temporário de memória de GPU durante a quantização. Para reduzir ainda mais o uso de memória de GPU, use enable_model_cpu_offload() no lugar. Você também pode quantizar componentes adicionais, como o text encoder.

torch.compile

O torchao oferece suporte a torch.compile, que pode acelerar a inferência com uma única linha de código.

import torch
from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConfig
from torchao.quantization import Int4WeightOnlyConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": TorchAoConfig(Int4WeightOnlyConfig(group_size=128))}
)
pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
    device_map="cuda"  # or "mps", "xpu", "cpu"
)

pipeline.transformer.compile(transformer, mode="max-autotune", fullgraph=True)

Consulte esta tabela para benchmarks de velocidade de inferência e uso de memória com Flux e CogVideoX. Mais benchmarks em diversos hardwares também estão disponíveis no repositório do torchao.

Os esquemas de quantização pós-treino em FP8 no torchao são eficazes para GPUs com capacidade de computação de pelo menos 8.9 (RTX-4090, Hopper, etc.). O FP8 frequentemente oferece o melhor equilíbrio entre velocidade, memória e qualidade ao gerar imagens e vídeos. Recomendamos combinar FP8 e torch.compile se sua GPU for compatível.

Tipos de quantização suportados

O torchao suporta quantização apenas de pesos e quantização de pesos e ativações dinâmicas para int8, float3-float8 e uint1-uint7.

A quantização apenas de pesos armazena os pesos do modelo em um tipo de dados de baixa bit, mas executa o cálculo com um tipo de dados de maior precisão, como bfloat16. Isso reduz os requisitos de memória dos pesos do modelo, mas mantém os picos de memória para o cálculo de ativações.

A quantização dinâmica de ativações armazena os pesos do modelo em um dtype de baixo bit, quantizando também as ativações em tempo real para economizar memória adicional. Isso reduz os requisitos de memória dos pesos do modelo, além de reduzir a sobrecarga de memória das computações de ativação. No entanto, isso pode implicar em um compromisso de qualidade em alguns casos, portanto, é recomendado testar diferentes modelos minuciosamente.

Consulte a documentação oficial do torchao para uma melhor compreensão dos métodos de quantização disponíveis. Uma lista exaustiva das opções de configuração está disponível aqui.

Alguns exemplos de configurações populares de quantização são os seguintes:

Serialização e desserialização de modelos quantizados

Para serializar um modelo quantizado em um dtype específico, primeiro carregue o modelo com o dtype de quantização desejado e, em seguida, salve-o usando o método save_pretrained().

import torch
from diffusers import AutoModel, TorchAoConfig
from torchao.quantization import Int8WeightOnlyConfig

quantization_config = TorchAoConfig(Int8WeightOnlyConfig())
transformer = AutoModel.from_pretrained(
    "black-forest-labs/Flux.1-Dev",
    subfolder="transformer",
    quantization_config=quantization_config,
    dtype=torch.bfloat16,
)
transformer.save_pretrained("/path/to/flux_int8wo", safe_serialization=False)

Para carregar um modelo quantizado serializado, use o método from_pretrained().

import torch
from diffusers import FluxPipeline, AutoModel

transformer = AutoModel.from_pretrained("/path/to/flux_int8wo", dtype=torch.bfloat16, use_safetensors=False)
pipe = FluxPipeline.from_pretrained("black-forest-labs/Flux.1-Dev", transformer=transformer, dtype=torch.bfloat16)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

prompt = "A cat holding a sign that says hello world"
image = pipe(prompt, num_inference_steps=30, guidance_scale=7.0).images[0]
image.save("output.png")

Se você estiver usando torch<=2.6.0, alguns métodos de quantização, como uint4 weight-only, não podem ser carregados diretamente e podem resultar em um UnpicklingError ao tentar carregar os modelos, mas funcionam como esperado ao salvá-los. Para contornar isso, é possível carregar o state dict manualmente no modelo. Observe, no entanto, que isso requer o uso de weights_only=False em torch.load, portanto, deve ser executado apenas se os pesos tiverem sido obtidos de uma fonte confiável.

import torch
from accelerate import init_empty_weights
from diffusers import FluxPipeline, AutoModel, TorchAoConfig
from torchao.quantization import IntxWeightOnlyConfig

# Serialize the model
transformer = AutoModel.from_pretrained(
    "black-forest-labs/Flux.1-Dev",
    subfolder="transformer",
    quantization_config=TorchAoConfig(IntxWeightOnlyConfig(dtype=torch.uint4)),
    dtype=torch.bfloat16,
)
transformer.save_pretrained("/path/to/flux_uint4wo", safe_serialization=False, max_shard_size="50GB")
# ...

# Load the model
state_dict = torch.load("/path/to/flux_uint4wo/diffusion_pytorch_model.bin", weights_only=False, map_location="cpu")
with init_empty_weights():
    transformer = AutoModel.from_config("/path/to/flux_uint4wo/config.json")
transformer.load_state_dict(state_dict, strict=True, assign=True)

A API AutoModel é suportada para PyTorch >= 2.6, conforme mostrado nos exemplos abaixo.

Recursos

• API de Quantização do TorchAO
• Exemplos de Diffusers-TorchAO