Ulysse délié : Parallélisme contextuel économe en mémoire via un découpage par tête

Untied Ulysses: Memory-Efficient Context Parallelism via Headwise Chunking

February 24, 2026

Auteurs: Ravi Ghadia, Maksim Abraham, Sergei Vorobyov, Max Ryabinin

cs.AI

Résumé

Le traitement efficace de longues séquences avec les modèles Transformer nécessite généralement une répartition des calculs sur plusieurs accélérateurs via le parallélisme de contexte. Les approches dominantes de cette famille de méthodes, telles que Ring Attention ou DeepSpeed Ulysses, permettent une mise à l'échelle selon la dimension contextuelle mais ne se concentrent pas sur l'efficacité mémoire, ce qui limite les longueurs de séquence qu'elles peuvent supporter. Des techniques plus avancées, comme le Transformer Distribué Entièrement Pipeline ou le déchargement des activations, peuvent encore étendre la longueur contextuelle possible au détriment du débit d'entraînement. Dans cet article, nous présentons UPipe, une technique de parallélisme de contexte simple mais efficace qui effectue un découpage fin au niveau des têtes d'attention. Cette technique réduit significativement l'utilisation de mémoire des activations dans l'auto-attention, brisant la barrière mémoire des activations et permettant d'atteindre des longueurs contextuelles bien supérieures. Notre approche réduit l'utilisation mémoire des tenseurs intermédiaires dans la couche d'attention jusqu'à 87,5 % pour des Transformers de 32 milliards de paramètres, tout en égalant les techniques de parallélisme de contexte précédentes en termes de vitesse d'entraînement. UPipe peut supporter une longueur contextuelle de 5 millions de tokens lors de l'entraînement de Llama3-8B sur un seul nœud 8×H100, améliorant les méthodes antérieures de plus de 25 %.

English

Efficiently processing long sequences with Transformer models usually requires splitting the computations across accelerators via context parallelism. The dominant approaches in this family of methods, such as Ring Attention or DeepSpeed Ulysses, enable scaling over the context dimension but do not focus on memory efficiency, which limits the sequence lengths they can support. More advanced techniques, such as Fully Pipelined Distributed Transformer or activation offloading, can further extend the possible context length at the cost of training throughput. In this paper, we present UPipe, a simple yet effective context parallelism technique that performs fine-grained chunking at the attention head level. This technique significantly reduces the activation memory usage of self-attention, breaking the activation memory barrier and unlocking much longer context lengths. Our approach reduces intermediate tensor memory usage in the attention layer by as much as 87.5% for 32B Transformers, while matching previous context parallelism techniques in terms of training speed. UPipe can support the context length of 5M tokens when training Llama3-8B on a single 8timesH100 node, improving upon prior methods by over 25%.