# Compilation et offloading de modèles quantifiés

L’optimisation des modèles implique souvent des compromis entre la [vitesse d’inférence](./fp16) et la [consommation mémoire](./memory). Par exemple, bien que le [cache](./cache) puisse accélérer l’inférence, il augmente également la consommation mémoire car il doit stocker les sorties des couches d’attention intermédiaires. Une stratégie d’optimisation plus équilibrée combine la quantification d’un modèle, [torch.compile](./fp16#torchcompile) et diverses [méthodes de déchargement](./memory#offloading).

> [!TIP]
> Consultez le guide [torch.compile](./fp16#torchcompile) pour en savoir plus sur la compilation et sur la manière dont elle peut être appliquée ici. Par exemple, la compilation régionale peut réduire considérablement le temps de compilation sans renoncer aux gains de vitesse.

Pour la génération d'images, la combinaison de la quantification et du [déchargement du modèle](./memory#model-offloading) offre souvent le meilleur compromis entre qualité, vitesse et mémoire. Le déchargement par groupes est moins efficace pour la génération d'images, car il n'est généralement pas possible de *complètement* superposer le transfert de données si le noyau de calcul se termine plus rapidement. Cela entraîne un certain surcoût de communication entre le CPU et le GPU.

Pour la génération de vidéos, combiner la quantification et le [group-offloading](./memory#group-offloading) est généralement plus efficace, car les modèles de vidéo sont plus limités par la puissance de calcul.

Le tableau ci-dessous fournit une comparaison des combinaisons de stratégies d’optimisation et de leur impact sur la latence et l’utilisation de la mémoire pour Flux.

| combinaison | latence (s) | utilisation de la mémoire (Go) |
|---|---|---|
| quantification  | 32.602 | 14.9453 |
| quantization, torch.compile  | 25.847 | 14.9448 |
| quantization, torch.compile, offloading du modèle sur CPU | 32.312 | 12.2369 |

Ces résultats sont mesurés sur Flux avec une RTX 4090. Les composants transformer et text_encoder sont quantifiés. Consultez le script de benchmark si vous souhaitez évaluer votre propre modèle.

Ce guide vous montre comment compiler et décharger un modèle quantifié avec [bitsandbytes](../quantization/bitsandbytes#torchcompile). Assurez-vous d'utiliser [PyTorch nightly](https://pytorch.org/get-started/locally/) et la dernière version de bitsandbytes.

Bien que nous utilisions bitsandbytes dans cet exemple, d'autres backends de quantification tels que [TorchAO](../quantization/torchao) prennent également en charge ces fonctionnalités.

```bash
pip install -U bitsandbytes
```

## Quantification et torch.compile

Commencez par [quantiser](../quantization/overview) un modèle pour réduire la mémoire nécessaire au stockage, puis [compiler](./fp16#torchcompile) ce modèle pour accélérer l'inférence.

Configurez [Dynamo](https://docs.pytorch.org/docs/stable/torch.compiler_dynamo_overview.html) avec `capture_dynamic_output_shape_ops = True` pour gérer les sorties dynamiques lors de la compilation des modèles bitsandbytes.

```py
import torch
from diffusers import DiffusionPipeline
from diffusers.quantizers import PipelineQuantizationConfig

torch._dynamo.config.capture_dynamic_output_shape_ops = True

# quantize
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"],
)
pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

# compile
pipeline.transformer.to(memory_format=torch.channels_last)
pipeline.transformer.compile(mode="max-autotune", fullgraph=True)
pipeline("""
    cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
    highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""
).images[0]
```

## Quantification, torch.compile et offloading

En plus de la quantification et de torch.compile, essayez le déchargement (offloading) si vous devez réduire davantage l'utilisation de la mémoire. Le déchargement déplace diverses couches ou composants du modèle du CPU vers le GPU au fur et à mesure des besoins de calcul.

Configurez `cache_size_limit` de [Dynamo](https://docs.pytorch.org/docs/stable/torch.compiler_dynamo_overview.html) lors du déchargement pour éviter une recompilation excessive et définissez `capture_dynamic_output_shape_ops = True` pour gérer les sorties dynamiques lors de la compilation des modèles bitsandbytes.

[Le déchargement du modèle sur CPU](./memory#model-offloading) déplace un composant individuel du pipeline, comme le modèle transformer, sur le GPU lorsqu'il est nécessaire pour le calcul. Sinon, il est déchargé sur le CPU.

```py
import torch
from diffusers import DiffusionPipeline
from diffusers.quantizers import PipelineQuantizationConfig

torch._dynamo.config.cache_size_limit = 1000
torch._dynamo.config.capture_dynamic_output_shape_ops = True

# quantize
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"],
)
pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

# model CPU offloading
pipeline.enable_model_cpu_offload()

# compile
pipeline.transformer.compile()
pipeline(
    "cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California, highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain"
).images[0]
```

[Group offloading](./memory#group-offloading) déplace les couches internes d'un composant individuel du pipeline, comme le modèle transformer, vers le GPU pour le calcul, puis les décharge lorsqu'elles ne sont plus nécessaires. En même temps, il utilise la fonctionnalité [CUDA stream](./memory#cuda-stream) pour précharger la couche suivante en vue de son exécution.

En superposant le calcul et le transfert de données, il est plus rapide que le déchargement du modèle sur CPU tout en économisant de la mémoire.

```py
# pip install ftfy
import torch
from diffusers import AutoModel, DiffusionPipeline
from diffusers.hooks import apply_group_offloading
from diffusers.utils import export_to_video
from diffusers.quantizers import PipelineQuantizationConfig
from transformers import UMT5EncoderModel

torch._dynamo.config.cache_size_limit = 1000
torch._dynamo.config.capture_dynamic_output_shape_ops = True

# quantize
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"],
)

text_encoder = UMT5EncoderModel.from_pretrained(
    "Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="text_encoder", dtype=torch.bfloat16
)
pipeline = DiffusionPipeline.from_pretrained(
    "Wan-AI/Wan2.1-T2V-14B-Diffusers",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

# group offloading
onload_device = torch.device("cuda")
offload_device = torch.device("cpu")

pipeline.transformer.enable_group_offload(
    onload_device=onload_device,
    offload_device=offload_device,
    offload_type="leaf_level",
    use_stream=True,
    non_blocking=True
)
pipeline.vae.enable_group_offload(
    onload_device=onload_device,
    offload_device=offload_device,
    offload_type="leaf_level",
    use_stream=True,
    non_blocking=True
)
apply_group_offloading(
    pipeline.text_encoder,
    onload_device=onload_device,
    offload_type="leaf_level",
    use_stream=True,
    non_blocking=True
)

# compile
pipeline.transformer.compile()

prompt = """
The camera rushes from far to near in a low-angle shot, 
revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in 
for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground. 
Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts dynamic 
shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
"""
negative_prompt = """
Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, 
low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, 
misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards
"""

output = pipeline(
    prompt=prompt,
    negative_prompt=negative_prompt,
    num_frames=81,
    guidance_scale=5.0,
).frames[0]
export_to_video(output, "output.mp4", fps=16)
```