# ParaAttention

    

    

Модели генерации изображений и видео большого размера, такие как [FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) и [HunyuanVideo](https://huggingface.co/tencent/HunyuanVideo), могут создавать проблемы при инференсе для приложений реального времени и развертывания из-за своих размеров.

[ParaAttention](https://github.com/chengzeyi/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 секунды.

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

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

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

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

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

    
    Как работает AdaCache, First Block Cache является его вариантом

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

```python
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.06 | FBCache rdt=0.08 | FBCache rdt=0.10 | FBCache rdt=0.12 |
| - | - | - | - | - | - |
| Предпросмотр | ![Original](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-original.png) | ![FBCache rdt=0.06](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-fbc-0.06.png) | ![FBCache rdt=0.08](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-fbc-0.08.png) | ![FBCache rdt=0.10](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-fbc-0.10.png) | ![FBCache rdt=0.12](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-fbc-0.12.png) |
| Время (с) | 26.36 | 21.83 | 17.01 | 16.00 | 13.78 |

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

Чтобы применить First Block Cache для HunyuanVideo, используйте `apply_cache_on_pipe`, как показано ниже. 0.06 — это значение по умолчанию для остаточной разницы в моделях HunyuanVideo.

```python
import time
import torch
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video

model_id = "tencent/HunyuanVideo"
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    dtype=torch.bfloat16,
    revision="refs/pr/18",
)
pipe = HunyuanVideoPipeline.from_pretrained(
    model_id,
    transformer=transformer,
    dtype=torch.float16,
    revision="refs/pr/18",
).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.6)

pipe.vae.enable_tiling()

begin = time.time()
output = pipe(
    prompt="A cat walks on the grass, realistic",
    height=720,
    width=1280,
    num_frames=129,
    num_inference_steps=30,
).frames[0]
end = time.time()
print(f"Time: {end - begin:.2f}s")

print("Saving video to hunyuan_video.mp4")
export_to_video(output, "hunyuan_video.mp4", fps=15)
```

  
  Ваш браузер не поддерживает тег видео.

 HunyuanVideo без FBCache

  
  Ваш браузер не поддерживает тег видео.

 HunyuanVideo с FBCache

First Block Cache сократил время инференса до 2271.06 секунд по сравнению с базовой линией, то есть ускорил его в 1.62 раза, при этом практически не теряя в качестве.

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

fp8 с динамической квантизацией дополнительно ускоряет инференс и снижает потребление памяти. Для использования 8-битных [NVIDIA Tensor Cores](https://www.nvidia.com/en-us/data-center/tensor-cores/) необходимо квантовать как активации, так и веса.

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

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

Установите [torchao](https://github.com/pytorch/ao/tree/main) с помощью команды ниже.

```bash
pip3 install -U torch torchao
```

[torch.compile](https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html) с `mode="max-autotune-no-cudagraphs"` или `mode="max-autotune"` выбирает оптимальное ядро для производительности. Компиляция может занять много времени при первом вызове модели, но это окупается после завершения компиляции.

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

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

```python
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 раза быстрее.

```python
import time
import torch
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video

model_id = "tencent/HunyuanVideo"
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    dtype=torch.bfloat16,
    revision="refs/pr/18",
)
pipe = HunyuanVideoPipeline.from_pretrained(
    model_id,
    transformer=transformer,
    dtype=torch.float16,
    revision="refs/pr/18",
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(pipe)

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.vae.enable_tiling()
# pipe.enable_model_cpu_offload()
# pipe.enable_sequential_cpu_offload()

for i in range(2):
    begin = time.time()
    output = pipe(
        prompt="A cat walks on the grass, realistic",
        height=720,
        width=1280,
        num_frames=129,
        num_inference_steps=1 if i == 0 else 30,
    ).frames[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 video to hunyuan_video.mp4")
export_to_video(output, "hunyuan_video.mp4", fps=15)
```

GPU NVIDIA L20 имеет только 48 ГБ памяти и может сталкиваться с ошибками нехватки памяти (OOM) после компиляции, если не вызывается `enable_model_cpu_offload`, поскольку HunyuanVideo при работе с высоким разрешением и большим количеством кадров формирует очень большие тензоры активаций. Для GPU с памятью менее 80 ГБ можно попробовать уменьшить разрешение и количество кадров, чтобы избежать ошибок OOM.

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

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

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

> [!TIP]
> Подробные инструкции и примеры масштабирования вывода на нескольких GPU приведены в репозитории [ParaAttention](https://github.com/chengzeyi/ParaAttention/tree/main).

Если процесс вывода должен быть постоянным и доступным для обслуживания, рекомендуется использовать [torch.multiprocessing](https://pytorch.org/docs/stable/multiprocessing.html) для написания собственного обработчика вывода. Это позволяет избежать накладных расходов на запуск процесса, а также на загрузку и повторную компиляцию модели.

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

```python
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](https://pytorch.org/docs/stable/elastic/run.html).

```bash
# 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 раза быстрее.

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

```python
import time
import torch
import torch.distributed as dist
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video

dist.init_process_group()

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

model_id = "tencent/HunyuanVideo"
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    dtype=torch.bfloat16,
    revision="refs/pr/18",
)
pipe = HunyuanVideoPipeline.from_pretrained(
    model_id,
    transformer=transformer,
    dtype=torch.float16,
    revision="refs/pr/18",
).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,
)
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)

# from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only
#
# torch._inductor.config.reorder_for_compute_comm_overlap = True
#
# 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.vae.enable_tiling()
# 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()
    output = pipe(
        prompt="A cat walks on the grass, realistic",
        height=720,
        width=1280,
        num_frames=129,
        num_inference_steps=1 if i == 0 else 30,
        output_type="pil" if dist.get_rank() == 0 else "pt",
    ).frames[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 video to hunyuan_video.mp4")
    export_to_video(output, "hunyuan_video.mp4", fps=15)

dist.destroy_process_group()
```

Сохраните в `run_hunyuan_video.py` и запустите с помощью [torchrun](https://pytorch.org/docs/stable/elastic/run.html).

```bash
# Use --nproc_per_node to specify the number of GPUs
torchrun --nproc_per_node=8 run_hunyuan_video.py
```

Скорость инференса составляет 649,23 секунды по сравнению с базовой линией, что на 5,66 раза быстрее, при использовании 8 GPU NVIDIA L20.

## Бенчмарки

| Тип GPU | Количество GPU | Оптимизации | Время (с) | Ускорение |
| - | - | - | - | - |
| NVIDIA L20 | 1 | Базовая линия | 26.36 | 1.00x |
| NVIDIA L20 | 1 | FBCache (rdt=0.08) | 17.01 | 1.55x |
| NVIDIA L20 | 1 | FP8 DQ | 13.40 | 1.96x |
| NVIDIA L20 | 1 | FBCache (rdt=0.12) + FP8 DQ | 7.56 | 3.48x |
| NVIDIA L20 | 2 | FBCache (rdt=0.12) + FP8 DQ + CP | 4.92 | 5.35x |
| NVIDIA L20 | 4 | FBCache (rdt=0.12) + FP8 DQ + CP | 3.90 | 6.75x |

| Тип GPU | Количество GPU | Оптимизации | Время (с) | Ускорение |
| - | - | - | - | - |
| NVIDIA L20 | 1 | Базовая линия | 3675.71 | 1.00x |
| NVIDIA L20 | 1 | FBCache | 2271.06 | 1.62x |
| NVIDIA L20 | 2 | FBCache + CP | 1132.90 | 3.24x |
| NVIDIA L20 | 4 | FBCache + CP | 718.15 | 5.12x |
| NVIDIA L20 | 8 | FBCache + CP | 649.23 | 5.66x |