Mengompilasi dan meng-offload model terkuantisasi

Optimasi model sering kali melibatkan pertukaran antara kecepatan inferensi dan penggunaan memori. Misalnya, meskipun penyimpanan sementara (caching) dapat meningkatkan kecepatan inferensi, hal ini juga meningkatkan konsumsi memori karena perlu menyimpan keluaran dari lapisan perhatian perantara. Strategi optimasi yang lebih seimbang menggabungkan kuantisasi model, torch.compile, dan berbagai metode offloading.

Periksa panduan torch.compile untuk mempelajari lebih lanjut tentang kompilasi dan bagaimana hal tersebut dapat diterapkan di sini. Sebagai contoh, kompilasi regional dapat secara signifikan mengurangi waktu kompilasi tanpa mengorbankan percepatan apa pun.

Untuk generasi gambar, menggabungkan kuantisasi dan model offloading sering kali memberikan trade-off terbaik antara kualitas, kecepatan, dan memori. Group offloading tidak seefektif itu untuk generasi gambar karena biasanya tidak mungkin untuk sepenuhnya menumpuk transfer data jika kernel komputasi selesai lebih cepat. Hal ini menyebabkan beberapa overhead komunikasi antara CPU dan GPU.

Untuk generasi video, menggabungkan kuantisasi dan group-offloading cenderung lebih baik karena model video lebih terikat pada komputasi.

Tabel di bawah ini menyajikan perbandingan kombinasi strategi optimasi serta dampaknya terhadap latensi dan penggunaan memori untuk Flux.

kombinasi	latensi (s)	penggunaan- memori (GB)
kuantisasi	32.602	14.9453
quantization, torch.compile	25.847	14.9448
kuantisasi, torch.compile, offloading model ke 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.

Panduan ini akan menunjukkan cara mengompilasi dan meng-offload model terkuantisasi dengan bitsandbytes. Pastikan Anda menggunakan PyTorch nightly dan versi terbaru dari bitsandbytes.

Meskipun kami menggunakan bitsandbytes dalam contoh ini, backend kuantisasi lain seperti TorchAO juga mendukung fitur-fitur ini.

pip install -U bitsandbytes

Kuantisasi dan torch.compile

Mulai dengan mengkuantisasi model untuk mengurangi memori yang dibutuhkan untuk penyimpanan dan mengompilasi model tersebut untuk mempercepat inferensi.

Konfigurasikan Dynamo dengan capture_dynamic_output_shape_ops = True untuk menangani output dinamis saat mengompilasi model bitsandbytes.

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]

Kuantisasi, torch.compile, dan offloading

Selain kuantisasi dan torch.compile, coba offloading jika Anda perlu mengurangi penggunaan memori lebih lanjut. Offloading memindahkan berbagai lapisan atau komponen model dari CPU ke GPU sesuai kebutuhan untuk komputasi.

Konfigurasikan Dynamo cache_size_limit saat offloading untuk menghindari kompilasi ulang yang berlebihan dan atur capture_dynamic_output_shape_ops = True untuk menangani output dinamis saat mengompilasi model bitsandbytes.

offloading model CPU

offloading grup

Model CPU offloading memindahkan komponen pipeline individual, seperti model transformer, ke GPU saat dibutuhkan untuk komputasi. Jika tidak, komponen tersebut di-offload ke 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]

Group offloading memindahkan lapisan internal dari komponen pipeline individual, seperti model transformer, ke GPU untuk komputasi dan melakukan offload saat tidak diperlukan. Pada saat yang sama, fitur CUDA stream digunakan untuk mengambil lapisan berikutnya secara prefetch untuk dieksekusi.

Dengan menumpuk komputasi dan transfer data, metode ini lebih cepat daripada offloading model ke CPU sekaligus menghemat memori.

# 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