Primeros pasos

La cuantización se centra en representar los datos con menos bits, al tiempo que intenta preservar la precisión de los datos originales. Esto suele implicar convertir un tipo de datos para representar la misma información con menos bits. Por ejemplo, si los pesos de su modelo se almacenan como números de coma flotante de 32 bits y se cuantizan a números de coma flotante de 16 bits, el tamaño del modelo se reduce a la mitad, lo que facilita su almacenamiento y reduce el uso de memoria. Una menor precisión también puede acelerar la inferencia, ya que se necesita menos tiempo para realizar cálculos con menos bits.

Diffusers admite varios backends de cuantización para hacer más accesibles los grandes modelos de difusión como Flux. Esta guía muestra cómo usar la clase PipelineQuantizationConfig para cuantizar un pipeline durante su inicialización a partir de un checkpoint preentrenado o no cuantizado.

Cuantización a nivel de pipeline

Hay dos formas de usar PipelineQuantizationConfig según el nivel de personalización que desees aplicar a la configuración de cuantización.

• para casos de uso básicos, defina los argumentos quant_backend, quant_kwargs y components_to_quantize
• Para un control de cuantización más granular, defina un quant_mapping que proporcione la configuración de cuantización para componentes individuales del modelo

Cuantización básica

Inicializa PipelineQuantizationConfig con los siguientes parámetros.

• quant_backend especifica qué backend de cuantización se debe utilizar. Los backends actualmente admitidos incluyen: bitsandbytes_4bit, bitsandbytes_8bit, gguf, quanto y torchao.
• quant_kwargs especifica los argumentos de cuantización a utilizar.

Estos argumentos de quant_kwargs son diferentes para cada backend. Consulte la documentación de la Quantization API para ver los argumentos de cada backend.

• components_to_quantize especifica qué componente(s) del pipeline se deben cuantizar. Por lo general, se debe cuantizar los componentes más intensivos en cómputo, como el transformer. El codificador de texto es otro componente a considerar para la cuantización si el pipeline tiene más de uno, como FluxPipeline. El siguiente ejemplo cuantiza el codificador de texto T5 en FluxPipeline manteniendo el modelo CLIP intacto.

  components_to_quantize acepta una lista para varios modelos o una cadena para un solo modelo.

El siguiente ejemplo carga el backend bitsandbytes con los siguientes argumentos de BitsAndBytesConfig, load_in_4bit, bnb_4bit_quant_type y 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"],
)

Pase pipeline_quant_config a from_pretrained() para cuantizar el 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]

Cuantización avanzada

El argumento quant_mapping ofrece más opciones sobre cómo cuantizar cada componente individual de un pipeline, como combinar diferentes backends de cuantización.

Inicialice PipelineQuantizationConfig y pase un quant_mapping a este. El quant_mapping le permite especificar las opciones de cuantización para cada componente de la pipeline, como el transformer y el text encoder.

El ejemplo siguiente utiliza dos backends de cuantización, QuantoConfig y transformers.BitsAndBytesConfig, para el transformer y el 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 un backend de bitsandbytes separado en Transformers. Debes importar y usar transformers.BitsAndBytesConfig para los componentes que provienen de Transformers. Por ejemplo, text_encoder_2 en FluxPipeline es un T5EncoderModel de Transformers, por lo que necesitas usar transformers.BitsAndBytesConfig en lugar de diffusers.BitsAndBytesConfig.

Utilice el método de cuantización básica anterior si no desea gestionar estas importaciones distintas o no está seguro de dónde proviene cada componente de la canalización.

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
        ),
    }
)

Pase pipeline_quant_config a from_pretrained() para cuantizar el 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]

Carga de puntos de control precuantizados

Algunos backends de cuantización admiten la carga de puntos de control que ya estaban cuantizados y subidos al Hub. En este caso, la configuración de cuantización se almacena en el config.json del modelo y se lee automáticamente: no es necesario crear un PipelineQuantizationConfig ni pasar un 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"

Los siguientes backends admiten la carga de puntos de control precuantizados de forma nativa.

Recursos

Echa un vistazo a los recursos de abajo para aprender más sobre cuantización.

• Si es nuevo en la cuantización, le recomendamos que consulte los siguientes cursos para principiantes, elaborados en colaboración con DeepLearning.AI.

  • Fundamentos de la cuantización con Hugging Face
  • Cuantización en profundidad

• Consulte la Guía para contribuir con un nuevo método de cuantización si le interesa agregar un nuevo método de cuantización.

• El Resumen de la cuantización de Transformers ofrece una visión general de las ventajas y desventajas de los distintos backends de cuantización.

• Lee la publicación del blog Exploring Quantization Backends in Diffusers para una breve introducción a cada backend de cuantización, cómo elegir un backend y cómo combinar la cuantización con otras optimizaciones de memoria.