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

How can a model train on the growing history an agent needs, or the many frames a video model must keep consistent? That is the problem behind Max Ryabinin’s work on breaking memory barriers in context parallelism. Ryabinin introduces himself as Together AI’s VP of Research and Development and places the project within its training and customization work. Together supplies GPU clusters for model creation, fine-tuning and reinforcement learning for model shaping, and serverless or dedicated inference. At the recording, he describes an inference portfolio of more than 200 models, but inference optimization is outside this discussion.

Agents create demand for longer histories that models can actually use. Video creates another pressure: multiple frames per second quickly consume tokens, while temporal consistency requires access to events seconds or minutes earlier. Those capabilities require models to process long contexts during training. Even when a workload never reaches millions of tokens, understanding where training memory goes can uncover capacity that can be reinvested in speed.

Extending a transformer’s sequence exposes two different bottlenecks. Attention computes pairwise interactions across sequence positions, so its computation grows quadratically. Separately, the sequence-dependent activations that must be stored grow with context length. Linear growth sounds manageable beside quadratic work, but sufficiently long tensors still exhaust GPU memory.

These are distinct constraints: avoiding a stored quadratic attention matrix does not remove the linear-sized tensors elsewhere in training. Ryabinin illustrates the memory pressure with a Hugging Face training-memory chart before asking how far established techniques, followed by additional optimizations, can push the limit.

The worked target is a three-million-token training sequence on one node containing eight H100 GPUs. Ryabinin calls the architecture Llama 3B in the spoken introduction; Figure 2 of the matching Untied Ulysses paper identifies the example as Llama3-8B. That distinction matters when interpreting the memory budget. He describes the initial configuration as running out of memory while placing the model, but the account does not establish that three-billion-parameter weights alone exceed the node’s capacity.

The first intervention is fully sharded data parallelism, or FSDP. Instead of keeping the model’s parameters replicated on every GPU, it shards them across the eight devices. Model memory falls, but the training configuration still runs out of memory because attention activations remain too large. Reducing parameter storage has exposed the next allocation that dominates the budget.

DeepSpeed Ulysses, introduced by Microsoft, applies parallelism to the context itself. GPUs begin with portions of the sequence. All-to-all communication rearranges the attention inputs so that each GPU receives the full sequence for an assigned subset of heads. Each device computes those heads, and a reverse redistribution returns the results to the sequence-partitioned layout. The heads are distributed; their access to the sequence is preserved.

That arrangement can use optimized attention kernels such as FlashAttention inside each device’s assigned computation, then aggregate the results. Ryabinin names FlashAttention versions one through four as examples; the measured throughput discussed later uses FlashAttention 3. In the eight-GPU walkthrough, Ryabinin reports approximately an eightfold reduction in memory utilization after applying Ulysses, still short of making the target fit on one node.

Activation checkpointing changes what training saves for the backward pass. Rather than retaining every intermediate activation, it retains enough information to recompute the needed values later. The tradeoff is additional computation, so checkpoint boundaries must be chosen to avoid excessive recomputation. In the worked example, Ryabinin reports a further roughly eightfold reduction in activation memory from checkpointing. That still leaves another storage problem.

Checkpointing needs inputs from which to reconstruct a transformer block’s activations. Those inputs need not remain on the GPU while other layers execute. They can be offloaded to CPU memory, then prefetched before backpropagation reaches the corresponding layer. Overlapping transfers with useful work limits their performance impact. Ryabinin credits Unsloth’s checkpoint offloading as the first such implementation to the team’s knowledge; that is a qualified attribution, not an established priority claim for activation offloading generally.

After offloading, Ryabinin puts the displayed memory usage near 37 GB, without specifying whether that figure is per device, aggregate, or a subset of allocations. Other operations can still create an out-of-memory peak. In particular, an MLP or loss computation can materialize an intermediate with a dimension spanning all three million sequence positions.

The next intervention is tiling across sequence positions. An MLP operates independently on each token’s feature vector, so it can process slices of the sequence instead of expanding all tokens into its larger hidden dimension at once. Loss computation can likewise be chunked, with the appropriate reduction across chunks. This does not make attention local: it applies to operations that do not require interactions between sequence positions.

A PyTorch implementation can combine sequence tiling with checkpointing so that the expanded MLP intermediates are recomputed during backward rather than retained for every tile:

import torch
from torch import nn
from torch.utils.checkpoint import checkpoint


def tiled_mlp(
    x: torch.Tensor,
    mlp: nn.Module,
    chunk_tokens: int = 4096,
) -> torch.Tensor:
    # x: [batch, sequence, hidden]
    if chunk_tokens <= 0:
        raise ValueError("chunk_tokens must be positive")

    outputs = [
        checkpoint(mlp, tile, use_reentrant=False)
        for tile in x.split(chunk_tokens, dim=1)
    ]
    return torch.cat(outputs, dim=1)

The output still spans the full sequence, and concatenation itself needs storage. The saving comes from limiting the MLP’s expanded intermediates, not eliminating every sequence-sized tensor. In the walkthrough, stacking tiling on the earlier techniques finally makes the three-million-token target fit.

Going beyond that target requires another look inside context parallelism. The additional method is Untied Ulysses, called UPipe in the paper and later in the talk. The key observation is that one group of attention heads can already saturate the GPU’s computational capacity. If a device has several groups assigned to it, allocating working space for all of them at once may consume memory without providing a corresponding throughput benefit.

UPipe processes those groups over time:

1. Recompute the attention inputs for one group of heads.
2. Compute attention for that group over the full sequence.
3. Store its partial result.
4. Advance to the next group, reusing the working buffers from the previous group.

The allocation becomes smaller because it serves one group at a time instead of every assigned head simultaneously. The earlier MLP tiling divided work along the sequence dimension; this optimization divides attention work along the head dimension while preserving full-sequence attention.

The benefit depends on whether each smaller group still supplies enough work to keep the GPU busy. Ryabinin reports additional activation-memory savings without a significant throughput impact at smaller scales, then turns to measurements across context-parallel training implementations.

Ryabinin describes competitive throughput at both eight-billion- and 32-billion-parameter scales, sometimes with better performance at shorter contexts. The five-million-token result needs a more specific reading: it applies to Llama3-8B, not to both model sizes. The original paper reports Llama3-8B training at five million tokens and 98.25 tokens/s/GPU on eight 80 GiB H100s. That measurement uses BF16, FlashAttention 3, FSDP, checkpoint offloading, Ulysses degree U=8, chunk count C=8, and 1.9 TiB of host RAM; offloading is unpinned at five million tokens. The paper’s Qwen3-32B configuration uses 16 H100s and runs out of memory at five million tokens. These results measure training feasibility and throughput, not successful reasoning or retrieval over a five-million-token input.

Head-group size provides a practical tuning choice:

Larger chunks consume more memory but can run the model faster. Smaller chunks are useful when the next allocation would otherwise exceed capacity. The goal is not the smallest possible buffer; it is enough memory savings while keeping each chunk computationally efficient.

UPipe therefore adds another adjustable layer to the existing stack. At a context length that already fits, the saved memory can be used elsewhere in training. At the larger target, it can make the difference between fitting and failing. The final comparison slide identifies Llama 3–8B at five million tokens on eight H100s: the first six configurations are marked out of memory, while the UPipe configuration’s stacked bar remains below 80 GiB. The complete stack matters—parameter sharding, context parallelism, checkpointing, offloading, tiling, and headwise buffer reuse address different allocations.

Long-context training repeatedly shifts the bottleneck. Once one allocation is reduced, another that previously looked secondary can become the reason training fails. Ryabinin recommends tools such as PyTorch Profiler to find those unexpected costs; the paper discusses profiling in detail. At the recording, the paper and additional results are public, and he announces a forthcoming thread explaining the method more deeply.

The audience Q&A returns to the tensors themselves. Asked whether QKV refers to quantization parameters, Ryabinin clarifies that it means the transformer layer’s query, key, and value matrices. Query–key interactions produce the pairwise attention computation; QKV is not a quantization setting.

With a three-million-token sequence, even an intermediate tensor that has three million entries along just one axis is a substantial allocation. That is why attention buffer reuse cannot carry the entire training system by itself. The computation needs several approaches working together, each controlling which tensors exist, where they live, and how long they remain allocated. The engineering problem is to execute the required computation without requiring all of its intermediate state to fit on the GPU at once.