Ulisse sciolto: parallelismo contestuale efficiente in memoria tramite suddivisione per testa

Untied Ulysses: Memory-Efficient Context Parallelism via Headwise Chunking

February 24, 2026

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

cs.AI

Abstract

L'elaborazione efficiente di sequenze lunghe con modelli Transformer richiede solitamente la suddivisione dei calcoli tra acceleratori tramite parallelismo contestuale. Gli approcci dominanti in questa famiglia di metodi, come Ring Attention o DeepSpeed Ulysses, consentono il ridimensionamento lungo la dimensione contestuale ma non si concentrano sull'efficienza della memoria, limitando le lunghezze di sequenza supportabili. Tecniche più avanzate, come Fully Pipelined Distributed Transformer o lo scaricamento delle attivazioni, possono estendere ulteriormente la lunghezza contestuale possibile a scapito del throughput di addestramento. In questo articolo presentiamo UPipe, una tecnica di parallelismo contestuale semplice ma efficace che esegue una suddivisione in blocchi a grana fine a livello di testa di attenzione. Questa tecnica riduce significativamente l'utilizzo della memoria delle attivazioni nel self-attention, superando la barriera della memoria delle attivazioni e sbloccando lunghezze contestuali molto maggiori. Il nostro approccio riduce l'utilizzo di memoria dei tensori intermedi nello strato di attenzione fino all'87,5% per Transformer da 32B, mantenendo al contempo prestazioni di addestramento equivalenti alle precedenti tecniche di parallelismo contestuale. UPipe può supportare una lunghezza contestuale di 5M token durante l'addestramento di Llama3-8B su un singolo nodo 8×H100, migliorando i metodi precedenti di oltre il 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%.