Context Parallelism

Le parallélisme de contexte est une technique de calcul distribué qui divise une longue séquence d'entrée entre plusieurs accélérateurs. Chaque GPU ne traite qu'une partie de la séquence tout en coopérant avec les autres lors de l'attention. Cela réduit la mémoire d'activation par appareil, permettant à un transformer de s'entraîner sur des entrées qui peuvent dépasser la mémoire d'un seul GPU, telles que des documents très longs, des vidéos prolongées ou de grandes collections de patchs d'image.

Contrairement au simple élargissement de la context window d'un modèle, le parallélisme de contexte ne modifie pas la quantité d'informations que l'architecture peut théoriquement accepter. En revanche, il rend le traitement de ce contexte pratiquement réalisable sur le plan informatique en distribuant la dimension de la séquence.

Fonctionnement du parallélisme de contexte#

Supposons qu'une séquence contienne 32 000 jetons et que le parallélisme de contexte utilise quatre GPU. Chaque appareil reçoit initialement environ 8 000 jetons et stocke les activations intermédiaires correspondantes.

La plupart des opérations, telles que la normalisation et les couches feed-forward, peuvent traiter ces blocs de séquence locaux de manière indépendante. Le défi réside dans le mechanism d'attention : une requête locale peut avoir besoin de prêter attention aux clés et aux valeurs détenues par tous les autres appareils.

Les implémentations échangent donc des blocs de clés-valeurs, ou KV, entre les GPU. Dans le ring attention, chaque appareil calcule une attention partielle à l'aide de ses données locales, transmet un bloc KV à l'appareil suivant et répète l'opération jusqu'à ce qu'il ait traité la séquence complète. Le PyTorch context parallel tutorial démontre ce comportement via une attention à produit scalaire mis à l'échelle et distribué, tandis que le NVIDIA context parallel package décrit les options de communication all-gather, reduce-scatter et basées sur l'anneau.

Le résultat final est mathématiquement équivalent à une attention sur séquence complète, sous réserve des différences numériques normales, mais aucun GPU unique ne doit conserver l'intégralité des activations de la séquence.

Pourquoi le parallélisme de contexte est important#

Les séquences longues posent deux problèmes majeurs de mise à l'échelle. Premièrement, les activations sauvegardées consomment plus de mémoire à mesure que la longueur de la séquence augmente. Deuxièmement, l'auto-attention standard compare les jetons de la séquence, générant un calcul substantiel et des données temporaires.

Le parallélisme de contexte résout le problème de la mémoire en divisant les activations entre les appareils. Il peut également diviser le travail d'attention, bien que la communication introduise un nouveau coût. Les systèmes efficaces superposent les transferts KV et les calculs à l'aide d'opérations telles que celles documentées dans le NCCL collective operations guide.

Cette technique est particulièrement utile lorsque la longueur de la séquence, plutôt que les poids du modèle ou la taille du lot, provoque une erreur de dépassement de mémoire (out-of-memory). Elle fait partie du domaine plus large du distributed training et est couramment combinée avec d'autres stratégies pour mettre à l'échelle plusieurs dimensions simultanément.

Parallélisme de contexte vs techniques apparentées#

• Tensor parallelism divise les opérations ou les matrices de poids au sein des couches individuelles. Le parallélisme de contexte divise plutôt les jetons le long de la dimension de la séquence.
• Pipeline parallelism attribue différents groupes de couches de modèle à différents appareils. Il partitionne la profondeur du modèle plutôt que la longueur de la séquence.
• Le Data parallelism réplique le modèle et attribue à chaque réplique différents exemples d'entraînement. Le PyTorch distributed overview le recommande lorsque le modèle complet et chaque échantillon tiennent sur un seul GPU.
• Le Sequence parallelism fragmente souvent les activations pour des opérations sélectionnées associées au parallélisme de tenseurs. Le parallélisme de contexte applique le partitionnement de séquence de manière plus large à travers les entrées et les activations du réseau.

Ces approches sont complémentaires. Le NVIDIA parallelism strategies guide montre comment le parallélisme de contexte, de tenseurs, de pipeline et de données peuvent former une disposition d'appareils multidimensionnelle.

Applications concrètes#

• IA pour documents longs : Un modèle de langage juridique ou médical peut avoir besoin de traiter un dossier complet, l'historique d'un patient ou un manuel technique. Le parallélisme de contexte distribue les milliers de jetons du document entre les accélérateurs, réduisant la pression sur la mémoire d'activation tout en préservant l'attention entre les sections distantes.

• Compréhension vidéo longue et multimodale : Les transformers vidéo et les large vision models peuvent représenter des images, des patchs d'image, des segments audio et du texte sous forme d'une seule longue séquence de jetons. La distribution de cette séquence aide les modèles à analyser des enregistrements prolongés sans réduire agressivement le nombre d'images ou le détail spatial. Le AWS Neuron context parallelism overview illustre comment des groupes d'accélérateurs peuvent échanger des fragments KV pour ces charges de travail à contexte long.

Pour les architectures de vision par ordinateur compactes telles qu'Ultralytics YOLO26, le parallélisme de contexte n'est généralement pas nécessaire. Le parallélisme de données multi-GPU standard via le Ultralytics model training workflow est généralement la méthode la plus appropriée pour accélérer l'entraînement.

Utilisation pratique et compromis#

L'exemple minimal suivant utilise l'API expérimentale de parallélisme de contexte de PyTorch et les scaled dot-product attention controls. Enregistre-le sous cp_example.py et lance-le sur deux GPU avec torchrun --standalone --nproc-per-node=2 cp_example.py.

import os

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.experimental import context_parallel
from torch.nn.attention import SDPBackend, sdpa_kernel

rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])

torch.cuda.set_device(rank)
torch.cuda.manual_seed(0)
dist.init_process_group("nccl")

mesh = init_device_mesh("cuda", (world_size,))
qkv = [torch.randn(1, 4, 4096, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True) for _ in range(3)]

with sdpa_kernel(SDPBackend.FLASH_ATTENTION), context_parallel(mesh, buffers=tuple(qkv), buffer_seq_dims=(2, 2, 2)):
    output = F.scaled_dot_product_attention(*qkv, is_causal=True)

output.float().square().mean().backward()
dist.destroy_process_group()

Ici, la dimension 2 est la dimension de la séquence, de sorte que chaque processus reçoit un fragment de séquence tandis que l'attention se coordonne à travers le maillage d'appareils.

En pratique, les ingénieurs doivent s'assurer que la réduction de la mémoire l'emporte sur le surcoût de communication, utiliser des interconnexions rapides et évaluer des longueurs de séquence représentatives. Pour les projets de vision standard, Ultralytics Platform propose des flux de travail cloud et locaux plus simples pour l'annotation de jeux de données, l'entraînement, le déploiement et la surveillance sans nécessiter de configuration manuelle du parallélisme de contexte.