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

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

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

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

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

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

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

These results are benchmarked on Flux with a RTX 4090. The transformer and text_encoder components are quantized. Refer to the benchmarking script if you're interested in evaluating your own model.

このガイドでは、bitsandbytes を使用して量子化モデルをコンパイルおよびオフロードする方法を説明します。PyTorch nightly と bitsandbytes の最新バージョンを使用していることを確認してください。

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

pip install -U bitsandbytes

Quantization と torch.compile

まず、モデルを量子化してストレージに必要なメモリを削減し、コンパイルして推論を高速化します。

bitsandbytes モデルをコンパイルする際に動的な出力を処理するため、Dynamo の capture_dynamic_output_shape_ops = True を設定してください。

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のcache_size_limitを設定して過剰な再コンパイルを回避し、bitsandbytesモデルをコンパイルする際の動的な出力に対応するためにcapture_dynamic_output_shape_ops = Trueを設定します。

モデルのCPUオフロード

グループオフローディング

モデルのCPUオフロードでは、計算に必要になったタイミングで、transformerモデルなどの個別のパイプラインコンポーネントをGPUに移動します。それ以外の時間は、CPUにオフロードされます。

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]

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

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

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

Update on GitHub