Untied Ulysses: Memory-Efficient Context Parallelism via Headwise Chunking

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 lowers the peak activation memory usage by as much as 82.5% for 70B Transformers, while matching previous context parallelism techniques in terms of training speed. UPipe can support maximum context lengths of up to 5M tokens for training 8B models on a single 8xH100 node, improving upon prior methods by 25%.

Training AI models such as LLMs on long inputs (such as hundred-page documents, videos, or large codebases) requires enormous GPU memory to store intermediate computations. We wanted to understand the systems limitations of model training on long contexts, finding and eliminating the bottlenecks that prevent us from scaling to longer sequences. Our work shows that one non-trivial bottleneck not addressed by prior work is caused by the attention mechanism, which is central to today's AI models. There are multiple existing methods to split the attention computation across multiple GPUs, but they do not prioritize memory efficiency. We introduce UPipe, a memory-efficient context method that processes attention in a distributed way, handling only a subset of inputs at a time. Because the inputs are long enough, processing parts of it sequentially does not incur any major runtime overhead, while significantly reducing memory pressure and unlocking longer contexts. UPipe enables training models Llama3-8B on a single H100 GPU node with up to 5M tokens, which improves upon prior work by 25%. Our method will allow everyone training AI models for long-sequence data to further increase the possible input sizes within the same hardware constraints.