feat: anchored EAGLE-3 training - #186
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7326fa6ed3
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| step_mask = _gather(loss_mask, positions) * keep_mask | ||
| step_target = _gather_target(target, positions) |
There was a problem hiding this comment.
When draft_vocab_size < vocab_size, compute_target_p_padded() uses PrecomputedTarget.position_mask to exclude positions whose full-vocabulary argmax is unavailable to the draft, and the dense Eagle3 path honors that mask. This path instead samples and masks solely with loss_mask, while _gather_target() also drops position_mask, so impossible targets silently contribute a renormalized KL loss and bias training. Use the target position mask as the effective mask for anchor sampling and each depth.
Useful? React with 👍 / 👎.
7326fa6 to
3faeb79
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3faeb79f4a
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
|
||
| # An anchor is only usable if its own position and the position right |
There was a problem hiding this comment.
When a dataset contains isolated supervised tokens or one-token assistant responses, this DFlash-specific adjacency predicate produces no Eagle anchors even though dense Eagle3 would train its depth-0 prediction; keep_mask is therefore entirely false and every returned loss is silently zero. More generally, it excludes the final supervised position of every response from depth 0 and biases the anchored objective. The Eagle path needs its own candidate predicate based on the positions valid for each TTT depth rather than reusing the DFlash requirement that both anchor and anchor + 1 be supervised.
Useful? React with 👍 / 👎.
3faeb79 to
a63726c
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a63726ce6b
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| lk_eta=getattr(self.args, "lk_eta", 3.0), | ||
| ) |
There was a problem hiding this comment.
When eagle3_num_anchors is combined with attention_backend="usp", the USP data fetcher supplies a sequence shard per rank, but AnchoredEagle3Model.forward bypasses the draft's USP attention implementation, ignores the supplied global position_ids, and runs local SDPA over that shard only. Anchors therefore cannot attend to preceding context held by other sequence-parallel ranks, silently corrupting training; reject this combination or implement the required USP communication and global-position handling before selecting the anchored model.
Useful? React with 👍 / 👎.
a63726c to
dba6e35
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: dba6e35101
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| hidden = current + attn.o_proj(attn_out) | ||
| hidden = hidden + layer.mlp(layer.post_attention_layernorm(hidden)) |
There was a problem hiding this comment.
When training.gradient_checkpointing=true, the anchored path still executes the attention and MLP for every depth directly, whereas the dense Eagle path checkpoints its backbone invocation. The option therefore only checkpoints the loss helper here and silently retains all layer activations across the recurrent depths, which can negate the expected memory reduction or cause long-sequence runs to OOM. Wrap the per-depth layer computation in checkpointing when the flag is enabled.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Anchoring already saves significant memory, so not a P0.
dba6e35 to
a719622
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a7196221b9
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| return sample_anchor_positions( | ||
| seq_len, loss_mask, self.num_anchors, self.block_size, device | ||
| ) |
There was a problem hiding this comment.
Delegating DFlash to this strided sampler silently changes its established random-anchor objective. For example, with 700 valid candidates and 512 anchors, valid_counts // num_anchors is 1, so each step selects a contiguous 512-token window; central candidates are selected nearly every time while edge candidates are rarely selected, instead of every candidate having probability 512/700 as in the previous implementation. This biases DFlash training on common sequences with between one and two anchor budgets of supervision; keep uniform random sampling for DFlash or make the new spacing behavior explicitly opt-in.
Useful? React with 👍 / 👎.
a719622 to
ab79f56
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ab79f568a9
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| cos, sin = attn.rotary_emb(key_cache, seq_len=seq_len + self.length) | ||
| cos, sin = cos.to(device), sin.to(device) | ||
| context_positions = torch.arange(seq_len, device=device).expand(bsz, -1) | ||
| key_cache, _ = apply_rotary_pos_emb(key_cache, key_cache, cos, sin, context_positions) |
There was a problem hiding this comment.
When anchoring is enabled for a supported Llama-style draft whose rope_scaling uses mrope, this call fails immediately because LlamaMutiRotaryEmbedding.forward accepts position_ids, not the seq_len keyword. The normal attention implementations also use apply_multimodal_rotary_pos_emb with the configured MRoPE sections, whereas this path always applies standard rotary embedding and ignores the supplied multidimensional position_ids; either implement that branch or reject MRoPE drafts during initialization.
Useful? React with 👍 / 👎.
ab79f56 to
b0f7ba9
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b0f7ba991b
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
b0f7ba9 to
dea977e
Compare
Add training.eagle3_num_anchors. When set, the TTT unroll is evaluated at N sampled anchor positions per sample instead of at every supervised position, so each depth costs O(N) rather than O(sequence length). Off by default; the dense Eagle3 path is untouched. The draft is a single layer, so its depth-0 keys and values are a pointwise projection of the layer input and can be built for the whole sequence without attention. Only the queries are gathered. KV is laid out as [context | chain depth 1 | ...] with N slots per depth; an anchor sees the context up to and including its own position plus its own slot in each chain block, and anchors never see each other. test_context_rows_match_the_dense_ eagle3_mask asserts a row equals the dense Eagle3 row for that position. Masking is a materialised [B, 1, N, KV] boolean passed to scaled_dot_product_attention. Anchored attention is small enough that a BlockMask costs more to build than the density it removes: at S=4096 with 4 depths the build alone is 2.1ms against 1.2ms for the whole SDPA path. Anchor sampling and the per-depth mask use the target position mask when the draft vocabulary is pruned, matching the dense path -- those positions have no representable argmax, so training on them fits a renormalized target the draft cannot emit. MLA drafts work alongside the usual k_proj/v_proj shape. Their keys come from the compressed kv_a/kv_b pair, which is still pointwise in the layer input, so the context K/V is built the same way; they rotate only the rope-side dims and carry their own softmax scale. The rotary table is built once over the whole sequence rather than through the draft's own helper, which sizes it from the query length -- here the anchor count, so every position past it would index off the end. configs/vllm_qwen3_8b_anchored.yaml uses an MLA draft by default. test_per_depth_losses_match_dense covers both draft shapes, and a separate case covers num_anchors far below seq_len -- the regime where that table sizing goes wrong and an equal-sized test cannot see it. MRoPE drafts are rejected: they rotate through apply_multimodal_rotary_pos_emb with the configured mrope_section rather than the standard rotary this path applies. Anchors are sampled with a spacing cap. training.eagle3_anchor_max_gap is how many supervised positions may sit between neighbours, so 0 puts them side by side. It matters because anchor a supervises token a+d at depth d: a token is seen at depth d only if the anchor d positions behind it was also picked, so anchors g apart leave each token ttt_length/(g+1) of its depths, and once g reaches ttt_length every token is pinned to a single depth. None (the default) means block_size, which spreads anchors as evenly as the budget allows. Gaps count supervised positions rather than raw tokens, so no slot is spent on an unsupervised region and anchors stay token-adjacent inside a span. Qwen3-8B, 10k PerfectBlend, 5 epochs, N=512: max_gap=0 reaches the dense simulated acceptance length of 0.78 at 1.09x the samples/s, where sampling without a cap plateaus at 0.47 from a third of the way in. The gap is not data volume -- both see the same 512 slots and the same tokens -- it is that isolated anchors never show a token more than one depth. The anchor sampler moves out of models/dflash.py to models/ops/anchors.py so both paths share it. DFlash picks up the default cap, which replaces its uniformly random draw with an evenly strided one at the same density: same expected depth coverage, without the isolated anchors that random gaps produce. Attention cost against dense TTT (TTT=4, N=256, one GB300, ms): seq 4096 dense 1.78 anchored 1.12 seq 6144 dense 3.82 anchored 1.10 seq 32768 dense 75.95 anchored 4.24 seq 131072 dense 1134.5 anchored 15.89 test_dflash's test_accuracy_improves trained for 30 steps, which left the accuracy in the noise: it improved on only 9 of 12 seeds, so the assertion turned on which anchors happened to be drawn rather than on whether the model learned. Both the old and new samplers score 9/12 there, so this is not a regression, but changing the sampler shifts the RNG stream and moves which seeds land where. At 100 steps every seed improves by at least 0.18 under either sampler. Claude-Session: https://claude.ai/code/session_01CqPZAxkfxK7vTCbBQniBni
dea977e to
e310c58
Compare
2 participants
Adds DFlash style
training.eagle3_num_anchors: the TTT unroll is evaluated at N sampled anchor positions per sample instead of at every supervised position, so each depth costs O(N) rather than O(sequence length). Off by default; the dense Eagle3 path is untouched.Why: Cheaper and uniform length-sequence training. Also a long sequence can dominate the loss heavily, therefore anchor sampling ensures each sample contributes roughly equally, removing the bias for model to lean towards longer samples.
How it works
Traditionally, EAGLE-3 trains with TTT rollout on every loss enabled token. The KV cache structure is shown below. This is expensive in terms of memory and runtime, especially for longer sequences.
DFlash has a different training method, it randomly samples "anchor tokens" and predicts N block continuations from each token. It achieves this by sorting each anchor token and create a bidirectional attention mask for each block prediction. This technique helps DFlash train stabily on longer sequences, since anchor token count doesn't change for longer sequences. Faster and potentially more stable training.
This PR combines both approaches, instead of doing TTT rollout on every token, we sample random anchor tokens to calculate loss from. Each token unsampled contributes through KV cache contributes anyway, so we don't need the full coverage to get similar results.
This means we calculate loss from significantly tokens, however results already show this doesn't hurt meaningfully. Below is an image of loss calculation coverage. The dark color means contributes more to the loss, since there is weight decay on same TTT (Same as DFlash)
Below is anchored, note that each token indirectly contributes via their KV cache anyway.
Since each token's weight is different, model can be trained for more epochs. In each epoch, data order will be random, so an un-seen token, or a seen but different TTT / weight token can contribute in the other epoch.
Runtime
Attention only (TTT=4, N=256, one GB300, ms / peak GB)
Dense time and memory grows quadratically, anchored is roughly flat. SDPA does not cost more memory, since the mask is built with
H=1and broadcast and SDPA tiles rather than materialising the score matrix.End-to-end training (Qwen3-8B, 20K open-perfectblend, 5 epochs)
Open-perfectblend doesn't have very long rows, so the benefits are minimal, but still visible.
DFlash is less data efficient, because predicted blocks are not overlapped with each other, so a token at position X is always treated as a token at (X mod block_size) position. Because position of a token also changes prediction hardness + loss weight, model doesn't see all tokens truly during training.
This has been validated by running EAGLE with non-overlapped anchors, the result was significantly worse than regular EAGLE, but it was better than DFlash and DFlash 2.
Also added block mask compilation for DFlash, which speeds up training by 10%,
Multi-layer sweep (N=1024, TTT=4, one GB300, ms / peak GB)
MLA
Real end to end run timees.