Ungebundener Ulysses: Speichereffiziente Kontext-Parallelität durch kopfweises Chunking
Untied Ulysses: Memory-Efficient Context Parallelism via Headwise Chunking
February 24, 2026
Autoren: Ravi Ghadia, Maksim Abraham, Sergei Vorobyov, Max Ryabinin
cs.AI
Zusammenfassung
Die effiziente Verarbeitung langer Sequenzen mit Transformer-Modellen erfordert in der Regel eine Aufteilung der Berechnungen über Beschleuniger mittels Context Parallelism. Die vorherrschenden Ansätze dieser Methodenfamilie, wie Ring Attention oder DeepSpeed Ulysses, ermöglichen eine Skalierung über die Kontextdimension, konzentrieren sich jedoch nicht auf Speichereffizienz, was die von ihnen unterstützten Sequenzlängen begrenzt. Fortgeschrittenere Techniken, wie Fully Pipelined Distributed Transformer oder das Auslagern von Aktivierungen, können die mögliche Kontextlänge weiter erhöhen, allerdings auf Kosten des Trainingsdurchsatzes. In diesem Artikel stellen wir UPipe vor, eine einfache, aber effektive Context-Parallelism-Technik, die eine feinkörnige Aufteilung auf Ebene der Attention-Heads vornimmt. Diese Technik reduziert den Speicherverbrauch der Aktivierungen im Self-Attention-Mechanismus erheblich, durchbricht die Beschränkung durch den Aktivierungsspeicher und ermöglicht damit deutlich längere Kontextlängen. Unser Ansatz verringert den Speicherverbrauch intermediärer Tensoren in der Attention-Schicht um bis zu 87,5 % für 32B-Transformer, bei gleicher Trainingsgeschwindigkeit wie bisherige Context-Parallelism-Techniken. UPipe kann eine Kontextlänge von 5 Millionen Tokens beim Training von Llama3-8B auf einem einzelnen 8xH100-Knoten unterstützen und übertrifft damit frühere Methoden um über 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%.