Ulisses Desatado: Paralelismo de Contexto Eficiente em Memória via Divisão por Cabeças

Untied Ulysses: Memory-Efficient Context Parallelism via Headwise Chunking

February 24, 2026

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

cs.AI

Resumo

O processamento eficiente de sequências longas com modelos Transformer geralmente requer a divisão dos cálculos entre aceleradores por meio de paralelismo de contexto. As abordagens dominantes nesta família de métodos, como Ring Attention ou DeepSpeed Ulysses, permitem a escalabilidade na dimensão de contexto, mas não se concentram na eficiência de memória, o que limita os comprimentos de sequência que podem suportar. Técnicas mais avançadas, como o Transformer Distribuído Totalmente Pipeline ou o descarregamento de ativações, podem estender ainda mais o comprimento de contexto possível ao custo da taxa de transferência de treinamento. Neste artigo, apresentamos o UPipe, uma técnica de paralelismo de contexto simples, mas eficaz, que realiza um fracionamento de granularidade fina ao nível do cabeçalho de atenção. Esta técnica reduz significativamente o uso de memória de ativação da auto-atenção, quebrando a barreira de memória de ativação e permitindo comprimentos de contexto muito maiores. Nossa abordagem reduz o uso de memória dos tensores intermediários na camada de atenção em até 87,5% para Transformers de 32B, enquanto iguala as técnicas anteriores de paralelismo de contexto em termos de velocidade de treinamento. O UPipe pode suportar um comprimento de contexto de 5M de tokens ao treinar o Llama3-8B em um único nó 8×H100, superando os métodos anteriores em mais 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%.