Buckets:

hf-doc-build/doc / diffusers /v0.43.0 /pt /optimization /context_parallelism.md
mishig's picture
|
download
raw
12.7 kB
# Paralelismo de contexto
[Paralelismo de contexto](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=context_parallelism) divide a sequência de entrada entre as GPUs, e cada GPU processa sua própria fatia. Isso reduz a memória por GPU, o que ajuda mais com sequências longas, como imagens de alta resolução e vídeo. Para compará-lo com as outras estratégias, veja [Inferência distribuída](../training/distributed_inference).
O paralelismo de contexto é executado dentro do backend de atenção, e apenas alguns backends o suportam. Se o seu não for compatível, `enable_parallelism` lança um erro que lista os backends compatíveis. Alterne os backends com [set_attention_backend()](/docs/diffusers/v0.43.0/pt/api/models/overview#diffusers.ModelMixin.set_attention_backend) e consulte os [backends disponíveis](./attention_backends#available-backends) para a lista completa.
Configure todas as variantes abaixo com [ContextParallelConfig](/docs/diffusers/v0.43.0/pt/api/parallel#diffusers.ContextParallelConfig).
| Variante | Config | Requisitos | Use quando |
|---|---|---|---|
| [Ring Attention](#ring-attention) | `ring_degree > 1` | o comprimento da sequência deve ser divisível por `ring_degree` | as sequências são longas e a largura de banda de interconexão é limitada |
| [Atenção Ulysses](#ulysses-attention) | `ulysses_degree > 1` | o comprimento da sequência e o número de cabeças de atenção devem ser divisíveis por `ulysses_degree` | suas GPUs possuem um interconector rápido |
| [Atenção Unificada](#unified-attention) | `ring_degree > 1` e `ulysses_degree > 1` | pelo menos 4 GPUs | você tem mais dispositivos do que a Ulysses Attention pode usar sozinha |
| [Ulysses Anything Attention](#ulysses-anything-attention) | `ulysses_degree > 1`, `ulysses_anything=True` | `ring_degree` deve ser 1 | o comprimento da sequência ou o número de cabeças não é divisível |
| [Ring Anything Attention](#ring-anything-attention) | `ring_degree > 1`, `ring_anything=True` | `ulysses_degree` deve ser 1, apenas inferência, `attn_mask` deve ser `None` | o comprimento da sequência não é divisível |
## Atenção em Anel
[Ring Attention](https://huggingface.co/papers/2310.01889) passa os blocos de chave (K) e valor (V) entre GPUs organizadas em um anel. Cada GPU calcula a atenção em relação ao bloco K/V que possui e, em seguida, passa esse bloco para a próxima GPU. Após uma volta completa, todas as GPUs terão visto o K/V de todos os tokens, e nenhuma GPU armazena a sequência completa.
Defina `ring_degree` para o número de GPUs e passe a configuração para o método `enable_parallelism()` do transformer.
```py
import os
import torch
from torch import distributed as dist
from diffusers import DiffusionPipeline, ContextParallelConfig
def setup_distributed():
if not dist.is_initialized():
dist.init_process_group(backend="nccl")
# torchrun sets LOCAL_RANK to this process's GPU index on its node
local_rank = int(os.environ["LOCAL_RANK"])
device = torch.device(f"cuda:{local_rank}")
torch.cuda.set_device(device)
return device
def main():
device = setup_distributed()
world_size = dist.get_world_size()
pipeline = DiffusionPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev", dtype=torch.bfloat16
).to(device)
pipeline.transformer.set_attention_backend("_native_cudnn")
cp_config = ContextParallelConfig(ring_degree=world_size)
pipeline.transformer.enable_parallelism(config=cp_config)
prompt = """
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
"""
# Must specify generator so all ranks start with same latents (or pass your own)
generator = torch.Generator().manual_seed(42)
image = pipeline(
prompt,
guidance_scale=3.5,
num_inference_steps=50,
generator=generator,
).images[0]
if dist.get_rank() == 0:
image.save("output.png")
if dist.is_initialized():
dist.destroy_process_group()
if __name__ == "__main__":
main()
```
Salve o script como `ring_attention.py` e execute-o com [torchrun](https://docs.pytorch.org/docs/stable/elastic/run.html) ou outro lançador compatível com PyTorch. Defina `--nproc-per-node` para o número de GPUs.
```bash
torchrun --nproc-per-node 2 ring_attention.py
```
Para tentar as outras variantes, substitua a chamada `enable_parallelism` em `ring_attention.py` pela da seção correspondente.
## Atenção Ulysses
[Ulysses Attention](https://huggingface.co/papers/2309.14509) usa comunicação *all-to-all*, na qual cada GPU troca dados com todas as outras GPUs. O primeiro all-to-all reorganiza os dados para que cada GPU possua todos os tokens de um subconjunto de cabeças de atenção. Cada GPU calcula a atenção para suas cabeças, e um segundo all-to-all reorganiza os resultados por token para a próxima camada.
Defina `ulysses_degree` para o número de GPUs e passe a configuração para `enable_parallelism()`.
```py
pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2))
```
## Atenção Unificada
[Unified Sequence Parallelism](https://huggingface.co/papers/2405.07719) combina Ring Attention e Ulysses Attention. Ele executa um all-to-all para redistribuir cabeças e tokens, executa Ring Attention no resultado e, em seguida, reverte o all-to-all para restaurar o layout original. Ulysses paraleliza entre cabeças de atenção e Ring paraleliza entre a sequência, de modo que, juntos, podem usar mais GPUs do que Ulysses sozinho.
Use Unified Attention quando você tiver pelo menos 4 GPUs. Defina tanto `ulysses_degree` quanto `ring_degree` para valores maiores que 1. As GPUs formam uma grade 2D de `ulysses_degree * ring_degree` dispositivos, e cada GPU pertence a um grupo Ulysses e a um grupo Ring.
```py
pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ring_degree=2))
```
O benchmark abaixo compara Ulysses, Ring e Unified Attention em 4 GPUs H100 com [este script](https://github.com/huggingface/diffusers/pull/12693#issuecomment-3694727532). Cada imagem usa 50 passos de inferência.
| Backend CP | Tempo / Imagem (ms) | Etapas / Seg | Memória de Pico (GB) |
|--------------------|------------------|-------------|------------------|
| ulysses | 6670.789 | 7.50 | 33.85 |
| ring | 13076.492 | 3.82 | 56.02 |
| unified_balanced | 11068.705 | 4.52 | 33.85 |
Ulysses tem a maior taxa de transferência, mas não pode usar mais GPUs do que o número de cabeças de atenção do modelo. Unified Attention remove esse limite.
## Ulysses Anything Attention
A Ulysses Attention exige que o comprimento da sequência e o número de cabeças de atenção sejam divisíveis por `ulysses_degree`. A [Ulysses Anything Attention](https://github.com/huggingface/diffusers/pull/12996) remove ambos os requisitos. Não é possível combiná-la com a Unified Attention, portanto `ring_degree` deve ser 1.
Defina `ulysses_degree` para um valor maior que 1 e `ulysses_anything=True`.
```py
pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ulysses_anything=True))
```
> [!TIP]
> Adicione um backend `gloo` para tensores de CPU ao [init_process_group](https://docs.pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group). Isso evita as sincronizações forçadas de CUDA causadas pelas transferências de host para dispositivo e de dispositivo para host.
```py
import torch.distributed as dist
dist.init_process_group(backend="cpu:gloo,cuda:nccl")
```
O benchmark abaixo executa o FLUX.1-dev em 4 GPUs L20 com [este script](https://github.com/huggingface/diffusers/pull/12996#issuecomment-3797695999).
| Backend CP | Tempo / Iteração (ms) | Etapas / Seg | Memória de Pico (GB) | Formato (HxW)|
|--------------------|------------------|-------------|------------------|------------|
| ulysses | 281.07 | 3.56 | 37.11 | 1024x1024 |
| ring | 351.34 | 2.85 | 37.01 | 1024x1024 |
| unified_balanced | 324.37 | 3.08 | 37.16 | 1024x1024 |
| ulysses_anything | 280.94 | 3.56 | 37.11 | 1024x1024 |
| ulysses | falhou | falhou | falhou | 1008x1008 |
| ring | falhou | falhou | falhou | 1008x1008 |
| unified_balanced | falhou | falhou | falhou | 1008x1008 |
| ulysses_anything | 278.40 | 3.59 | 36.99 | 1008x1008 |
Em 1008x1008, onde o comprimento da sequência não é divisível uniformemente, apenas o Ulysses Anything Attention é executado. Em 1024x1024, ele apresenta o mesmo desempenho que o Ulysses Attention padrão.
## Ring Anything Attention
Ring Attention exige que o comprimento da sequência seja divisível por `ring_degree`. [Ring Anything Attention](https://github.com/huggingface/diffusers/pull/13545#issuecomment-4302195582) suporta qualquer comprimento de sequência. Ele preenche o K/V local de cada GPU até o comprimento local mais longo, faz o all-gather do K/V preenchido e fatia cada bloco de volta para seu comprimento verdadeiro antes de calcular a atenção.
Você não pode combiná-lo com Unified Attention, portanto `ulysses_degree` deve ser 1. Defina `ring_degree` para um valor maior que 1 e `ring_anything=True`.
```py
pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=2, ring_anything=True))
```
Ring Anything Attention também se beneficia do backend `gloo` mostrado em [Ulysses Anything Attention](#ulysses-anything-attention).
> [!NOTE]
> Ring Anything Attention suporta apenas inferência, e `attn_mask` deve ser `None`.
O benchmark abaixo executa o FLUX.1-dev em 4 GPUs RTX 4090 (48GB).
| Backend CP | Tempo / Iteração (ms) | Etapas / Seg | Memória de Pico (GB) | Formato (HxW)|
|--------------------|------------------|-------------|------------------|------------|
| ulysses | 259.07 | 3.86 | 33.83 | 1024x1024 |
| ring | 338.98 | 2.95 | 33.83 | 1024x1024 |
| unified_balanced | 321.54 | 3.11 | 33.83 | 1024x1024 |
| ulysses_anything | 259.07 | 3.86 | 33.83 | 1024x1024 |
| ring_anything | 340.14 | 2.94 | 33.83 | 1024x1024 |
| ulysses | falhou | falhou | falhou | 1008x1008 |
| ring | falhou | falhou | falhou | 1008x1008 |
| unified_balanced | falhou | falhou | falhou | 1008x1008 |
| ulysses_anything | 253.16 | 3.95 | 33.75 | 1008x1008 |
| ring_anything | 335.57 | 2.98 | 33.75 | 1008x1008 |
Ring Anything Attention é executado em 1008x1008 e apresenta desempenho aproximadamente igual ao do Ring Attention padrão. Para a motivação e os trade-offs, veja [este comentário](https://github.com/huggingface/diffusers/pull/13545#issuecomment-4304104462).
## Ative o paralelismo de contexto durante o carregamento
Em vez de chamar `enable_parallelism` após o carregamento, passe o [ContextParallelConfig](/docs/diffusers/v0.43.0/pt/api/parallel#diffusers.ContextParallelConfig) para o argumento `parallel_config` de [from_pretrained()](/docs/diffusers/v0.43.0/pt/api/models/overview#diffusers.ModelMixin.from_pretrained). Inicialize o grupo de processos primeiro, pois `from_pretrained` levanta um erro se `torch.distributed` não estiver inicializado. Inicie o script com `torchrun`, como no exemplo de [Ring Attention](#ring-attention).
```py
import os
import torch
from torch import distributed as dist
from diffusers import AutoModel, ContextParallelConfig, DiffusionPipeline
CKPT_ID = "black-forest-labs/FLUX.1-dev"
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
device = torch.device(f"cuda:{local_rank}")
torch.cuda.set_device(device)
cp_config = ContextParallelConfig(ring_degree=dist.get_world_size())
transformer = AutoModel.from_pretrained(
CKPT_ID,
subfolder="transformer",
dtype=torch.bfloat16,
parallel_config=cp_config
)
pipeline = DiffusionPipeline.from_pretrained(
CKPT_ID, transformer=transformer, dtype=torch.bfloat16,
).to(device)
```

Xet Storage Details

Size:
12.7 kB
·
Xet hash:
4e7402069a4ea76d33194625459491b8b0e55a01d42da1c93baf3430ddc858cb

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