# 量子化モデルのコンパイルとオフロード

モデルの最適化では、[推論速度](./fp16)と[メモリ使用量](./memory)のトレードオフが生じます。例えば、[キャッシュ](./cache)は推論速度を向上させますが、中間アテンション層の出力を保存する必要があるため、メモリ消費量も増加します。よりバランスの取れた最適化戦略としては、モデルの量子化、[torch.compile](./fp16#torchcompile)、およびさまざまな[オフロード手法](./memory#offloading)を組み合わせる方法があります。

> [!TIP]
> コンパイルの詳細や、ここでどのように適用できるかについては、[torch.compile](./fp16#torchcompile)のガイドを参照してください。例えば、地域別コンパイルを使用すると、速度向上を犠牲にすることなく、コンパイル時間を大幅に短縮できます。

画像生成において、量子化と[モデルのオフロード](./memory#model-offloading)を組み合わせることで、品質、速度、メモリ使用量のバランスを最適化できることが多いです。グループオフロードは画像生成ではあまり効果的ではありません。これは、計算カーネルがより早く完了した場合、データ転送を*完全に*オーバーラップさせることが通常できないためです。その結果、CPUとGPU間の通信オーバーヘッドが発生します。

動画生成では、動画モデルは計算量に制約されることが多いため、量子化と[グループオフローディング](./memory#group-offloading)を組み合わせる方が一般的に効果的です。

以下の表は、Flux における最適化戦略の組み合わせと、それらがレイテンシおよびメモリ使用量に与える影響を比較しています。

| 組み合わせ | レイテンシ (s) | メモリ使用量 (GB) |
|---|---|---|
| 量子化  | 32.602 | 14.9453 |
| 量子化、torch.compile  | 25.847 | 14.9448 |
| 量子化、torch.compile、モデルのCPUオフロード | 32.312 | 12.2369 |

これらの結果は、Flux を RTX 4090 でベンチマークしたものです。transformer と text_encoder のコンポーネントは量子化されています。ご自身のモデルを評価したい場合は、ベンチマーキングスクリプトを参照してください。

このガイドでは、[bitsandbytes](../quantization/bitsandbytes#torchcompile) を使用して量子化モデルをコンパイルおよびオフロードする方法を説明します。[PyTorch nightly](https://pytorch.org/get-started/locally/) と bitsandbytes の最新バージョンを使用していることを確認してください。

この例ではbitsandbytesを使用していますが、[TorchAO](../quantization/torchao)などの他の量子化バックエンドもこれらの機能をサポートしています。

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

## Quantization と torch.compile

まず、モデルを[量子化](../quantization/overview)してストレージに必要なメモリを削減し、[コンパイル](./fp16#torchcompile)して推論を高速化します。

bitsandbytes モデルをコンパイルする際に動的な出力を処理するため、[Dynamo](https://docs.pytorch.org/docs/stable/torch.compiler_dynamo_overview.html) の `capture_dynamic_output_shape_ops = True` を設定してください。

```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]
```

## 量子化、torch.compile、オフロード

量子化やtorch.compileに加えて、メモリ使用量をさらに削減する必要がある場合は、オフローディングを試してください。オフローディングでは、計算に必要なタイミングで、さまざまなレイヤーやモデルのコンポーネントをCPUからGPUへ移動します。

オフローディング時に[Dynamo](https://docs.pytorch.org/docs/stable/torch.compiler_dynamo_overview.html)の`cache_size_limit`を設定して過剰な再コンパイルを回避し、bitsandbytesモデルをコンパイルする際の動的な出力に対応するために`capture_dynamic_output_shape_ops = True`を設定します。

[モデルのCPUオフロード](./memory#model-offloading)では、計算に必要になったタイミングで、transformerモデルなどの個別のパイプラインコンポーネントをGPUに移動します。それ以外の時間は、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]
```

[グループオフローディング](./memory#group-offloading)は、transformerモデルのような個別のパイプラインコンポーネントの内部レイヤーをGPUに移動して計算を実行し、不要なタイミングでオフロードします。同時に、[CUDA stream](./memory#cuda-stream)機能を利用して、次の実行対象のレイヤーを事前取得（プリフェッチ）します。

計算とデータ転送をオーバーラップさせることで、モデルのCPUオフロードよりも高速であり、かつメモリを節約できます。

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