Começando

A quantização foca em representar dados com menos bits, ao mesmo tempo em que tenta preservar a precisão dos dados originais. Isso geralmente significa converter um tipo de dados para representar a mesma informação com menos bits. Por exemplo, se os pesos do seu modelo forem armazenados como números de ponto flutuante de 32 bits e forem quantizados para números de ponto flutuante de 16 bits, o tamanho do modelo é reduzido pela metade, o que facilita o armazenamento e diminui o uso de memória. A menor precisão também pode acelerar a inferência, pois leva menos tempo para realizar cálculos com menos bits.

O Diffusers suporta vários backends de quantização para tornar grandes modelos de difusão, como o Flux, mais acessíveis. Este guia mostra como usar a classe PipelineQuantizationConfig para quantizar um pipeline durante sua inicialização a partir de um checkpoint pré-treinado ou não quantizado.

Quantização em nível de pipeline

Existem duas formas de usar PipelineQuantizationConfig, dependendo do nível de personalização que você deseja aplicar à configuração de quantização.

para casos de uso básicos, defina os argumentos quant_backend, quant_kwargs e components_to_quantize
para controle granular de quantização, defina um quant_mapping que fornece a configuração de quantização para componentes individuais do modelo

Quantização básica

Inicialize PipelineQuantizationConfig com os seguintes parâmetros.

quant_backend especifica qual backend de quantização usar. Os backends atualmente suportados incluem: bitsandbytes_4bit, bitsandbytes_8bit, gguf, quanto e torchao.
quant_kwargs especifica os argumentos de quantização a serem usados.

Esses argumentos de quant_kwargs são diferentes para cada backend. Consulte a documentação da Quantization API para ver os argumentos de cada backend.

components_to_quantize especifica qual(is) componente(s) do pipeline devem ser quantizados. Normalmente, você deve quantizar os componentes mais intensivos em computação, como o transformer. O codificador de texto é outro componente a ser considerado para quantização se o pipeline tiver mais de um, como no caso de FluxPipeline. O exemplo abaixo quantiza o codificador de texto T5 em FluxPipeline, mantendo o modelo CLIP intacto.

components_to_quantize aceita uma lista para vários modelos ou uma string para um único modelo.

O exemplo abaixo carrega o backend bitsandbytes com os seguintes argumentos de BitsAndBytesConfig, load_in_4bit, bnb_4bit_quant_type e bnb_4bit_compute_dtype.

import torch
from diffusers import DiffusionPipeline
from diffusers.quantizers import PipelineQuantizationConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_backend="bitsandbytes_4bit",
    quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
    components_to_quantize=["transformer", "text_encoder_2"],
)

Passe o pipeline_quant_config para from_pretrained() para quantizar o pipeline.

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

image = pipe("photo of a cute dog").images[0]

Quantização avançada

O argumento quant_mapping oferece mais opções sobre como quantizar cada componente individual em um pipeline, como combinar diferentes backends de quantização.

Inicialize PipelineQuantizationConfig e passe um quant_mapping para ele. O quant_mapping permite especificar as opções de quantização para cada componente do pipeline, como o transformer e o text encoder.

O exemplo abaixo usa dois backends de quantização, QuantoConfig e transformers.BitsAndBytesConfig, para o transformer e o codificador de texto.

import torch
from diffusers import DiffusionPipeline
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from diffusers.quantizers.quantization_config import QuantoConfig
from diffusers.quantizers import PipelineQuantizationConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={
        "transformer": QuantoConfig(weights_dtype="int8"),
        "text_encoder_2": TransformersBitsAndBytesConfig(
            load_in_4bit=True, compute_dtype=torch.bfloat16
        ),
    }
)

Existe um backend bitsandbytes separado no Transformers. Você precisa importar e usar transformers.BitsAndBytesConfig para componentes que vêm do Transformers. Por exemplo, text_encoder_2 em FluxPipeline é um T5EncoderModel do Transformers, então você precisa usar transformers.BitsAndBytesConfig em vez de diffusers.BitsAndBytesConfig.

Use o método de quantização básica acima se você não quiser gerenciar esses imports distintos ou não tiver certeza de onde cada componente do pipeline vem.

import torch
from diffusers import DiffusionPipeline
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from diffusers.quantizers import PipelineQuantizationConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={
        "transformer": DiffusersBitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16),
        "text_encoder_2": TransformersBitsAndBytesConfig(
            load_in_4bit=True, compute_dtype=torch.bfloat16
        ),
    }
)

Passe o pipeline_quant_config para from_pretrained() para quantizar o pipeline.

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

image = pipe("photo of a cute dog").images[0]

Carregando checkpoints pré-quantizados

Alguns backends de quantização suportam o carregamento de checkpoints que já foram quantizados e enviados ao Hub. Nesse caso, a configuração de quantização é armazenada no config.json do modelo e é lida automaticamente — você não precisa criar um PipelineQuantizationConfig nem passar um argumento quantization_config.

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "rootonchair/ERNIE-Image-Turbo-nunchaku-lite-nvfp4",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

Os seguintes backends suportam o carregamento de checkpoints pré-quantizados nativamente.

Backend	Observações
bitsandbytes	Config é salvo em `config.json`; nenhum argumento extra é necessário. Suporta quantização em tempo de execução, bem como o carregamento de checkpoints pré-quantizados.
torchao	Igual ao acima.
GGUF	Usa `from_single_file` com classes de Model; o carregamento em nível de pipeline não é suportado.
BFL	Carrega checkpoints de arquivo único `fp8r` e `nvfp4` da Black Forest Labs com `from_single_file`; o carregamento em nível de pipeline não é suportado.
AutoRound	Apenas o carregamento é suportado; quantize primeiro com a CLI ou a API Python do AutoRound.
Nunchaku Lite	Config é salvo em `config.json`; requer o pacote `kernels`. Apenas o carregamento é suportado.
ModelOpt	Apenas o carregamento de modelos pré-quantizados é suportado.
SDNQ	Suporta tanto a quantização em tempo de execução quanto o carregamento de modelos pré-quantizados; requer o pacote `sdnq`.

Recursos

Confira os recursos abaixo para saber mais sobre quantização.

Se você é novo em quantização, recomendamos conferir os seguintes cursos introdutórios, desenvolvidos em colaboração com a DeepLearning.AI.
- Fundamentos de Quantização com Hugging Face
- Quantização em profundidade
Consulte o Guia para contribuir com um novo método de quantização se você estiver interessado em adicionar um novo método de quantização.
A Visão geral da quantização do Transformers fornece uma visão geral das vantagens e desvantagens dos diferentes backends de quantização.
Leia o post no blog Exploring Quantization Backends in Diffusers para uma breve introdução a cada backend de quantização, como escolher um backend e como combinar a quantização com outras otimizações de memória.

Update on GitHub