Diffusers documentation
Biên dịch và chuyển các mô hình lượng tử
Biên dịch và chuyển các mô hình lượng tử
Việc tối ưu hóa mô hình thường liên quan đến sự đánh đổi giữa tốc độ suy luận và lượng bộ nhớ sử dụng. Ví dụ, trong khi lưu trữ đệm (caching) có thể tăng tốc độ suy luận, nó cũng làm tăng mức tiêu thụ bộ nhớ vì cần lưu trữ đầu ra của các lớp chú ý trung gian. Một chiến lược tối ưu hóa cân bằng hơn kết hợp việc lượng tử hóa mô hình, torch.compile và các phương pháp offloading khác nhau.
Xem hướng dẫn torch.compile để tìm hiểu thêm về biên dịch và cách áp dụng chúng ở đây. Ví dụ, biên dịch theo vùng có thể giảm đáng kể thời gian biên dịch mà không đánh đổi bất kỳ tốc độ tăng tốc nào.
Đối với việc tạo ảnh, việc kết hợp lượng tử hóa và chuyển mô hình thường mang lại sự cân bằng tốt nhất giữa chất lượng, tốc độ và bộ nhớ. Việc chuyển theo nhóm không hiệu quả bằng đối với việc tạo ảnh vì thường không thể trùng lặp hoàn toàn việc truyền dữ liệu nếu nhân tính toán kết thúc nhanh hơn. Điều này dẫn đến một số chi phí truyền thông giữa CPU và GPU.
Đối với việc tạo video, việc kết hợp lượng tử hóa và group-offloading thường cho kết quả tốt hơn vì các mô hình video phụ thuộc nhiều hơn vào khả năng tính toán.
Bảng dưới đây cung cấp so sánh các tổ hợp chiến lược tối ưu hóa và tác động của chúng đến độ trễ và mức sử dụng bộ nhớ đối với Flux.
| tổ hợp | độ trễ (giây) | lượng bộ nhớ sử dụng (GB) |
|---|---|---|
| lượng tử hóa | 32.602 | 14.9453 |
| quantization, torch.compile | 25.847 | 14.9448 |
| lượng tử hóa, torch.compile, chuyển mô hình ra CPU | 32.312 | 12.2369 |
Hướng dẫn này sẽ cho bạn thấy cách biên dịch và offload một mô hình đã lượng tử hóa bằng bitsandbytes. Hãy đảm bảo bạn đang sử dụng PyTorch nightly và phiên bản mới nhất của bitsandbytes.
Mặc dù chúng ta sử dụng bitsandbytes trong ví dụ này, các backend lượng tử hóa khác như TorchAO cũng hỗ trợ các tính năng này.
pip install -U bitsandbytes
Quantization và torch.compile
Bắt đầu bằng cách lượng tử hóa một mô hình để giảm bộ nhớ cần thiết cho việc lưu trữ và biên dịch nó để tăng tốc suy luận.
Cấu hình Dynamo capture_dynamic_output_shape_ops = True để xử lý các đầu ra động khi biên dịch các mô hình 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]Quantization, torch.compile, và offloading
Ngoài lượng tử hóa và torch.compile, hãy thử offloading nếu bạn cần giảm thêm mức sử dụng bộ nhớ. Offloading di chuyển các lớp hoặc thành phần mô hình khác nhau từ CPU sang GPU khi cần cho các phép tính.
Cấu hình Dynamo cache_size_limit trong quá trình offloading để tránh biên dịch lại quá mức và đặt capture_dynamic_output_shape_ops = True để xử lý các đầu ra động khi biên dịch các mô hình bitsandbytes.
Model CPU offloading di chuyển một thành phần pipeline riêng lẻ, chẳng hạn như mô hình transformer, sang GPU khi cần thiết cho việc tính toán. Ngược lại, thành phần đó sẽ được offload sang 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 di chuyển các lớp bên trong của một thành phần pipeline riêng lẻ, chẳng hạn như mô hình transformer, sang GPU để tính toán và chuyển nó ra ngoài khi không còn cần thiết. Đồng thời, tính năng này sử dụng CUDA stream để nạp trước lớp tiếp theo để thực thi.
Bằng cách chồng chéo giữa tính toán và truyền dữ liệu, phương pháp này nhanh hơn so với việc offload mô hình ra CPU, đồng thời cũng tiết kiệm bộ nhớ.
# 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)