AI Engineer Europe 2026
Road to 5 Million Tokens: Breaking Barriers in Long Context Training — Max Ryabinin, Together AI
Read the talk
The Road to Five Million Tokens Runs Through Activation Memory
Long-context training becomes feasible by shrinking what must live on a GPU at once: first parameters, then saved activations, intermediate buffers, and finally groups of attention heads.
From a talk by Max Ryabinin
Before you start: Familiarity with transformer attention, GPU memory, and the forward and backward passes of neural-network training will help.
What does it take to train on a much longer context?
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.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Quadratic work, growing live memory
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.
| Bottleneck | Source | Consequence |
|---|---|---|
| Quadratic computation | Pairwise attention interactions | More work as the sequence grows |
| Linear memory growth | Sequence-dependent activations | Larger live allocations |
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.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Recompute, offload, then tile
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:
python
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.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Reuse the attention buffer across head groups
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:
- Recompute the attention inputs for one group of heads.
- Compute attention for that group over the full sequence.
- Store its partial result.
- 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.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
What the five-million-token result measures
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:
| Heads computed together | Working memory | Throughput tendency |
|---|---|---|
| Fewer | Lower | May be lower |
| More | Higher | Can be higher |
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.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Find the allocation that becomes dominant next
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.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
QKV means query, key, and value
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.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Resources
From the talk
Original paper explaining headwise chunking, memory reuse and long-context training benchmarks.
Explains how all-to-all communication converts sequence partitions into attention-head partitions and back.
Unsloth's account of asynchronous activation offloading, with long-context fine-tuning examples and benchmark conditions.
Current instructions for measuring operator execution time, tracking memory and exporting traces.
Further reading
Official implementation with installation instructions and single-node and multinode benchmark scripts for Llama 8B and Qwen3 32B.
- The Ultra-Scale PlaybookArticle
A detailed guide to distributed model training, memory accounting and parallelism strategies.
Read the complete timestamped transcript
- 0:00
[upbeat music] Hi, everyone.
- 0:15
My name is Max. I am VP of Research and Development at Together AI, and today I'm going to tell you about our research project, which is called Road to Five Million Sequence Length: Breaking Memory Barriers in Context Parallelism.
- 0:30
So to begin, uh, I'll first say a few words about Together AI and who we are. Together AI is an AI-native cloud which provides, uh, services and infrastructure for AI developers and builders at all stages of development, starting from, uh, creating a model where you just might need a GPU cluster with heavily optimized computes and, uh,
- 0:55
highly reliable, uh, systems, uh, all the way through model shaping, where you can take existing models and customize them for your tasks in terms of performance, in terms of speed, in terms of quality through services such as the fine-tuning or, uh, uh, reinforcement learning.
- 1:15
Also, we are an inference provider, so if you have an app which is reliant on open source model inference, uh, you can, uh, work with us, and we'll provide you with the fastest way to launch and use AI models with more than 200 models in our portfolio, uh, options for deployment, which include serverless and
- 1:40
dedicated inference, and a ton of advanced optimizations, which I will not be able to speak about today. Um, the purpose of this talk is, uh, focused on model training, customization, and fine-tuning in particular.
- 1:56
And, uh, I'll start by asking a question. Uh, I think in the last few months or, like, at least a year or so, we're seeing a lot of interest in the community, uh, both on the system side and on the research side in training long context models.
- 2:13
Um, the primary reasons for that are twofold, I would say. First of all, with the, uh, like, explosion in popularity of agents, you can see a lot of different applications where you might want to put as many tokens as you want in your context, and you want the model to leverage that context effectively.
- 2:35
Um, second, with the development of applications such as video generation, you might often need to keep track of multiple, uh, like, uh, different frames, which... or, like, even, uh, multiple frames per second, which can, um, occupy quite a few tokens in your context pretty quickly.
- 2:57
And you also need models that have good, uh, sense of, uh, temporal consistency, which means that they're able to see what was happening a few seconds or ideally a few minutes ago.
- 3:10
Um, to do that all effectively, you need to make sure that the models are able to process that context and work with it correctly at the training time. Um, but even if you're not at the scales of, like, millions of, uh, tokens in the context length, it's still quite important to understand where the memory goes, because
- 3:35
who knows, maybe you might be able to reinvest it in some other ways and speed up your training overall.
- 3:42
So the problem here is that if you are taking a standard transformer-based language model and trying to extend its context, you can run [coughs] into two bottlenecks.
- 3:57
Bottleneck number one is that you are faced with quadratic computation because, well, long story short, for transformer-based models, uh, you have pairwise interactions a-across all the elements in a sequence.
- 4:11
The second problem is more insidious, one might say. Uh, as you continue scaling your context, your memory keeps growing linearly, which is not as bad, but still pretty difficult to deal with unless you apply a range of specific techniques.
- 4:30
Um, and this is an example from Hugging Face's blog post on model training, which shows that the sequence length, uh, growth can affect your memory limits pretty considerably.
- 4:46
Um, yeah, here's the slide. And our goal of that project was to see, uh, how far exactly we are able to get with a range of existing techniques that are pretty well known to some in the community, as well as some further optimizations that we wanted to leverage to push this a bit, uh, further ahead.
- 5:11
So, um, let's say you're taking a model which is a standard Llama 3B architecture. You're trying to fit three million training tokens into your context, and you're taking, uh, all of this on an 8, uh, X, uh, H100 GPU node.
- 5:30
Um, the first stage you'll see is that even with, uh, just the model parameters, you're not able to fit, uh, it into the GPU. Uh, you run out of memory just by trying to place the model.
- 5:43
Of course, the next stage is to apply fully sharded data parallelism, where all the parameters are, like, basically chunked across the eight GPUs that you have, which is great, but, uh, still doesn't solve the problem.
- 5:57
You see that, uh, the memory usage for the model drops quite significantly, but you still are running out of memory because of all the attention activations.
- 6:08
The next point that, uh, we leveraged and I encourage you to use as well is, uh, taking advantage of, uh, context parallelism. In particular, there is a pretty well-known technique called DeepSpeed Ulysses, first, uh, introduced by Microsoft.
- 6:26
The idea is that instead of computing, uh, all of your multi-head attention, like, on every GPU separately for the whole sequence, you can do something more clever. In particular, you can try to compute the attention for different heads at different points in time, uh, or, like, on different GPUs, uh, through communicating these
- 6:50
activations as, uh, they are required in such a way that one GPU is only responsible for one attention head here, but it's still computing the attention over the whole sequence.
- 7:04
Um, that technique is quite effective at, uh, addressing the problem, and it also allows you to utilize the best possible attention implementation, like Flash Attention one, two, three, four, um, to optimize that part of the computation, and then you aggregate the results as, uh, you would've, uh, previously.
- 7:27
So if you apply Ulysses context parallelism, the utilization drops quite significantly, like approximately 8X, uh, here, as, uh, it should. But we are still quite far from our goal of being able to fit that onto just a single, uh, H100 node.
- 7:48
So what happens next is that we can try to recompute the activations as they, uh, are needed to us at the backward pass. That technique is known as activation checkpointing, and, uh, it's available in pretty much all of the different frameworks these days that you could use.
- 8:08
You just need to enable it in a correct way that does not, uh, impose too much of a computational burden on you.
- 8:18
Um, with that, with activation checkpointing, you can drop the activation usage by, like, a further factor of eight, but still something else needs to be done. Uh, the next optimization is, um, also connected to the storage of activations.
- 8:38
You can try to, uh, store some of the inputs to each transformer block, not on the GPU, but instead offload them to CPU when they are not required. Uh, this is not very impactful for the performance because you can offload the...
- 8:58
it, uh, and prefetch when you are trying to back propagate to that, to the corresponding layer. Um, this optimization, to the best of our knowledge, was first implemented, uh, by Unsloth, and, uh, it allows you to drastically expand the context window.
- 9:16
Um, the next point is that you're getting, uh, with offloading next to 37, uh, gigabytes of data, but then comes the other part of out-of-memory usage. Uh, what happens next is that you can essentially tile all, all of the computations across the sequence length in case they are element-wise.
- 9:39
So all the loss computations, all the MLPs, they can be chunked to avoid creating these huge buffers that would be 3 million along one of the dimensions. Um, that's our tech sequence length training, and even with these optimizations, you're finally getting to a point where 3 million is possible.
- 10:00
But what if you wanted to go further? And here we actually need to do something else, which is the primary, uh, optimization we've done in our work, dubbed the Untitled Ulysses, and you could describe it as a further deeper, uh, analysis and expansion of, uh, this context parallelism technique.
- 10:26
So what we found was that even trying to compute one set of heads at a time, uh, is already enough to saturate the computational capacity of the GPU with- within one iteration, which means that if you have multiple different heads scheduled to be executed on, uh, one GPU, you can divide it in chunks and then, uh,
- 10:51
essentially iterate through these chunks over time. So you have one group of heads which are being recomputed, then you compute attention over them, you store the partial result, then you follow up with the next stage, which can reuse all the buffers you've allocated at the previous stage.
- 11:11
Um, so the advantage here is that instead of allocating this huge buffer as, uh, you would've before, like here in this slide, you allocate a buffer which is smaller, but you reuse it across, like, two or more different iterations.
- 11:28
And that allows you to further save on the activation memory for your training without any significant impact to the throughput at smaller scales. So here you can see the results that we've, uh, measured across different context parallelism techniques.
- 11:46
As you can see, both at the 8 billion scale and at the 32 billion scale, we are matching quite closely the most memory optimized implementations of transformer training while being able to scale even further, like 5 million tokens, and, uh, sometimes even being more performant at small, uh, at shorter context lengths.
- 12:10
Um, the relation between the chunk size or the number of heads you compute at the same time and the throughput is quite straightforward. So if your chunk is larger, your memory utilization is higher, but at the same time, you can run the whole model a bit faster.
- 12:29
Um, so by stacking all of these, uh, techniques together and applying UPipe on top, you can, for example, uh, free up a bit of additional memory in, uh, your training if you need it and reinvest it somewhere else, uh, for example, among the stages.
- 12:46
Or, uh, you could say that we're interested in training across not 3 by 5 million context lengths, and then, uh, UPipe is the technique that will save you, uh, by contrast to everything else.
- 13:00
So, um, as a takeaway, uh, I think one of the things that, uh, could be quite insightful here is that training models with large context lengths is a very interesting and challenging goal.
- 13:16
Uh, but the bottlenecks might appear where you least expect. So tooling like the PyTorch Profiler, which, uh, we, uh, elaborate on a ton in our paper, uh, or other te- or other techniques can help you a lot.
- 13:31
Uh, and also check out our paper for more results. All of that is public at the moment, and we have an upcoming thread which will illustrate the method in more depth.
- 13:43
Thank you very much for listening, and now we are ready for questions.
- 13:52
Thank you. Do you guys have any questions 'cause we're Together AI employees? [laughs]
- 14:01
Sorry, we, we joined in from the middle, so we, we lack some certain context.
- 14:06
Yeah. Got it.
- 14:07
But I'm just curious about the, the... So the QKV-
- 14:11
Yeah
- 14:11
... that was, uh, quantization, uh, parameters. Is that correctly understood?
- 14:15
Uh, not exactly. It was just the query, key, and value matrices-
- 14:19
Okay
- 14:20
... of the transformer layer. So, uh, you multiply them here in the attention part, which, uh, creates most of the complexity because all of the queries have to, like, result in all of these pairwise activ- uh, interactions with, uh, uh, like, keys.
- 14:37
Uh, and the problem is that if you have a se- [coughs]
- 14:42
a sequence which is, like, 3 million in length, it means that technically, like, in the standard most, uh, vanilla way, you would've just, like, allocated that whole big tensor which has 3 million in...
- 14:55
Which is 3 million in size, uh, along one of the axes. And, like, that's pretty significant, as you could imagine, which means that you have to resort to, like, not just one technique, which is UPipe, but a range of other approaches to somehow help you, like, uh, lever- like, execute these computations without running out of your memory.
- 15:19
So yeah, that's the key idea and the key challenge of, uh, working with transformers at, uh, this scale.
- 15:30
Cool. Um, in that case, thank you very much for questions and for listening. [outro music]