Quanto

El backend Quanto está en desuso y se eliminará en la versión 1.0.0. Considere cambiar a uno de los otros backends de cuantización admitidos, como bitsandbytes o torchao.

Quanto es un backend de cuantización de PyTorch para Optimum. Ha sido diseñado con versatilidad y simplicidad en mente:

Todas las funciones están disponibles en modo eager (funciona con modelos no trazables)
Admite el entrenamiento consciente de la cuantización
Los modelos cuantizados son compatibles con torch.compile
Los modelos cuantizados son agnósticos del dispositivo (p. ej., CUDA, XPU, MPS, CPU).

Para usar el backend Quanto, primero debe instalar optimum-quanto>=0.2.6 y accelerate

pip install optimum-quanto accelerate

Ahora puede cuantizar un modelo pasando el objeto QuantoConfig al método from_pretrained(). Aunque la biblioteca Quanto permite cuantizar los módulos nn.Conv2d y nn.LayerNorm, actualmente Diffusers solo admite la cuantización de los pesos en las capas nn.Linear de un modelo. El siguiente fragmento de código muestra cómo aplicar la cuantización float8 con 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")

Omitir la cuantización en módulos específicos

Es posible omitir la aplicación de cuantización en ciertos módulos utilizando el argumento modules_to_not_convert en QuantoConfig. Asegúrese de que los módulos pasados a este argumento coincidan con las claves de los módulos en el 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,
)

Usar from_single_file con el backend Quanto

QuantoConfig es compatible con ~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)

Guardar modelos cuantizados

Diffusers admite la serialización de modelos de Quanto mediante el método ~ModelMixin.save_pretrained.

Los requisitos de serialización y carga son diferentes para los modelos cuantizados directamente con la librería Quanto y los modelos cuantizados con Diffusers usando Quanto como backend. Actualmente no es posible cargar en Diffusers modelos cuantizados directamente con Quanto 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>")

Usar torch.compile con Quanto

Actualmente, el backend de Quanto admite torch.compile para los siguientes tipos de cuantización:

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 cuantización admitidos

Pesos

float8
int8
int4
int2

Update on GitHub