Context Parallelism

Le parallélisme de contexte est une technique de calcul distribué qui répartit une longue séquence d’entrée sur plusieurs accélérateurs. Chaque GPU traite une partie de la séquence tout en collaborant avec les autres pendant l’attention. Cela réduit la mémoire d’activation par appareil et permet d’entraîner un Transformer sur des entrées susceptibles de dépasser la mémoire d’un seul GPU, comme des documents très longs, des vidéos étendues ou de grandes collections de fragments d’image.

Contrairement au simple agrandissement de la fenêtre de contexte d’un modèle, le parallélisme de contexte ne modifie pas la quantité d’informations que l’architecture peut théoriquement accepter. Il rend plutôt le traitement de ce contexte viable sur le plan du calcul en répartissant la dimension de 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, comme la normalisation et les couches feed-forward, peuvent traiter indépendamment ces fragments locaux de séquence. Le défi concerne le mécanisme d’attention : une requête locale peut devoir porter sur les clés et les valeurs détenues par tous les autres appareils.

Les implémentations échangent donc des blocs clé-valeur, ou KV, entre les GPU. Dans l’attention en anneau, 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’au traitement de la séquence entière. Le tutoriel PyTorch sur le parallélisme de contexte présente ce fonctionnement avec l’attention distribuée à produit scalaire à l’échelle, tandis que le package NVIDIA de parallélisme de contexte décrit les options de communication par all-gather, reduce-scatter et anneau.

Le résultat final est mathématiquement équivalent à l’attention sur la séquence entière, à de normales différences numériques près, mais aucun GPU ne doit conserver toutes les activations de la séquence.

Pourquoi le parallélisme de contexte est important#

Les longues séquences créent deux problèmes majeurs de mise à l’échelle. D’abord, les activations sauvegardées occupent davantage de mémoire à mesure que la séquence s’allonge. Ensuite, l’auto-attention standard compare les jetons de la séquence, ce qui génère d’importants calculs et beaucoup de données temporaires.

Le parallélisme de contexte résout le problème de mémoire en répartissant les activations entre les appareils. Il peut aussi répartir le travail lié à l’attention, mais la communication entraîne un nouveau coût. Les systèmes efficaces font se chevaucher les transferts KV et les calculs à l’aide d’opérations telles que celles décrites dans le guide des opérations collectives NCCL.

Cette technique est surtout utile lorsque la longueur des séquences, plutôt que les poids du modèle ou la taille des lots, provoque une erreur de mémoire insuffisante. Elle relève du domaine plus vaste de l’entraînement distribué et se combine couramment à d’autres stratégies pour mettre simultanément à l’échelle plusieurs dimensions.

Parallélisme de contexte et techniques apparentées#

• Le parallélisme tensoriel répartit les opérations ou les matrices de poids au sein de couches individuelles. Le parallélisme de contexte répartit plutôt les jetons le long de la dimension de séquence.
• Le parallélisme en pipeline attribue différents groupes de couches du modèle à différents appareils. Il répartit la profondeur du modèle plutôt que la longueur des séquences.
• Le parallélisme des données réplique le modèle et fournit des exemples d’entraînement différents à chaque réplique. La présentation de PyTorch sur le calcul distribué le recommande lorsque le modèle complet et chaque échantillon tiennent sur un seul GPU.
• Le parallélisme des séquences répartit souvent les activations de certaines opérations associées au parallélisme tensoriel. Le parallélisme de contexte applique plus largement le partitionnement des séquences aux entrées et aux activations du réseau.

Ces approches sont complémentaires. Le guide NVIDIA sur les stratégies de parallélisme montre comment le parallélisme de contexte, tensoriel, en pipeline et des données peut former une disposition multidimensionnelle des appareils.

Applications concrètes#

• IA appliquée aux documents longs : Un modèle de langage juridique ou médical peut devoir traiter un dossier entier, un historique de patient ou un manuel technique. Le parallélisme de contexte répartit les milliers de jetons du document entre les accélérateurs, réduisant la pression sur la mémoire des activations tout en préservant l’attention entre des sections éloignées.

• Compréhension des longues vidéos et des données multimodales : Les Transformers vidéo et les grands modèles de vision peuvent représenter des images vidéo, des fragments d’image, des segments audio et du texte sous la forme d’une longue séquence de jetons. Répartir cette séquence aide les modèles à analyser de longs enregistrements sans réduire fortement le nombre d’images ou le niveau de détail spatial. La présentation AWS Neuron du parallélisme de contexte illustre comment des groupes d’accélérateurs peuvent échanger des fragments KV pour ces tâches à long contexte.

Pour les architectures compactes de vision par ordinateur, comme Ultralytics YOLO26, le parallélisme de contexte est généralement inutile. Le parallélisme standard des données sur plusieurs GPU via le flux de travail d’entraînement des modèles Ultralytics est généralement plus adapté 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 commandes de contrôle de l’attention à produit scalaire à l’échelle. 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 séquence : chaque processus reçoit donc un fragment de séquence, tandis que l’attention est coordonnée dans le maillage d’appareils.

En pratique, les ingénieurs doivent vérifier que la réduction de mémoire compense les coûts de communication, utiliser des interconnexions rapides et tester 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 annoter des jeux de données, entraîner et déployer des modèles, et assurer leur suivi, sans configuration manuelle du parallélisme de contexte.