bitsandbytes

bitsandbytes es la opción más sencilla para cuantizar un modelo a 8 y 4 bits. La cuantización a 8 bits multiplica los valores atípicos en fp16 con los no atípicos en int8, convierte los valores no atípicos de vuelta a fp16 y luego los suma para devolver los pesos en fp16. Esto reduce el efecto degradante que los valores atípicos tienen sobre el rendimiento de un modelo.

La cuantización de 4 bits comprime aún más un modelo y se utiliza comúnmente con QLoRA para ajustar LLMs cuantizados.

Esta guía demuestra cómo la cuantización permite ejecutar FLUX.1-dev con menos de 16 GB de VRAM e incluso en una instancia gratuita de Google Colab.

Para usar bitsandbytes, asegúrate de tener instaladas las siguientes bibliotecas:

pip install diffusers transformers accelerate bitsandbytes -U

Ahora puedes cuantizar un modelo pasando una BitsAndBytesConfig a from_pretrained(). Esto funciona para cualquier modelo de cualquier modalidad, siempre que admita la carga con Accelerate y contenga capas torch.nn.Linear.

Cuantizar un modelo en 8 bits reduce a la mitad el uso de memoria:

bitsandbytes es compatible tanto con Transformers como con Diffusers, por lo que puedes cuantizar tanto el FluxTransformer2DModel como el T5EncoderModel.

Para GPUs de la serie Ada y superiores, recomendamos cambiar dtype a torch.bfloat16.

El CLIPTextModel y el AutoencoderKL no están cuantizados porque ya son pequeños en tamaño y porque AutoencoderKL solo tiene unas pocas capas torch.nn.Linear.

Copiado

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
import torch
from diffusers import AutoModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(load_in_8bit=True,)

text_encoder_2_8bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True,)

transformer_8bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

Por defecto, todos los demás módulos, como torch.nn.LayerNorm, se convierten a torch.float16. Puedes cambiar el tipo de datos de estos módulos con el parámetro dtype.

Copiado

transformer_8bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
+   dtype=torch.float32,
)

Generemos una imagen usando nuestros modelos cuantizados.

Establecer device_map="auto" llena automáticamente todo el espacio disponible primero en la GPU o las GPUs, luego en la CPU y, por último, en el disco duro (la opción más lenta) si aún no hay suficiente memoria.

Copiado

from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=transformer_8bit,
    text_encoder_2=text_encoder_2_8bit,
    dtype=torch.float16,
    device_map="auto",
)

pipe_kwargs = {
    "prompt": "A cat holding a sign that says hello world",
    "height": 1024,
    "width": 1024,
    "guidance_scale": 3.5,
    "num_inference_steps": 50,
    "max_sequence_length": 512,
}

image = pipe(**pipe_kwargs, generator=torch.manual_seed(0),).images[0]

Cuando hay suficiente memoria, también puede mover directamente el pipeline a la GPU con .to("cuda") y aplicar enable_model_cpu_offload() para optimizar el uso de memoria de la GPU.

Una vez que un modelo está cuantizado, puedes subir el modelo al Hub con el método push_to_hub(). El archivo de configuración de cuantización config.json se sube primero, seguido de los pesos del modelo cuantizado. También puedes guardar los modelos 8-bit serializados localmente con save_pretrained().

El entrenamiento con pesos de 8 bits y 4 bits solo se admite para entrenar parámetros extra.

Comprueba tu huella de memoria con el método get_memory_footprint:

print(model.get_memory_footprint())

Tenga en cuenta que esto solo le indica la huella de memoria de los parámetros del modelo y no estima los requisitos de memoria para la inferencia.

Los modelos cuantizados se pueden cargar desde el método from_pretrained() sin necesidad de especificar los parámetros de quantization_config:

from diffusers import AutoModel, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(load_in_4bit=True)

model_4bit = AutoModel.from_pretrained(
    "hf-internal-testing/flux.1-dev-nf4-pkg", subfolder="transformer"
)

8-bit (algoritmo LLM.int8())

¡Aprende más sobre los detalles de la cuantización de 8 bits en esta entrada de blog!

Esta sección explora algunas de las características específicas de los modelos de 8 bits, como los umbrales de valores atípicos y el omitir la conversión de módulos.

Umbral de valores atípicos

Un “outlier” es un valor de estado oculto que supera cierto umbral, y estos valores se calculan en fp16. Si bien los valores suelen estar distribuidos de forma normal ([-3.5, 3.5]), esta distribución puede ser muy diferente en modelos grandes ([-60, 6] o [6, 60]). La cuantización de 8 bits funciona bien para valores ~5, pero más allá de eso, hay un penalización significativa en el rendimiento. Un valor de umbral predeterminado adecuado es 6, pero puede ser necesario un umbral más bajo para modelos más inestables (modelos pequeños o finetuning).

Para encontrar el mejor umbral para su modelo, le recomendamos experimentar con el parámetro llm_int8_threshold en BitsAndBytesConfig:

from diffusers import AutoModel, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(
    load_in_8bit=True, llm_int8_threshold=10,
)

model_8bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quantization_config,
)

Omitir la conversión de módulos

Para algunos modelos, no es necesario cuantizar todos los módulos a 8 bits, ya que esto puede causar inestabilidad. Por ejemplo, en modelos de difusión como Stable Diffusion 3, el módulo proj_out puede omitirse mediante el parámetro llm_int8_skip_modules en BitsAndBytesConfig:

from diffusers import SD3Transformer2DModel, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(
    load_in_8bit=True, llm_int8_skip_modules=["proj_out"],
)

model_8bit = SD3Transformer2DModel.from_pretrained(
    "stabilityai/stable-diffusion-3-medium-diffusers",
    subfolder="transformer",
    quantization_config=quantization_config,
)

4-bit (algoritmo QLoRA)

Obtenga más detalles en esta entrada de blog.

Esta sección explora algunas de las características específicas de los modelos de 4 bits, como cambiar el tipo de datos de cálculo, usar el tipo de datos Normal Float 4 (NF4) y usar cuantización anidada.

Tipo de datos de cálculo

Para acelerar el cálculo, puedes cambiar el tipo de datos de float32 (el valor predeterminado) a bf16 usando el parámetro bnb_4bit_compute_dtype en BitsAndBytesConfig:

import torch
from diffusers import BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)

Float 4 normal (NF4)

NF4 es un tipo de datos de 4 bits del artículo QLoRA, adaptado para pesos inicializados desde una distribución normal. Debe usar NF4 para entrenar modelos base de 4 bits. Esto se puede configurar con el parámetro bnb_4bit_quant_type en BitsAndBytesConfig:

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

from diffusers import AutoModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
)

text_encoder_2_4bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
)

transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

Para la inferencia, el bnb_4bit_quant_type no tiene un gran impacto en el rendimiento. Sin embargo, para mantener la consistencia con los pesos del modelo, debe usar los valores de bnb_4bit_compute_dtype y dtype.

Cuantización anidada

La cuantización anidada es una técnica que puede ahorrar memoria adicional sin un coste de rendimiento extra. Esta función realiza una segunda cuantización de los pesos ya cuantizados para ahorrar 0,4 bits/parámetro adicionales.

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

from diffusers import AutoModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)

text_encoder_2_4bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)

transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

Descomprimir modelos de bitsandbytes

Una vez cuantizado, puedes descuantizar un modelo a su precisión original, aunque esto podría provocar una pequeña pérdida de calidad. Asegúrate de tener suficiente RAM de GPU para cargar el modelo descuantizado.

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

from diffusers import AutoModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)

text_encoder_2_4bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)

transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

text_encoder_2_4bit.dequantize()
transformer_4bit.dequantize()

torch.compile

Acelera la inferencia con torch.compile. Asegúrate de tener la última versión de bitsandbytes instalada y también recomendamos instalar PyTorch nightly.

torch._dynamo.config.capture_dynamic_output_shape_ops = True

quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)
transformer_4bit.compile(fullgraph=True)

quant_config = DiffusersBitsAndBytesConfig(load_in_4bit=True)
transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)
transformer_4bit.compile(fullgraph=True)

En una RTX 4090 con compilación, la generación de Flux de 4 bits se completó en 25.809 segundos frente a 32.570 segundos sin ella.

Consulte el script de benchmarking para más detalles.

Recursos

• Cuaderno de principio a fin que muestra la inferencia de Flux.1 Dev en un Colab de nivel gratuito
• Entrenamiento