Paralelismo de contexto

Paralelismo de contexto 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.

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() e consulte os backends disponíveis para a lista completa.

Configure todas as variantes abaixo com ContextParallelConfig.

Atenção em Anel

Ring Attention 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.

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 ou outro lançador compatível com PyTorch. Defina --nproc-per-node para o número de GPUs.

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 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().

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2))

Atenção Unificada

Unified Sequence Parallelism 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.

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. Cada imagem usa 50 passos de inferência.

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 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.

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ulysses_anything=True))

Adicione um backend gloo para tensores de CPU ao 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.

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.

Em 1008x1008, onde o comprimento da sequência não é divisível uniformemente, apenas a Ulysses Anything Attention é executada. Em 1024x1024, ela apresenta o mesmo desempenho que a 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 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 real 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.

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.

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).

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.

Ative o paralelismo de contexto durante o carregamento

Em vez de chamar enable_parallelism após o carregamento, passe o ContextParallelConfig para o argumento parallel_config de 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.

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)