Training a standard LLaMA 3B model with a 3 million token context on a single 8xH100 node fails before you even start: the model parameters alone exhaust GPU memory. Max Ryabinin from Together AI walks through the full stack of techniques needed to get there: fully sharded data parallelism, DeepSpeed Ulysses context parallelism for an 8x activation reduction, activation checkpointing for another 8x, CPU offloading for transformer block inputs, and chunked sequence training to avoid allocating buffers 3 million tokens wide.

Even that stack falls short at 5 million tokens. The novel contribution, Untied Ulysses, goes deeper into the context parallelism step: instead of allocating one large buffer per attention head group, it chunks the heads further and reuses those buffers across iterations, cutting activation memory with negligible throughput impact. At both 8B and 32B scale the results match the most memory optimized transformer training baselines while pushing sequence length 25% further than prior Ulysses implementations.

Speaker info:
- https://www.linkedin.com/in/max-ryabinin/
- https://x.com/m_ryabinin

AI Engineer

Together AI's VP of R&D presents research on scaling transformer training to 5 million token context lengths. The talk walks through a stack of memory optimization techniques applied to a Llama 3B model on an 8xH100 node: fully sharded data parallelism, DeepSpeed Ulysses context parallelism, activation checkpointing, CPU offloading (pioneered by Unsloth), and chunked sequence-length tiling (Arctic sequence length training). The primary novel contribution is 'Untitled Ulysses' (U-Pipe), which further reduces activation memory by chunking attention head computation within a single GPU iteration, reusing buffers across chunks. Results show competitive or superior throughput vs. other memory-optimized implementations at both 8B and 32B parameter scales while reaching 5M token context lengths.

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