# Road to 5 Million Tokens: Breaking Barriers in Long Context Training Max Ryabinin, Together AI | AI Engineer Europe 2026 | 15:50 Source: https://www.youtube.com/watch?v=TUnPNY4E2fw Channel: AI Engineer (https://www.youtube.com/@aiDotEngineer). Summarised by AIE Talks. Page: https://aietalks.com/talks/road-to-5-million-tokens-breaking-barriers-in-long-context-training Published: 2026-06-08 Tags: gpus, inference, long-context ## TL;DR - Long-context training runs into quadratic attention computation and activation memory that grows with sequence length. - A 3 million-token Llama 3B run on one 8xH100 node requires sharding, Ulysses context parallelism, activation checkpointing, CPU offloading, and chunked sequence training. - Untied Ulysses reduces memory further by processing smaller groups of attention heads and reusing their buffers across iterations, allowing training at 5 million tokens. ## Summary Max Ryabinin explains how Together AI fit extremely long-context training onto a single 8xH100 node. A standard Llama 3B model cannot fit even its parameters at the start of a 3 million-token run. Fully sharded data parallelism distributes the parameters, while DeepSpeed Ulysses context parallelism reduces attention activations by about 8x. Activation checkpointing provides another roughly 8x reduction. CPU offloading moves transformer block inputs out of GPU memory, and chunked sequence training prevents element-wise operations from creating buffers millions of tokens wide. Together, these methods make 3 million tokens possible. To reach 5 million, Untied Ulysses splits each GPU's attention heads into smaller groups, computes them over separate iterations, and reuses the allocated buffers. Ryabinin reports results close to memory-optimized transformer training methods at 8B and 32B scale, with further context-length scaling and little throughput loss at smaller chunk sizes. ## Key ideas ### Long contexts create both compute and memory bottlenecks [03:31](https://www.youtube.com/watch?v=TUnPNY4E2fw&t=211s) Ryabinin identifies two problems when extending a transformer's context. Attention has quadratic computation because the model forms pairwise interactions across sequence elements. Memory grows linearly with sequence length, which becomes difficult at millions of tokens. He gives agents as one reason to use large contexts, since they may need to retain more information, and video generation as another, since multiple frames can quickly consume tokens. Training must teach the model to process and use that context correctly. Understanding memory use can also free resources for other training improvements at shorter sequence lengths. ### Fully sharded data parallelism solves parameter placement but leaves attention activations [05:14](https://www.youtube.com/watch?v=TUnPNY4E2fw&t=314s) In Ryabinin's example, a standard Llama 3B model is trained with 3 million tokens on an 8xH100 node. The model parameters alone exhaust GPU memory before computation starts. Fully sharded data parallelism chunks the parameters across the eight GPUs and lowers model-memory use substantially. It does not solve the run, because attention activations still consume enough memory to cause an out-of-memory failure. This makes the example useful as a stack of memory interventions rather than a single technique that fixes long-context training. ### DeepSpeed Ulysses distributes attention heads across GPUs [06:04](https://www.youtube.com/watch?v=TUnPNY4E2fw&t=364s) DeepSpeed Ulysses changes how multi-head attention is computed across the GPUs. Each GPU works on different attention heads while communicating the activations needed to compute attention over the full sequence. Ryabinin describes a GPU as being responsible for one attention head in the example, followed by aggregation of the results. The method reduces activation use by approximately 8x and remains compatible with optimized attention implementations such as Flash Attention. Even after this reduction, the 3 million-token run still needs more memory savings. ### Checkpointing and CPU offloading remove most of the remaining activation pressure [07:48](https://www.youtube.com/watch?v=TUnPNY4E2fw&t=468s) Activation checkpointing avoids keeping every activation for the backward pass. Instead, the needed activations are recomputed during backpropagation, reducing activation use by another factor of about eight when configured correctly. Ryabinin then describes offloading transformer-block inputs to the CPU when they are not needed on the GPU. They can be prefetched before the corresponding backward computation, so the technique has limited performance impact in his description. He attributes the first implementation of this offloading approach, to their knowledge, to Unsloth. ### Chunked sequence training avoids buffers millions of tokens wide [09:33](https://www.youtube.com/watch?v=TUnPNY4E2fw&t=573s) Some operations do not need to process the entire sequence at once. Ryabinin says loss computation and MLP operations can be tiled across sequence length, avoiding buffers with a dimension of 3 million tokens. He calls this chunked sequence length training. This final part of the initial stack makes the 3 million-token Llama 3B example possible on the single 8xH100 node. The technique applies to element-wise computation, while attention requires the separate context-parallelism methods discussed earlier. ### Untied Ulysses reuses smaller attention buffers across iterations [09:57](https://www.youtube.com/watch?v=TUnPNY4E2fw&t=597s) Untied Ulysses examines the context-parallel attention step more closely. Ryabinin says one set of heads can already saturate a GPU's compute capacity within one iteration. If a GPU has several head groups scheduled, those groups can be split into chunks and processed over time. The system recomputes one group, computes its attention, stores the partial result, and then reuses the allocated buffers for the next group. This replaces one large buffer with a smaller buffer reused across two or more iterations, reducing activation memory without a significant throughput impact at small scales. ### Chunk size trades memory use against throughput [11:43](https://www.youtube.com/watch?v=TUnPNY4E2fw&t=703s) The results compare several context-parallel training methods at 8B and 32B model scale. Ryabinin says Untied Ulysses comes close to the most memory-optimized transformer training implementations while reaching 5 million tokens, and can be more performant at shorter contexts in some cases. The chunk-size tradeoff is direct: larger chunks use more memory but let the model run faster. Smaller chunks reduce memory and require more iterations. The broader stack can also leave memory available for other parts of training, such as pipeline stages. ## Notable quotes - "The first stage you'll see is that even with just model parameters, you're not able to fit it into the GPU." (05:14) - "With activation checkpointing, you can drop the activation usage by like a further factor of eight, but still something else needs to be done." (08:08) - "The advantage here is that instead of allocating this huge buffer as you would have before, you allocate a buffer which is smaller, but you reuse it across two or more different iterations." (11:11) - "The bottlenecks might appear where you least expect." (12:56) ## Tools & references mentioned - Together AI - Llama 3B - DeepSpeed Ulysses - Microsoft - Flash Attention - Unsloth - PyTorch profiler ## Who should watch - You are training transformer models with context windows large enough for attention activations to dominate GPU memory. - You are deciding whether sharding, checkpointing, offloading, or context parallelism will address a long-context out-of-memory failure. - You want to understand the memory and throughput tradeoff in Untied Ulysses before working with multi-million-token sequences.