양자화 모델 컴파일 및 오프로딩

모델 최적화는 추론 속도와 메모리 사용량 사이의 트레이드오프를 수반하는 경우가 많습니다. 예를 들어, 캐싱은 추론 속도를 높일 수 있지만, 중간 어텐션 레이어의 출력을 저장해야 하므로 메모리 사용량도 증가시킵니다. 보다 균형 잡힌 최적화 전략은 모델 양자화, torch.compile 및 다양한 오프로딩 방법을 결합하는 것입니다.

컴파일링 및 여기에서 어떻게 적용할 수 있는지에 대해 자세히 알아보려면 torch.compile 가이드를 확인하세요. 예를 들어, 지역적 컴파일링은 속도 향상을 포기하지 않으면서도 컴파일링 시간을 크게 줄일 수 있습니다.

이미지 생성의 경우, 양자화와 모델 오프로딩을 결합하면 품질, 속도, 메모리 간의 최적의 트레이드오프를 제공하는 경우가 많습니다. 컴퓨트 커널이 더 빠르게 종료되면 데이터 전송을 완전히 오버랩하는 것이 일반적으로 불가능하기 때문에, 그룹 오프로딩은 이미지 생성에 그다지 효과적이지 않습니다. 이로 인해 CPU와 GPU 사이에 일부 통신 오버헤드가 발생합니다.

비디오 생성의 경우, 비디오 모델은 컴퓨트 바운드(compute-bound) 특성이 더 강하기 때문에 양자화와 group-offloading을 결합하는 것이 일반적으로 더 좋습니다.

아래 표는 Flux의 지연 시간과 메모리 사용량에 대한 최적화 전략 조합의 비교와 그 영향을 제공합니다.

조합	지연 시간 (초)	메모리 사용량 (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

양자화 및 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 오프로딩은 트랜스포머 모델과 같은 개별 파이프라인 구성 요소를 계산이 필요할 때 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]

Group offloading은 transformer 모델과 같은 개별 파이프라인 컴포넌트의 내부 레이어를 계산 시 GPU로 이동하고, 필요하지 않을 때 offload합니다. 동시에 CUDA stream 기능을 사용하여 다음 실행을 위한 레이어를 미리 가져옵니다(prefetch).

계산과 데이터 전송을 오버랩하여 모델 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