Paralelismo de contexto

El paralelismo de contexto divide la secuencia de entrada entre las GPU y cada GPU procesa su propia porción. Esto reduce la memoria por GPU, lo cual resulta más útil para secuencias largas, como imágenes de alta resolución y vídeo. Para compararlo con las demás estrategias, consulte Inferencia distribuida.

El paralelismo de contexto se ejecuta dentro del backend de atención, y solo algunos backends lo admiten. Si el suyo no lo hace, enable_parallelism lanza un error que enumera los backends compatibles. Cambie de backend con set_attention_backend() y consulte los backends disponibles para ver la lista completa.

Configure cada variante de abajo con ContextParallelConfig.

Ring Attention

Ring Attention transmite los bloques de claves (K) y valores (V) entre las GPU dispuestas en un anillo. Cada GPU calcula la atención contra el bloque K/V que posee y luego pasa ese bloque a la GPU siguiente. Después de un ciclo completo, cada GPU ha visto los K/V de cada token y ninguna GPU almacena la secuencia completa.

Establezca ring_degree en el número de GPUs y pase la configuración al método enable_parallelism() del 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()

Guarde el script como ring_attention.py y ejecútelo con torchrun o cualquier otro lanzador compatible con PyTorch. Establezca --nproc-per-node en el número de GPUs.

torchrun --nproc-per-node 2 ring_attention.py

Para probar las otras variantes, reemplace la llamada a enable_parallelism en ring_attention.py con la de su sección.

Ulysses Attention

Ulysses Attention utiliza comunicación all-to-all, donde cada GPU intercambia datos con todas las demás GPUs. El primer all-to-all reagrupa los datos para que cada GPU contenga todos los tokens de un subconjunto de cabezas de atención. Cada GPU calcula la atención para sus cabezas, y un segundo all-to-all reagrupa los resultados por token para la siguiente capa.

Establezca ulysses_degree en el número de GPUs y pase la configuración a enable_parallelism().

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

Atención unificada

Unified Sequence Parallelism combina Ring Attention y Ulysses Attention. Ejecuta un all-to-all para redistribuir las cabezas y los tokens, ejecuta Ring Attention sobre el resultado y luego invierte el all-to-all para restaurar el diseño original. Ulysses paraleliza entre las cabezas de atención y Ring paraleliza a lo largo de la secuencia, por lo que, juntos, pueden utilizar más GPUs que Ulysses por sí solo.

Use Unified Attention cuando tenga al menos 4 GPUs. Establezca tanto ulysses_degree como ring_degree en un valor mayor que 1. Las GPUs forman una cuadrícula 2D de ulysses_degree * ring_degree dispositivos, y cada GPU pertenece a un grupo Ulysses y a un grupo Ring.

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

El siguiente benchmark compara Ulysses, Ring y Unified Attention en 4 GPUs H100 con este script. Cada imagen utiliza 50 pasos de inferencia.

Ulysses tiene el mayor rendimiento, pero no puede usar más GPUs que las cabezas de atención que tiene el modelo. Unified Attention elimina ese límite.

Ulysses Anything Attention

Ulysses Attention requiere que la longitud de la secuencia y el número de cabezas de atención sean divisibles exactamente por ulysses_degree. Ulysses Anything Attention elimina ambos requisitos. No se puede combinar con Unified Attention, por lo que ring_degree debe ser 1.

Establezca ulysses_degree en un valor mayor que 1 y ulysses_anything=True.

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

Agregue un backend gloo para tensores de CPU a init_process_group. Esto evita las sincronizaciones forzadas de CUDA que provocan las transferencias de host a dispositivo y de dispositivo a host.

import torch.distributed as dist

dist.init_process_group(backend="cpu:gloo,cuda:nccl")

El benchmark siguiente ejecuta FLUX.1-dev en 4 GPUs L20 con este script.

En 1008x1008, donde la longitud de la secuencia no se divide exactamente, solo se ejecuta Ulysses Anything Attention. En 1024x1024, su rendimiento es el mismo que el de Ulysses Attention estándar.

Ring Anything Attention

Ring Attention requiere que la longitud de la secuencia sea divisible exactamente por ring_degree. Ring Anything Attention admite cualquier longitud de secuencia. Rellena los K/V locales de cada GPU hasta la longitud local más larga, realiza un all-gather de los K/V rellenos y recorta cada bloque a su longitud real antes de calcular la atención.

No se puede combinar con Unified Attention, por lo que ulysses_degree debe ser 1. Establezca ring_degree en un valor mayor que 1 y ring_anything=True.

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=2, ring_anything=True))

Ring Anything Attention también se beneficia del backend gloo mostrado en Ulysses Anything Attention.

Ring Anything Attention solo admite inferencia y attn_mask debe ser None.

El benchmark siguiente ejecuta FLUX.1-dev en 4 GPUs RTX 4090 (48GB).

Ring Anything Attention se ejecuta a 1008x1008 y rinde aproximadamente igual que el Ring Attention estándar. Para conocer la motivación y las compensaciones, consulte este comentario.

Activar el paralelismo de contexto durante la carga

En lugar de llamar a enable_parallelism después de cargar, pase el ContextParallelConfig al argumento parallel_config de from_pretrained(). Inicialice primero el grupo de procesos, ya que from_pretrained genera un error si torch.distributed no está inicializado. Ejecute el script con torchrun, como en el ejemplo 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)