torchao

torchao ofrece dtypes de alto rendimiento y optimizaciones basadas en cuantización y esparsidad para la inferencia y el entrenamiento de modelos de PyTorch. Es compatible con cualquier modelo de cualquier modalidad, siempre que admita la carga con Accelerate y contenga capas torch.nn.Linear.

Asegúrate de que Pytorch 2.5+ y torchao estén instalados con el comando siguiente.

uv pip install -U torch torchao

Cada dtype de cuantización está disponible como una instancia separada de la clase AOBaseConfig. Esto ofrece opciones de configuración más flexibles al exponer más argumentos disponibles.

Pase el AOBaseConfig de un dtype de cuantización, como Int4WeightOnlyConfig, a TorchAoConfig en 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" cuantiza cada capa en la GPU mientras se carga. Esto es rápido, pero requiere temporalmente memoria adicional de GPU para los pesos originales y cuantizados. Si el modelo ya utiliza la mayor parte de la memoria de tu GPU, la carga puede fallar con un error de memoria insuficiente. En ese caso, elimina device_map y mueve el pipeline a la GPU después de cargarlo:

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"

Sin device_map, Diffusers cuantiza las capas en la CPU. Esto es más lento, pero evita el pico temporal de memoria de GPU durante la cuantización. Para reducir aún más el uso de memoria de GPU, usa enable_model_cpu_offload() en su lugar. También puedes cuantizar componentes adicionales, como el codificador de texto.

torch.compile

torchao admite torch.compile, que puede acelerar la inferencia con una sola línea 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 tabla para los benchmarks de velocidad de inferencia y uso de memoria con Flux y CogVideoX. También hay más benchmarks en varios hardware disponibles en el repositorio de torchao.

Los esquemas de cuantización post-entrenamiento en FP8 de torchao son efectivos para GPUs con capacidad de cómputo de al menos 8.9 (RTX-4090, Hopper, etc.). FP8 suele ofrecer el mejor equilibrio entre velocidad, memoria y calidad al generar imágenes y videos. Recomendamos combinar FP8 y torch.compile si tu GPU es compatible.

Tipos de cuantización admitidos

torchao admite la cuantización solo de pesos y la cuantización de pesos y activaciones dinámicas para int8, float3-float8 y uint1-uint7.

La cuantización solo de pesos almacena los pesos del modelo en un tipo de datos de bajo número de bits, pero realiza los cálculos con un tipo de datos de mayor precisión, como bfloat16. Esto reduce los requisitos de memoria derivados de los pesos del modelo, pero mantiene los picos de memoria para el cálculo de las activaciones.

La cuantización dinámica de activaciones almacena los pesos del modelo en un dtype de bajo número de bits, mientras que también cuantiza las activaciones sobre la marcha para ahorrar memoria adicional. Esto reduce los requisitos de memoria de los pesos del modelo, al tiempo que reduce la sobrecarga de memoria de los cálculos de activación. Sin embargo, esto puede implicar un compromiso en la calidad en algunos casos, por lo que se recomienda probar a fondo diferentes modelos.

Consulte la documentación oficial de torchao para comprender mejor los métodos de cuantización disponibles. Una lista exhaustiva de las opciones de configuración está disponible aquí.

Algunas configuraciones de cuantización populares y de ejemplo son las siguientes:

Serialización y deserialización de modelos cuantizados

Para serializar un modelo cuantizado en un dtype dado, primero cargue el modelo con el dtype de cuantización deseado y luego guárdelo usando el 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 cargar un modelo cuantizado serializado, use el 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")

Si estás usando torch<=2.6.0, algunos métodos de cuantización, como uint4 solo de pesos, no se pueden cargar directamente y pueden provocar un UnpicklingError al intentar cargar los modelos, aunque funcionan como se espera al guardarlos. Para evitar esto, se puede cargar el state dict manualmente en el modelo. Ten en cuenta, sin embargo, que esto requiere usar weights_only=False en torch.load, por lo que solo debe ejecutarse si los pesos se obtuvieron de una fuente de confianza.

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)

La API de AutoModel es compatible con PyTorch >= 2.6, como se muestra en los ejemplos a continuación.

Recursos

• API de cuantización de TorchAO
• Ejemplos de Diffusers-TorchAO