Road to 5 Million Tokens: Breaking Barriers in Long Context Training — Max Ryabinin, Together AI

Key Concepts

• Context Parallelism: A technique to distribute the computation of long sequences across multiple GPUs.
• DeepSpeed Ulysses: A method that partitions attention heads across GPUs to handle long-context sequences.
• Activation Checkpointing: A memory-saving technique that recomputes intermediate activations during the backward pass instead of storing them.
• U-Pipe: A novel optimization developed by Together AI that chunks attention head computations to further reduce memory overhead.
• Flash Attention: An I/O-aware exact attention algorithm that speeds up training and reduces memory usage.
• Fully Sharded Data Parallelism (FSDP): A method to shard model parameters across multiple GPUs to fit large models into memory.

---

1. Main Topics and Challenges

Max, VP of R&D at Together AI, discusses the "Road to 5 Million Sequence Length." The primary motivation for long-context training is the rise of AI agents and video generation, which require models to maintain temporal consistency and process massive amounts of data.

The Bottlenecks:

• Quadratic Computation: Standard transformers involve pairwise interactions across all elements in a sequence, leading to $O(n^2)$ complexity.
• Linear Memory Growth: As sequence length increases, memory usage grows linearly, quickly exceeding the capacity of standard GPU hardware (e.g., H100 clusters).

2. Step-by-Step Optimization Framework

To reach a 3-to-5 million token context length on an 8x H100 node, the team applied a layered optimization strategy:

1. FSDP (Fully Sharded Data Parallelism): Distributes model parameters across GPUs to solve the initial "out of memory" error when loading the model.
2. DeepSpeed Ulysses: Instead of computing all multi-head attention on every GPU, this technique assigns specific attention heads to specific GPUs, communicating activations as needed.
3. Activation Checkpointing: Recomputes activations during the backward pass to save memory, reducing activation storage requirements by approximately 8x.
4. CPU Offloading: Moves inputs of transformer blocks to the CPU when not in use, prefetching them only when needed for backpropagation.
5. Sequence Tiling: Chunks element-wise computations (like MLPs and loss functions) to avoid creating massive buffers that scale with the sequence length.
6. U-Pipe (The Core Innovation): An extension of Ulysses that schedules chunks of attention heads over time. By iterating through head groups, the system reuses memory buffers, allowing for significantly larger context windows without sacrificing throughput.

3. Research Findings and Performance

• Memory Efficiency: By stacking these techniques, the team successfully fit 3 million tokens into a single 8x H100 node.
• Scalability: The U-Pipe method allows for scaling up to 5 million tokens.
• Throughput: The relationship between chunk size and throughput is direct; larger chunks increase memory utilization but improve speed. U-Pipe maintains performance parity with standard optimized training even at shorter context lengths.
• Tooling: The speaker emphasizes the use of the PyTorch Profiler to identify hidden bottlenecks, noting that memory issues often appear in unexpected places.

4. Notable Quotes

• "The problem here is that if you are taking a standard transformer-based language model and trying to extend its context, you can run into two bottlenecks... quadratic computation... [and] memory keeps growing linearly."
• "What we found was that even trying to compute one set of heads at a time is already enough to saturate the computational capacity of the GPU... we have one group of heads which are being recomputed... then you follow up with the next stage, which can reuse all the buffers."

5. Synthesis and Conclusion

The "Road to 5 Million" project demonstrates that achieving extreme context lengths is not the result of a single "silver bullet" but rather the orchestration of multiple memory-saving techniques. By combining established methods like DeepSpeed Ulysses and activation checkpointing with their custom U-Pipe optimization, Together AI has effectively broken the memory barriers that previously limited transformer training. The key takeaway is that developers should focus on granular memory management—specifically through head-chunking and buffer reuse—to push the boundaries of what current GPU hardware can handle.