Buckets:

mishig's picture
|
download
raw
17.6 kB
# ParaAttention
Modelos grandes de geração de imagem e vídeo, como [FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) e [HunyuanVideo](https://huggingface.co/tencent/HunyuanVideo), podem representar um desafio de inferência para aplicações em tempo real e implantação devido ao seu tamanho.
[ParaAttention](https://github.com/chengzeyi/ParaAttention) é uma biblioteca que implementa **paralelismo de contexto** e **cache do primeiro bloco** e pode ser combinada com outras técnicas (torch.compile, quantização dinâmica fp8) para acelerar a inferência.
Este guia mostrará como aplicar o ParaAttention ao FLUX.1-dev e ao HunyuanVideo em GPUs NVIDIA L20. Nenhuma otimização é aplicada em nosso benchmark de referência, exceto para o HunyuanVideo, a fim de evitar erros de memória insuficiente.
Nosso benchmark de referência mostra que o FLUX.1-dev é capaz de gerar uma imagem com resolução de 1024x1024 em 28 passos em 26,36 segundos, e o HunyuanVideo é capaz de gerar 129 quadros em resolução 720p em 30 passos em 3675,71 segundos.
> [!TIP]
> Para inferência ainda mais rápida com paralelismo de contexto, tente usar GPUs NVIDIA A100 ou H100 (se disponíveis) com suporte a NVLink, especialmente quando há um grande número de GPUs.
## Cache do Primeiro Bloco
O cache da saída dos blocos do transformers no modelo e a reutilização deles nas etapas de inferência subsequentes reduzem o custo de computação e tornam a inferência mais rápida.
No entanto, é difícil decidir quando reutilizar o cache para garantir a qualidade das imagens ou vídeos gerados. O ParaAttention usa diretamente a **diferença residual da saída do primeiro bloco do transformer** para aproximar a diferença entre as saídas do modelo. Quando a diferença é pequena o suficiente, a diferença residual das etapas de inferência anteriores é reutilizada. Em outras palavras, a etapa de desruido é ignorada.
Isso alcança uma aceleração de 2x na inferência do FLUX.1-dev e do HunyuanVideo, com muito boa qualidade.
Como o AdaCache funciona, o First Block Cache é uma variante dele
Para aplicar o cache do primeiro bloco no FLUX.1-dev, chame `apply_cache_on_pipe` conforme mostrado abaixo. 0,08 é o valor padrão de diferença residual para os modelos 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")
```
| Otimizações | Original | FBCache rdt=0.06 | FBCache rdt=0.08 | FBCache rdt=0.10 | FBCache rdt=0.12 |
| - | - | - | - | - | - |
| Pré-visualização | ![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) |
| Tempo Real (s) | 26.36 | 21.83 | 17.01 | 16.00 | 13.78 |
O First Block Cache reduziu a velocidade de inferência para 17,01 segundos em comparação com a linha de base, ou seja, 1,55x mais rápido, mantendo uma perda de qualidade quase nula.
Para aplicar o First Block Cache no HunyuanVideo, use `apply_cache_on_pipe` conforme mostrado abaixo. 0,06 é o valor padrão de diferença residual para os modelos 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)
```
Seu navegador não suporta a tag de vídeo.
HunyuanVideo sem FBCache
Seu navegador não suporta a tag de vídeo.
HunyuanVideo com FBCache
O First Block Cache reduziu a velocidade de inferência para 2271,06 segundos em comparação com a linha de base, ou seja, 1,62x mais rápido, mantendo uma perda de qualidade quase nula.
## quantização fp8
O fp8 com quantização dinâmica acelera ainda mais a inferência e reduz o uso de memória. Tanto as ativações quanto os pesos devem ser quantizados para usar os [NVIDIA Tensor Cores](https://www.nvidia.com/en-us/data-center/tensor-cores/) de 8 bits.
Use `float8_weight_only` e `float8_dynamic_activation_float8_weight` para quantizar o codificador de texto e o modelo transformer.
O método de quantização padrão é a quantização por tensor, mas se sua GPU suportar quantização por linha, você também pode experimentá-la para obter melhor precisão.
Instale o [torchao](https://github.com/pytorch/ao/tree/main) com o comando abaixo.
```bash
pip3 install -U torch torchao
```
[torch.compile](https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html) com `mode="max-autotune-no-cudagraphs"` ou `mode="max-autotune"` seleciona o melhor kernel para desempenho. A compilação pode levar muito tempo se for a primeira vez que o modelo é chamado, mas vale a pena uma vez que o modelo tenha sido compilado.
Este exemplo apenas quantiza o modelo transformer, mas você também pode quantizar o codificador de texto para reduzir ainda mais o uso de memória.
> [!TIP]
> A quantização dinâmica pode alterar significativamente a distribuição da saída do modelo, portanto, você precisa alterar o `residual_diff_threshold` para um valor maior para que ela tenha efeito.
```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")
```
A quantização dinâmica fp8 e o torch.compile reduziram o tempo de inferência para 7,56 segundos em comparação com a linha de base, ou seja, 3,48x mais rápido.
```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)
```
A GPU NVIDIA L20 possui apenas 48GB de memória e pode enfrentar erros de falta de memória (OOM) após a compilação, caso `enable_model_cpu_offload` não seja chamado, pois o HunyuanVideo possui tensores de ativação muito grandes ao ser executado com alta resolução e um grande número de quadros. Para GPUs com menos de 80GB de memória, você pode tentar reduzir a resolução e o número de quadros para evitar erros de OOM.
Modelos grandes de geração de vídeo geralmente são limitados pelos cálculos de atenção, em vez das camadas totalmente conectadas. Esses modelos não se beneficiam significativamente da quantização e do torch.compile.
## Paralelismo de Contexto
O Context Parallelism paraleliza a inferência e escala com múltiplas GPUs. O design composicional do ParaAttention permite combinar o Context Parallelism com o First Block Cache e a quantização dinâmica.
> [!TIP]
> Consulte o repositório [ParaAttention](https://github.com/chengzeyi/ParaAttention/tree/main) para obter instruções detalhadas e exemplos de como escalar a inferência com múltiplas GPUs.
Se o processo de inferência precisar ser persistente e acionável, é sugerido usar [torch.multiprocessing](https://pytorch.org/docs/stable/multiprocessing.html) para escrever seu próprio processador de inferência. Isso pode eliminar o overhead de iniciar o processo e de carregar e recompilar o modelo.
O exemplo de código abaixo combina First Block Cache, quantização dinâmica fp8, torch.compile e Context Parallelism para a velocidade de inferência mais rápida.
```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()
```
Salve em `run_flux.py` e execute com [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
```
A velocidade de inferência é reduzida para 8,20 segundos em comparação com a linha de base, ou seja, 3,21x mais rápida, com 2 GPUs NVIDIA L20. Em 4 L20s, a velocidade de inferência é de 3,90 segundos, ou seja, 6,75x mais rápida.
O exemplo de código abaixo combina First Block Cache e Context Parallelism para a velocidade de inferência mais rápida.
```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()
```
Salve em `run_hunyuan_video.py` e execute com [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
```
A velocidade de inferência é reduzida para 649,23 segundos em comparação com a linha de base, ou seja, 5,66x mais rápida, com 8 GPUs NVIDIA L20.
## Benchmarks
| Tipo de GPU | Número de GPUs | Otimizações | Tempo Real (s) | Aceleração |
| - | - | - | - | - |
| NVIDIA L20 | 1 | Linha de base | 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 |
| Tipo de GPU | Número de GPUs | Otimizações | Tempo Real (s) | Aceleração |
| - | - | - | - | - |
| NVIDIA L20 | 1 | Linha de base | 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 |

Xet Storage Details

Size:
17.6 kB
·
Xet hash:
3831c5eb1591eedb14c5419b5a38bfb779973abd2de4145ca2e048dfa299c964

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.