torchao

torchao fournit des dtypes et des optimisations haute performance basés sur la quantification et la sparsité pour l’inférence et l’entraînement de modèles PyTorch. Il prend en charge tout modèle, quelle que soit sa modalité, à condition qu’il soit chargeable avec Accelerate et qu’il contienne des couches torch.nn.Linear.

Vérifiez que Pytorch 2.5+ et torchao sont installés à l'aide de la commande ci-dessous.

uv pip install -U torch torchao

Chaque type de quantisation est disponible en tant qu'instance distincte de la classe AOBaseConfig. Cela offre des options de configuration plus flexibles en exposant un plus grand nombre d'arguments disponibles.

Passez le AOBaseConfig d'un dtype de quantification, comme Int4WeightOnlyConfig, à TorchAoConfig dans 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" quantise chaque couche sur le GPU au moment du chargement. C’est rapide, mais cela nécessite temporairement de la mémoire GPU supplémentaire pour les poids originaux et quantifiés. Si le modèle utilise déjà la majeure partie de votre mémoire GPU, le chargement peut échouer avec une erreur de mémoire insuffisante. Dans ce cas, supprimez device_map et déplacez le pipeline sur le GPU après le chargement :

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"

Sans device_map, Diffusers quantise les couches sur le CPU. C'est plus lent, mais cela évite le pic temporaire de mémoire GPU lors de la quantisation. Pour réduire encore plus l'utilisation de la mémoire GPU, utilisez plutôt enable_model_cpu_offload(). Vous pouvez également quantifier d'autres composants, tels que l'encodeur de texte.

torch.compile

torchao prend en charge torch.compile, qui peut accélérer l'inférence avec une seule ligne de code.

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)

Reportez-vous à ce tableau pour les benchmarks de vitesse d’inférence et d’utilisation de la mémoire avec Flux et CogVideoX. D’autres benchmarks sur divers matériels sont également disponibles dans le dépôt torchao.

Les schémas de quantification post-entraînement FP8 dans torchao sont efficaces pour les GPU ayant une capacité de calcul d'au moins 8.9 (RTX-4090, Hopper, etc.). FP8 offre souvent le meilleur compromis entre vitesse, mémoire et qualité lors de la génération d'images et de vidéos. Nous recommandons de combiner FP8 et torch.compile si votre GPU est compatible.

Types de quantification pris en charge

torchao prend en charge la quantification des poids uniquement ainsi que la quantification des poids et des activations dynamiques pour int8, float3-float8 et uint1-uint7.

La quantification des poids uniquement stocke les poids du modèle dans un type de données à faible nombre de bits, mais effectue les calculs avec un type de données de plus haute précision, comme bfloat16. Cela réduit les besoins en mémoire liés aux poids du modèle, mais conserve les pics de mémoire pour le calcul des activations.

La quantification dynamique des activations stocke les poids du modèle dans un type de données à faible nombre de bits, tout en quantifiant les activations à la volée pour économiser de la mémoire supplémentaire. Cela réduit les besoins en mémoire liés aux poids du modèle, ainsi que la surcharge mémoire liée aux calculs d'activation. Cependant, cela peut parfois entraîner un compromis en termes de qualité, il est donc recommandé de tester soigneusement différents modèles.

Reportez-vous à la documentation officielle de torchao pour une meilleure compréhension des méthodes de quantification disponibles. Une liste exhaustive des options de configuration est disponible ici.

Voici quelques exemples de configurations de quantification populaires :

Sérialisation et désérialisation des modèles quantifiés

Pour sérialiser un modèle quantifié dans un dtype donné, chargez d'abord le modèle avec le dtype de quantification souhaité, puis enregistrez-le à l'aide de la méthode 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)

Pour charger un modèle quantifié sérialisé, utilisez la méthode 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 vous utilisez torch<=2.6.0, certaines méthodes de quantification, telles que uint4 weight-only, ne peuvent pas être chargées directement et peuvent entraîner une UnpicklingError lors de la tentative de chargement des modèles, mais fonctionnent comme prévu lors de leur sauvegarde. Pour contourner ce problème, il est possible de charger manuellement le state dict dans le modèle. Notez toutefois que cela nécessite d’utiliser weights_only=False dans torch.load, il ne faut donc l’exécuter que si les poids proviennent d’une source de confiance.

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)

L'API AutoModel est prise en charge pour PyTorch >= 2.6, comme le montrent les exemples ci-dessous.

Ressources

• API de quantification TorchAO
• Exemples Diffusers-TorchAO