#DAG packing
Efficiently packing sequences with common prefixes
One of the essential optimization steps in pretraining is simple document packing, which allows a constant context length (required for kernel optimization) to be maximally utilized by packing multiple documents into it while effectively training on them separately via a block-diagonal causal attention mask that keeps tokens from different documents from attending to each other. Developing this idea further, we can address repeating token sequences, which are common in boilerplate text/code and (more recently) in branching autonomous agents — and it turns out the same logic applies to supervised fine-tuning of language models on agent session data (more on that later).
Let's imagine a use case that is quite different from predicting text: a world model trying to predict the next state of a simple game (yes its ARC-AGI-3, yes I have spent an unhealthy amount of time on this benchmark).
Here are a few sequences from the game ls20. One thing we might notice is that this representation is not very efficient: game state 1 in each sequence is identical, and game states 2 and 3 in sequences 1 and 2 are also identical — the two only diverge at the fourth state. Sequence 3 went off after a different first action. Since we have so neatly already represented each sequence as a chain, let's use some more proper notation and convert this into a DAG where the nodes are the game states and the edges are the actions that connect them. (Strictly, a node is a state paired with the action that produced it; why this distinction matters, and what a node even is for a language model, is in the appendix. Not every use case has actions as edges.)
Ok, that looks a lot better! But this is nothing new, I fully acknowledge that. What is new is what we can do next. Let's assume we are training a model on a corpus of many sequences just like these ones. The classic approach to document packing would be a fixed context window, and if a single sequence doesn't fill it, padding it with another sequence or its part.
This is good enough for most use cases, but one thing it doesn't consider is that almost all of these sequences start with the same few steps, and all of them start from the same start state. If we were working on a compression algorithm trying to fit as many sequences into as few context windows as possible, this would be very redundant. And since our model is, in essence, trying to represent the full transition graph of the game — the graph of all states and actions, i.e. all possible behavior — in a limited amount of parameters, seeing this redundancy in training is not very efficient either.
So let's work with the DAG directly, and try to fit as many of its tokens into a context window as we can.
The walk
Instead of packing whole sequences, we walk the DAG. The version I found practical is a depth-first walk: we keep a stack of nodes along the current path. At each step, we look at the node on top of the stack, pick one of its outgoing transitions we haven't walked yet, and emit one training position for it — the node as input. Then we move to the child and push it onto the stack. When the node on top has no unexplored transitions left, we pop it and backtrack. We stop when the window is full.
One thing worth pointing out: the walk crosses every shared node exactly once, so the duplicated prefixes from the classic approach disappear on their own — we never chose to "remove tokens from later sequences", we simply stopped paying for them twice. The one exception is the last common node, the node where sequences diverge: the walk copies it, so it appears once as input per child.
Why it has to: each copy now predicts a different next token, and if we merged them into one position the model would face contradictory targets for the same input (more on targets below).
(Note: this walk is just the version I found practical, not the optimal packing of a DAG into windows. Any walk that crosses every edge would do; I suspect the optimal one looks different, but I have not found it.)
The mask
One problem the walk exposes: with a plain causal mask, tokens from one branch would attend to tokens of unrelated branches sitting next to them in the flat sequence — the same cross-contamination document packing carefully avoids, except now it happens at nearly every position, because the boundaries are everywhere.
So we change the mask: each token may attend only to itself and to the nodes on its own path — the positions of its ancestors along the current walk. In the simple example above, this reduces to what we wanted anyway: the later tokens attend back to the shared prefix and to nothing else. And the copies of the last common node are isolated from each other: they look identical, but each attends only along its own path — see the mask in the figure above.
(No compute is wasted on predicting the shared prefix multiple times, and the later tokens are still properly contextualized — they did not lose any information.)
Positions
One more thing to consider: positional encodings. We just made tokens attend to positions that are disjointed in the flat sequence, and we don't want their positional encodings to remain as if they are far apart. So instead of flat indices, we use the graph depth — the number of transitions from the root. Every valid transition then has relative distance 1, no matter where in the window it physically sits.
Targets
Now, this is an optimized input, but we still need a target. In standard pretraining we would take the sequence without the last token and shift it by one. Here we don't need to: the walk already emitted, at every step, the parent node as input and the child node as target. Every position in the window is a complete training pair, so there is no shifting and no loss masking — every token is a valid example. This is also where copying the last common node pays off: each copy predicts its own child, so no position ever faces contradictory targets for the same input. In the packed-window figure, every column is already a complete (X, Y) pair.
Sneak peek: language models
Everything above leaned on actions as edges, but the DAG itself doesn't require them: a language model is the degenerate case where the node is just the token and there is no edge label at all. That makes the same machinery directly applicable to supervised fine-tuning — and agent data is where it gets interesting.
Agents gather a lot of sequential data, and modern agent harnesses (specifically the pi-coding-agent) routinely branch coding agents onto different tasks from a common context: you explore a codebase or do research once, then spawn several agents that all start from that shared context but do different things with it. As a DAG, that shared exploration is one common prefix, and each spawned agent is a sequence diverging from the last common node — exactly the shape above. Fine-tuning for a specific harness helps even more: almost every session shares the same (fairly large) system prompt, which document packing pays for in every single training example, and DAG packing pays for once.
We can visualize this the same way as the ARC-AGI-3 stages above: take the events from a
saved session in ~/.pi/agent/sessions/ and draw the tree the same way.
(Figure TODO: a DAG drawn from the events of a saved pi session, styled like the figures above. To generate: run a dummy-task session with a small model — deepseek v4 flash at low reasoning — and render it once the visualizations exist.)
For the full write-up — the precise packing algorithm, mask construction, and pseudocode — go to this page.
Appendix
What exactly is a node?
Throughout the main text I called the nodes game states and put the actions on the edges, with a footnote that strictly a node is a state paired with the action that produced it. Here is why that pairing matters: two different actions can lead between the same pair of states, and since the node itself is the input token, an action not folded into the node makes two genuinely different transitions look like one indistinguishable input. Folding the action in keeps the transition deterministic — from a given node, each outgoing edge is a distinct (action, next state) choice — and the model never has to guess what it just did.
The general rule: a node is the minimal information the next-node distribution conditions on. For a game that is the (state, action) pair. For a language model there is no separate action to condition on — the history is the prefix itself — so the node is just the token, and there is no edge label.
When is merging valid?
Merging two paths that arrive at the same node asserts that the next-node distribution depends only on the node itself. If the next step depends on history beyond the node, the merge is invalid and we should not make it — the DAG degrades to a tree, and everything above still holds unchanged.
For a language model this question answers itself: the model conditions on the token prefix, so two sequences that share a token prefix genuinely share their context, and merging them is valid by construction — the DAG over text is just a trie of the corpus. For game environments this is the Markov property of the environment, and it can genuinely fail; the fix is the same as above — don't merge, keep the DAG a tree.
Assumptions made in the main text
- Single root. The example has all sequences starting from the same start state. A general corpus has many roots: either pack each DAG of the forest independently, or add a virtual root joining them (each start state then sits at depth 1).
- Merging is valid. The Markov assumption above; if it fails, keep a tree.
- Every transition is worth predicting. The no-shifting, no-loss-masking claim assumes all edges are valid examples. Transitions that are noise can be dropped at graph-build time if needed.
- Graph depth is a meaningful distance. Position IDs assume that adjacent nodes should have relative distance 1. If the notion of distance you care about is something else, this needs rethinking.
- The walk is practical, not optimal. Noted in the walk section; the optimal packing of a DAG into windows is an open question here.
- Fixed context length. Inherited from document packing (kernel optimization).