Hugging Face's logo

Diffusers documentation

ParaAttention

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

ParaAttention

Модели генерации изображений и видео большого размера, такие как FLUX.1-dev и HunyuanVideo, могут создавать проблемы при инференсе для приложений реального времени и развертывания из-за своих размеров.

ParaAttention — это библиотека, реализующая context parallelism и first block cache, которая может использоваться в сочетании с другими методами (torch.compile, fp8 dynamic quantization) для ускорения инференса.

В этом руководстве показано, как применить ParaAttention к FLUX.1-dev и HunyuanVideo на GPU NVIDIA L20. Для базового бенчмарка оптимизации не применяются, за исключением HunyuanVideo, чтобы избежать ошибок нехватки памяти.

Наш базовый бенчмарк показывает, что FLUX.1-dev способен генерировать изображение с разрешением 1024x1024 за 28 шагов за 26,36 секунды, а HunyuanVideo — 129 кадров с разрешением 720p за 30 шагов за 3675,71 секунды.

Для ещё более быстрого вывода с параллелизмом контекста попробуйте использовать GPU NVIDIA A100 или H100 (если они доступны) с поддержкой NVLink, особенно при большом количестве GPU.

Кэш первого блока

Кэширование выходов блоков модели transformers и их повторное использование на следующих шагах инференса снижает вычислительные затраты и ускоряет инференс.

Однако сложно определить, когда следует переиспользовать кэш, чтобы гарантировать качество генерируемых изображений или видео. ParaAttention напрямую использует остаточную разницу выхода первого блока трансформера для аппроксимации различий между выходами модели. Если разница достаточно мала, остаточная разница предыдущих шагов инференса переиспользуется. Иными словами, шаг денойзинга пропускается.

Это обеспечивает двукратное ускорение вывода для FLUX.1-dev и HunyuanVideo при очень хорошем качестве.

Cache in Diffusion Transformer
Как работает AdaCache, First Block Cache является его вариантом
FLUX-1.dev
HunyuanVideo

Чтобы применить кэш первого блока для FLUX.1-dev, вызовите apply_cache_on_pipe, как показано ниже. 0.08 — это значение по умолчанию для остаточной разницы в моделях FLUX.

import time
import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(pipe, residual_diff_threshold=0.08)

# Enable memory savings
# pipe.enable_model_cpu_offload()
# pipe.enable_sequential_cpu_offload()

begin = time.time()
image = pipe(
    "A cat holding a sign that says hello world",
    num_inference_steps=28,
).images[0]
end = time.time()
print(f"Time: {end - begin:.2f}s")

print("Saving image to flux.png")
image.save("flux.png")
ОптимизацииОригиналFBCache rdt=0.06FBCache rdt=0.08FBCache rdt=0.10FBCache rdt=0.12
ПредпросмотрOriginalFBCache rdt=0.06FBCache rdt=0.08FBCache rdt=0.10FBCache rdt=0.12
Время (с)26.3621.8317.0116.0013.78

First Block Cache сократил время вывода до 17,01 секунды по сравнению с базовой линией, что на 1,55 раза быстрее, при этом практически не теряя в качестве.

квантование fp8

fp8 с динамической квантизацией дополнительно ускоряет инференс и снижает потребление памяти. Для использования 8-битных NVIDIA Tensor Cores необходимо квантовать как активации, так и веса.

Используйте float8_weight_only и float8_dynamic_activation_float8_weight для квантования текстового энкодера и трансформерной модели.

Метод квантования по умолчанию — квантование на уровне тензора, но если ваша GPU поддерживает квантование по строкам, вы также можете попробовать его для повышения точности.

Установите torchao с помощью команды ниже.

pip3 install -U torch torchao

torch.compile с mode="max-autotune-no-cudagraphs" или mode="max-autotune" выбирает оптимальное ядро для производительности. Компиляция может занять много времени при первом вызове модели, но это окупается после завершения компиляции.

В этом примере квантуется только трансформерная модель, но вы также можете квантовать текстовый энкодер, чтобы ещё больше снизить потребление памяти.

Динамическое квантование может значительно изменить распределение выходных данных модели, поэтому для его вступления в силу необходимо увеличить значение residual_diff_threshold.

FLUX-1.dev
HunyuanVideo
import time
import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(
    pipe,
    residual_diff_threshold=0.12,  # Use a larger value to make the cache take effect
)

from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only

quantize_(pipe.text_encoder, float8_weight_only())
quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
pipe.transformer = torch.compile(
   pipe.transformer, mode="max-autotune-no-cudagraphs",
)

# Enable memory savings
# pipe.enable_model_cpu_offload()
# pipe.enable_sequential_cpu_offload()

for i in range(2):
    begin = time.time()
    image = pipe(
        "A cat holding a sign that says hello world",
        num_inference_steps=28,
    ).images[0]
    end = time.time()
    if i == 0:
        print(f"Warm up time: {end - begin:.2f}s")
    else:
        print(f"Time: {end - begin:.2f}s")

print("Saving image to flux.png")
image.save("flux.png")

Динамическое квантование fp8 и torch.compile сократили время инференса до 7,56 секунд по сравнению с базовой линией, что на 3,48 раза быстрее.

Параллелизм контекста

Context Parallelism параллелизует вывод и масштабируется на нескольких GPU. Композиционный дизайн ParaAttention позволяет сочетать Context Parallelism с First Block Cache и динамической квантизацией.

Подробные инструкции и примеры масштабирования вывода на нескольких GPU приведены в репозитории ParaAttention.

Если процесс вывода должен быть постоянным и доступным для обслуживания, рекомендуется использовать torch.multiprocessing для написания собственного обработчика вывода. Это позволяет избежать накладных расходов на запуск процесса, а также на загрузку и повторную компиляцию модели.

FLUX-1.dev
HunyuanVideo

Приведённый ниже пример кода объединяет First Block Cache, динамическое квантование fp8, torch.compile и Context Parallelism для достижения максимальной скорости вывода.

import time
import torch
import torch.distributed as dist
from diffusers import FluxPipeline

dist.init_process_group()

torch.cuda.set_device(dist.get_rank())

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.context_parallel import init_context_parallel_mesh
from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
from para_attn.parallel_vae.diffusers_adapters import parallelize_vae

mesh = init_context_parallel_mesh(
    pipe.device.type,
    max_ring_dim_size=2,
)
parallelize_pipe(
    pipe,
    mesh=mesh,
)
parallelize_vae(pipe.vae, mesh=mesh._flatten())

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(
    pipe,
    residual_diff_threshold=0.12,  # Use a larger value to make the cache take effect
)

from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only

quantize_(pipe.text_encoder, float8_weight_only())
quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
torch._inductor.config.reorder_for_compute_comm_overlap = True
pipe.transformer = torch.compile(
   pipe.transformer, mode="max-autotune-no-cudagraphs",
)

# Enable memory savings
# pipe.enable_model_cpu_offload(gpu_id=dist.get_rank())
# pipe.enable_sequential_cpu_offload(gpu_id=dist.get_rank())

for i in range(2):
    begin = time.time()
    image = pipe(
        "A cat holding a sign that says hello world",
        num_inference_steps=28,
        output_type="pil" if dist.get_rank() == 0 else "pt",
    ).images[0]
    end = time.time()
    if dist.get_rank() == 0:
        if i == 0:
            print(f"Warm up time: {end - begin:.2f}s")
        else:
            print(f"Time: {end - begin:.2f}s")

if dist.get_rank() == 0:
    print("Saving image to flux.png")
    image.save("flux.png")

dist.destroy_process_group()

Сохраните в run_flux.py и запустите с помощью torchrun.

# Use --nproc_per_node to specify the number of GPUs
torchrun --nproc_per_node=2 run_flux.py

Скорость инференса составляет 8,20 секунды по сравнению с базовой линией, что на 3,21 раза быстрее, при использовании 2 GPU NVIDIA L20. На 4 GPU L20 скорость инференса составляет 3,90 секунды, что на 6,75 раза быстрее.

Бенчмарки

FLUX-1.dev
HunyuanVideo
Тип GPUКоличество GPUОптимизацииВремя (с)Ускорение
NVIDIA L201Базовая линия26.361.00x
NVIDIA L201FBCache (rdt=0.08)17.011.55x
NVIDIA L201FP8 DQ13.401.96x
NVIDIA L201FBCache (rdt=0.12) + FP8 DQ7.563.48x
NVIDIA L202FBCache (rdt=0.12) + FP8 DQ + CP4.925.35x
NVIDIA L204FBCache (rdt=0.12) + FP8 DQ + CP3.906.75x
Update on GitHub

←CacheDiT Pruna→