Diffusers documentation
Компиляция и выгрузка квантованных моделей
Компиляция и выгрузка квантованных моделей
Оптимизация моделей часто связана с компромиссом между скоростью вывода и потреблением памяти. Например, хотя кэширование может ускорить вывод, оно также увеличивает потребление памяти, поскольку требует хранения выходов промежуточных слоев внимания. Более сбалансированная стратегия оптимизации сочетает квантование модели, torch.compile и различные методы выгрузки.
Ознакомьтесь с руководством по torch.compile, чтобы узнать больше о компиляции и о том, как её можно применить здесь. Например, региональная компиляция может значительно сократить время компиляции, не теряя при этом в ускорении.
Для генерации изображений сочетание квантования и выгрузки модели часто обеспечивает наилучший баланс между качеством, скоростью и потреблением памяти. Групповая выгрузка менее эффективна для генерации изображений, поскольку обычно невозможно полностью перекрыть передачу данных, если ядро вычислений завершается быстрее. Это приводит к некоторому накладному расходу на коммуникацию между CPU и GPU.
Для генерации видео сочетание квантования и group-offloading обычно даёт лучший результат, поскольку модели для видео более ограничены по вычислительным мощностям.
Таблица ниже содержит сравнение комбинаций стратегий оптимизации и их влияния на задержку и потребление памяти для Flux.
| комбинация | задержка (с) | потребление памяти (ГБ) |
|---|---|---|
| квантование | 32.602 | 14.9453 |
| квантование, torch.compile | 25.847 | 14.9448 |
| квантование, torch.compile, выгрузка модели на CPU | 32.312 | 12.2369 |
В этом руководстве показано, как скомпилировать и выгрузить квантованную модель с помощью bitsandbytes. Убедитесь, что вы используете PyTorch nightly и последнюю версию bitsandbytes.
Хотя в этом примере мы используем bitsandbytes, другие бэкенды квантования, такие как TorchAO, также поддерживают эти функции.
pip install -U bitsandbytes
Квантование и torch.compile
Начните с квантования модели для сокращения объёма памяти, необходимого для её хранения, и компиляции для ускорения вывода.
Настройте Dynamo с параметром capture_dynamic_output_shape_ops = True, чтобы обрабатывать динамические выходные данные при компиляции моделей 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]Квантование, torch.compile и выгрузка
Помимо квантования и torch.compile, попробуйте offloading, если вам нужно дополнительно снизить потребление памяти. Offloading перемещает различные слои или компоненты модели с CPU на GPU по мере необходимости для вычислений.
Настройте Dynamo cache_size_limit при offloading, чтобы избежать чрезмерной перекомпиляции, и установите capture_dynamic_output_shape_ops = True для обработки динамических выходов при компиляции моделей bitsandbytes.
Перемещение модели на 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 перемещает внутренние слои отдельного компонента конвейера, например трансформерной модели, на 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)