Quanto

O backend Quanto está obsoleto e será removido na versão 1.0.0. Considere alternar para um dos outros backends de quantização suportados, como bitsandbytes ou torchao.

Quanto é um backend de quantização do PyTorch para o Optimum. Ele foi projetado com foco em versatilidade e simplicidade:

Todos os recursos estão disponíveis no modo eager (funciona com modelos que não podem ser rastreados)
Suporta treinamento com quantização
Modelos quantizados são compatíveis com torch.compile
Modelos quantizados são agnósticos em relação ao dispositivo (por exemplo, CUDA, XPU, MPS, CPU)

Para usar o backend Quanto, você primeiro precisa instalar optimum-quanto>=0.2.6 e accelerate

pip install optimum-quanto accelerate

Agora você pode quantizar um modelo passando o objeto QuantoConfig para o método from_pretrained(). Embora a biblioteca Quanto permita quantizar os módulos nn.Conv2d e nn.LayerNorm, atualmente, o Diffusers suporta apenas a quantização dos pesos nas camadas nn.Linear de um modelo. O trecho a seguir demonstra como aplicar a quantização float8 com o Quanto.

import torch
from diffusers import FluxTransformer2DModel, QuantoConfig

model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="float8")
transformer = FluxTransformer2DModel.from_pretrained(
      model_id,
      subfolder="transformer",
      quantization_config=quantization_config,
      dtype=torch.bfloat16,
)

pipe = FluxPipeline.from_pretrained(model_id, transformer=transformer, dtype=dtype)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

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

Pular a Quantização em módulos específicos

É possível pular a aplicação de quantização em determinados módulos usando o argumento modules_to_not_convert no QuantoConfig. Certifique-se de que os módulos passados para este argumento correspondam às chaves dos módulos no state_dict

import torch
from diffusers import FluxTransformer2DModel, QuantoConfig

model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="float8", modules_to_not_convert=["proj_out"])
transformer = FluxTransformer2DModel.from_pretrained(
      model_id,
      subfolder="transformer",
      quantization_config=quantization_config,
      dtype=torch.bfloat16,
)

Usando from_single_file com o Backend Quanto

QuantoConfig é compatível com ~FromOriginalModelMixin.from_single_file.

import torch
from diffusers import FluxTransformer2DModel, QuantoConfig

ckpt_path = "https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/flux1-dev.safetensors"
quantization_config = QuantoConfig(weights_dtype="float8")
transformer = FluxTransformer2DModel.from_single_file(ckpt_path, quantization_config=quantization_config, dtype=torch.bfloat16)

Salvando modelos quantizados

O Diffusers suporta a serialização de modelos Quanto usando o método ~ModelMixin.save_pretrained.

Os requisitos de serialização e carregamento são diferentes para modelos quantizados diretamente com a biblioteca Quanto e modelos quantizados com Diffusers usando Quanto como backend. Atualmente, não é possível carregar modelos quantizados diretamente com Quanto no Diffusers usando ~ModelMixin.from_pretrained

import torch
from diffusers import FluxTransformer2DModel, QuantoConfig

model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="float8")
transformer = FluxTransformer2DModel.from_pretrained(
      model_id,
      subfolder="transformer",
      quantization_config=quantization_config,
      dtype=torch.bfloat16,
)
# save quantized model to reuse
transformer.save_pretrained("<your quantized model save path>")

# you can reload your quantized model with
model = FluxTransformer2DModel.from_pretrained("<your quantized model save path>")

Usando torch.compile com Quanto

Atualmente, o backend Quanto suporta torch.compile para os seguintes tipos de quantização:

pesos int8

import torch
from diffusers import FluxPipeline, FluxTransformer2DModel, QuantoConfig

model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="int8")
transformer = FluxTransformer2DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    quantization_config=quantization_config,
    dtype=torch.bfloat16,
)
transformer = torch.compile(transformer, mode="max-autotune", fullgraph=True)

pipe = FluxPipeline.from_pretrained(
    model_id, transformer=transformer, dtype=dtype
)
pipe.to("cuda")  # or "mps", "xpu", "cpu"
images = pipe("A cat holding a sign that says hello").images[0]
images.save("flux-quanto-compile.png")

Tipos de Quantização Suportados

Pesos

float8
int8
int4
int2

Update on GitHub