Parallélisme de contexte

Le parallélisme de contexte répartit la séquence d’entrée sur plusieurs GPU, et chaque GPU traite sa propre portion. Cela réduit la mémoire par GPU, ce qui est particulièrement utile pour les séquences longues, comme les images haute résolution et les vidéos. Pour le comparer aux autres stratégies, consultez Inférence distribuée.

Le parallélisme de contexte s’exécute à l’intérieur du backend d’attention, et seuls certains backends le prennent en charge. Si le vôtre ne le fait pas, enable_parallelism lève une erreur qui énumère les backends compatibles. Changez de backend avec set_attention_backend() et consultez les backends disponibles pour la liste complète.

Configurez chaque variante ci-dessous avec ContextParallelConfig.

Ring Attention

Ring Attention transmet les blocs de clés (K) et de valeurs (V) entre des GPU disposés en anneau. Chaque GPU calcule l'attention par rapport au bloc K/V qu'il détient, puis transmet ce bloc au GPU suivant. Après un tour complet, chaque GPU a vu les K/V de tous les jetons, et aucun GPU ne détient jamais la séquence complète.

Définissez ring_degree sur le nombre de GPU et passez la configuration à la méthode enable_parallelism() du transformeur.

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

Enregistrez le script sous le nom ring_attention.py et lancez-le avec torchrun ou un autre lanceur compatible PyTorch. Définissez --nproc-per-node sur le nombre de GPU.

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

Pour essayer les autres variantes, remplacez l'appel enable_parallelism dans ring_attention.py par celui de leur section.

Ulysses Attention

Ulysses Attention utilise la communication all-to-all, où chaque GPU échange des données avec tous les autres GPU. Le premier all-to-all regroupe les données de sorte que chaque GPU détient tous les jetons pour un sous-ensemble de têtes d'attention. Chaque GPU calcule l'attention pour ses têtes, et un second all-to-all regroupe les résultats par jeton pour la couche suivante.

Définissez ulysses_degree sur le nombre de GPU et passez la configuration à enable_parallelism().

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

Attention unifiée

Unified Sequence Parallelism combine Ring Attention et Ulysses Attention. Il exécute un all-to-all pour redistribuer les têtes et les jetons, applique Ring Attention au résultat, puis inverse l'all-to-all pour restaurer la disposition d'origine. Ulysses parallélise sur les têtes d'attention et Ring sur la séquence, ce qui leur permet, ensemble, d'utiliser plus de GPU qu'Ulysses seul.

Utilisez Unified Attention lorsque vous disposez d’au moins 4 GPU. Définissez à la fois ulysses_degree et ring_degree sur une valeur supérieure à 1. Les GPU forment une grille 2D de ulysses_degree * ring_degree dispositifs, et chaque GPU appartient à un groupe Ulysses et à un groupe Ring.

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

Le benchmark ci-dessous compare Ulysses, Ring et Unified Attention sur 4 GPU H100 à l'aide de ce script. Chaque image utilise 50 étapes d'inférence.

Ulysses offre le débit le plus élevé, mais il ne peut pas utiliser plus de GPU que le nombre de têtes d'attention du modèle. Unified Attention supprime cette limite.

Ulysses Anything Attention

L’attention Ulysses exige que la longueur de la séquence et le nombre de têtes d’attention soient divisibles par ulysses_degree. Ulysses Anything Attention supprime ces deux contraintes. Il ne peut pas être combiné avec Unified Attention, donc ring_degree doit être égal à 1.

Définissez ulysses_degree sur une valeur supérieure à 1 et ulysses_anything=True.

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

Ajoutez un backend gloo pour les tenseurs CPU à init_process_group. Cela évite les synchronisations CUDA forcées causées par les transferts hôte-vers-appareil et appareil-vers-hôte.

import torch.distributed as dist

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

Le benchmark ci-dessous exécute FLUX.1-dev sur 4 GPU L20 avec ce script.

À 1008x1008, où la longueur de la séquence ne se divise pas exactement, seul Ulysses Anything Attention s’exécute. À 1024x1024, il offre des performances identiques à celles de l’Ulysses Attention standard.

Ring Anything Attention

Ring Attention exige que la longueur de la séquence soit divisible par ring_degree. Ring Anything Attention prend en charge toute longueur de séquence. Il remplit les K/V locaux de chaque GPU jusqu'à la longueur locale maximale, effectue un all-gather sur les K/V remplis, puis recoupe chaque bloc à sa longueur réelle avant de calculer l'attention.

Vous ne pouvez pas le combiner avec Unified Attention, donc ulysses_degree doit être 1. Définissez ring_degree sur une valeur supérieure à 1 et ring_anything=True.

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

Ring Anything Attention bénéficie également de l’arrière-plan gloo présenté dans Ulysses Anything Attention.

Ring Anything Attention prend en charge uniquement l'inférence, et attn_mask doit être None.

Le benchmark ci-dessous exécute FLUX.1-dev sur 4 GPU RTX 4090 (48 Go).

Ring Anything Attention s’exécute à 1008x1008 et offre des performances similaires à celles de Ring Attention standard. Pour la motivation et les compromis, consultez ce commentaire.

Activer le parallélisme de contexte lors du chargement

Au lieu d’appeler enable_parallelism après le chargement, passez le ContextParallelConfig à l’argument parallel_config de from_pretrained(). Initialisez d’abord le groupe de processus, car from_pretrained lève une erreur si torch.distributed n’est pas initialisé. Lancez le script avec torchrun, comme dans l’exemple 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)