Kantize modelleri derleme ve aktarma

Model optimizasyonu genellikle çalıştırma hızı ile Bellek kullanımı arasında ödünleşimler gerektirir. Örneğin, önbelleğe alma çalıştırma hızını artırabilir ancak ara dikkat katmanlarının çıktılarını depolaması gerektiğinden bellek tüketimini de artırır. Daha dengeli bir optimizasyon stratejisi, modelin kuantizlenmesi, torch.compile ve çeşitli aktarma yöntemleri ile birleştirilir.

Derleme ve bu derlemelerin burada nasıl uygulanabileceğini öğrenmek için torch.compile kılavuzunu inceleyin. Örneğin, bölgesel derleme, hızlanmalardan vazgeçmeden derleme süresini önemli ölçüde azaltabilir.

Görsel üretiminde, kuantizasyonun ve model offloading'in birleştirilmesi genellikle kalite, hız ve bellek arasında en iyi dengeyi sağlar. Grup offloading, hesaplama çekirdeği daha hızlı bittiyse veri aktarımının tamamen örtüşmesi genellikle mümkün olmadığından görsel üretimde o kadar etkili değildir. Bu durum, CPU ile GPU arasında bazı iletişim maliyetlerine yol açar.

Video üretimi için, video modelleri daha çok hesaplama sınırlı olduğundan, kuantizasyonun ve group-offloading'in birleştirilmesi genellikle daha iyi sonuç verir.

Aşağıdaki tablo, Flux için optimizasyon stratejisi kombinasyonlarını ve bunların gecikme ile bellek kullanımı üzerindeki etkilerini karşılaştırmaktadır.

kombinasyon	gecikme (s)	bellek-kullanımı (GB)
quantization	32.602	14.9453
quantization, torch.compile	25.847	14.9448
quantization, torch.compile, model CPU offloading	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.

Bu kılavuz, bitsandbytes ile bir kuantize modelin nasıl derleneceğini ve offload edileceğini gösterecektir. PyTorch nightly ve bitsandbytes'in en son sürümünü kullandığınızdan emin olun.

Bu örnekte bitsandbytes kullanıyoruz, ancak TorchAO gibi diğer kuantizasyon altyapıları da bu özellikleri destekler.

pip install -U bitsandbytes

Quantization ve torch.compile

Bir modelin bellek gereksinimini azaltmak için kantizasyonunu yapıp çıkarımı hızlandırmak için derlemesini yaparak başlayın.

bitsandbytes modellerini derlerken dinamik çıktıları işlemek için Dynamo yapılandırmasını capture_dynamic_output_shape_ops = True olarak ayarlayın.

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]

Kantizasyon, torch.compile ve offloading

Kantizasyon ve torch.compile'in yanı sıra, bellek kullanımını daha da azaltmanız gerekiyorsa offloading'i deneyin. Offloading, hesaplama gereksinimlerine göre çeşitli katmanları veya model bileşenlerini CPU'dan GPU'ya taşır.

Offloading sırasında aşırı yeniden derlemeyi önlemek için Dynamo cache_size_limit değerini yapılandırın ve bitsandbytes modellerini derlerken dinamik çıktıları işlemek için capture_dynamic_output_shape_ops = True ayarlayın.

model CPU'ya aktarımı

grup boşaltma

Model CPU offloading, transformer modeli gibi bireysel bir pipeline bileşenini, hesaplama için gerektiğinde GPU'ya taşır. Aksi takdirde, bileşen CPU'ya aktarılır.

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]

Grup dışa aktarma, transformer modeli gibi tek bir pipeline bileşeninin iç katmanlarını hesaplama için GPU'ya taşır ve gerekmediğinde dışa aktarır. Aynı zamanda, bir sonraki katmanı çalıştırma için önceden yüklemek üzere CUDA stream özelliğini kullanır.

Hesaplama ve veri aktarımını üst üste bindirerek, model CPU dışa aktarımından daha hızlıdır ve aynı zamanda bellek tasarrufu sağlar.

# 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