bitsandbytes

bitsandbytes est l’option la plus simple pour quantiser un modèle en 8 et 4 bits. La quantification en 8 bits multiplie les valeurs aberrantes en fp16 par les valeurs non aberrantes en int8, convertit les valeurs non aberrantes de nouveau en fp16, puis les additionne pour restituer les poids en fp16. Cela réduit l’effet dégradant des valeurs aberrantes sur les performances du modèle.

La quantification 4 bits comprime encore davantage un modèle et est couramment utilisée avec QLoRA pour l’ajustement fin des LLM quantifiés.

Ce guide montre comment la quantification permet d’exécuter FLUX.1-dev avec moins de 16 Go de VRAM, voire sur une instance Google Colab gratuite.

Pour utiliser bitsandbytes, assurez-vous d'avoir installé les bibliothèques suivantes :

pip install diffusers transformers accelerate bitsandbytes -U

Vous pouvez désormais quantiser un modèle en passant un BitsAndBytesConfig à from_pretrained(). Cela fonctionne pour n'importe quel modèle de n'importe quelle modalité, à condition qu'il prenne en charge le chargement avec Accelerate et qu'il contienne des couches torch.nn.Linear.

La quantification d'un modèle en 8 bits réduit la consommation mémoire de moitié :

bitsandbytes est pris en charge à la fois dans Transformers et Diffusers, vous permettant donc de quantifier à la fois FluxTransformer2DModel et T5EncoderModel.

Pour les GPU de la série Ada et supérieures, nous recommandons de changer dtype en torch.bfloat16.

Les modèles CLIPTextModel et AutoencoderKL ne sont pas quantifiés car ils sont déjà de petite taille et parce que AutoencoderKL ne contient que quelques couches torch.nn.Linear.

Copié

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

Par défaut, tous les autres modules tels que torch.nn.LayerNorm sont convertis en torch.float16. Vous pouvez modifier le type de données de ces modules avec le paramètre dtype.

Copié

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

Générons une image à l'aide de nos modèles quantifiés.

Définir device_map="auto" remplit automatiquement tout l’espace disponible sur le ou les GPU d’abord, puis sur le CPU, et enfin sur le disque dur (l’option la plus lente) s’il reste encore insuffisamment de mémoire.

Copié

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]

Lorsque la mémoire est suffisante, vous pouvez également déplacer directement le pipeline sur le GPU avec .to("cuda") et appliquer enable_model_cpu_offload() pour optimiser l'utilisation de la mémoire GPU.

Une fois un modèle quantisé, vous pouvez le publier sur le Hub à l'aide de la méthode push_to_hub(). Le fichier de configuration de quantisation config.json est d'abord publié, suivi des poids du modèle quantisé. Vous pouvez également enregistrer localement les modèles 8 bits sérialisés avec save_pretrained().

L'entraînement avec des poids en 8 bits et 4 bits n'est pris en charge que pour l'entraînement de paramètres supplémentaires.

Vérifiez votre empreinte mémoire avec la méthode get_memory_footprint :

print(model.get_memory_footprint())

Notez que cela ne vous indique que l’empreinte mémoire des paramètres du modèle et n’estime pas les besoins en mémoire pour l’inférence.

Les modèles quantisés peuvent être chargés à partir de la méthode from_pretrained() sans avoir besoin de spécifier les paramètres 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 bits (algorithme LLM.int8())

En savoir plus sur les détails de la quantification 8 bits dans ce article de blog !

Cette section explore certaines fonctionnalités spécifiques des modèles 8 bits, telles que les seuils de valeurs aberrantes et le saut de la conversion de modules.

Seuil des valeurs aberrantes

Un « outlier » est une valeur d’état caché supérieure à un certain seuil, et ces valeurs sont calculées en fp16. Bien que les valeurs soient généralement distribuées normalement ([-3.5, 3.5]), cette distribution peut être très différente pour les grands modèles ([-60, 6] ou [6, 60]). La quantification 8 bits fonctionne bien pour des valeurs ~5, mais au-delà, il y a une pénalité de performance significative. Une valeur de seuil par défaut appropriée est 6, mais un seuil plus bas peut être nécessaire pour les modèles plus instables (petits modèles ou finetuning).

Pour trouver le meilleur seuil pour votre modèle, nous recommandons d'expérimenter avec le paramètre llm_int8_threshold dans 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,
)

Passer la conversion du module

Pour certains modèles, il n’est pas nécessaire de quantiser chaque module en 8 bits, ce qui peut d’ailleurs provoquer une instabilité. Par exemple, pour les modèles de diffusion comme Stable Diffusion 3, le module proj_out peut être ignoré à l’aide du paramètre llm_int8_skip_modules dans 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 bits (algorithme QLoRA)

En savoir plus sur ses détails dans ce article de blog.

Cette section explore certaines fonctionnalités spécifiques des modèles 4 bits, telles que le changement du type de données de calcul, l'utilisation du type de données Normal Float 4 (NF4) et l'utilisation de la quantification imbriquée.

Type de données de calcul

Pour accélérer le calcul, vous pouvez changer le type de données de float32 (la valeur par défaut) en bf16 à l'aide du paramètre bnb_4bit_compute_dtype dans 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 est un type de données sur 4 bits issu de l'article QLoRA, adapté aux poids initialisés à partir d'une distribution normale. Vous devez utiliser NF4 pour l'entraînement de modèles de base sur 4 bits. Cela peut être configuré avec le paramètre bnb_4bit_quant_type dans 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,
)

Pour l'inférence, le bnb_4bit_quant_type n'a pas d'impact majeur sur les performances. Toutefois, pour rester cohérent avec les poids du modèle, vous devez utiliser les valeurs de bnb_4bit_compute_dtype et dtype.

Quantification imbriquée

La quantification imbriquée est une technique qui permet d’économiser de la mémoire supplémentaire sans coût de performance additionnel. Cette fonctionnalité effectue une seconde quantification des poids déjà quantifiés afin d’économiser 0,4 bit/paramètre supplémentaire.

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

Déquantisation des modèles bitsandbytes

Une fois quantisé, vous pouvez déquantiser un modèle pour revenir à sa précision d’origine, mais cela peut entraîner une légère perte de qualité. Assurez-vous de disposer de suffisamment de RAM GPU pour charger le modèle déquantisé.

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

Accélérez l'inférence avec torch.compile. Assurez-vous d'avoir la dernière version de bitsandbytes installée et nous recommandons également d'installer 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)

Sur une RTX 4090 avec compilation, la génération Flux en 4 bits s’est terminée en 25,809 secondes contre 32,570 secondes sans compilation.

Consultez le script de benchmarking pour plus de détails.

Ressources

• Carnet de notes de bout en bout montrant l'inférence de Flux.1 Dev dans un Colab gratuit
• Entraînement