feat(eagle3): add Eagle 3.1 training support

[orestis-z]

Purpose

Adds Eagle 3.1 training support with two config flags (matching vLLM's config fields):

• --norm-before-fc: Apply a single RMSNorm to the concatenated target hidden states before the FC projection layer — i.e. Norm(concat(h₂, h₁₈, h₃₃)). This was already implemented on main for gpt-oss checkpoint compatibility; this PR adds the training CLI flag and wires it through P-EAGLE. Note: the original Eagle 3.1 paper specifies per-layer normalization (concat(Norm(h₂), Norm(h₁₈), Norm(h₃₃))), which corresponds to vLLM's fc_norm field. Per-layer norm support (--fc-norm) will be added in a follow-up PR.
• --norm-output: Feed post-norm hidden states back across TTT steps to stabilize magnitude drift across speculation depths.

Both flags default to False (opt-in).

Also fixes P-EAGLE missing input_norm before FC and wires norm_before_fc / norm_output through the converter and P-EAGLE's from_training_args.

vLLM compatibility

• norm_before_fc works out of the box with vLLM
• norm_output requires a one-line patch in vLLM's algos.py to propagate the config field (model code already supports it via getattr): [Speculative Decoding] Propagate norm_output and fc_norm config for Eagle3 speculators vllm#46488

Tests

pytest tests/unit/test_config.py -k "eagle3_config" -v    # config roundtrip
pytest tests/integration/models/test_model_forward.py -k "NormOutput" -v  # forward pass

• test_norm_output — both flags enabled, verifies input_norm + forward + backward
• test_norm_output_without_norm_before_fc — norm_output alone (no input_norm)
• test_peagle_norm_before_fc — P-EAGLE with norm_before_fc
• 4 config roundtrip tests (model_dump, to_dict, save/load_pretrained, defaults)

Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan/results, such as providing test command and pasting the results.
• The necessary documentation update.
• I (a human) have written or reviewed the code in this pr to the best of my ability.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 0df26bb3-9130-4afa-93c9-b655dd6f130c

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds norm_output and norm_before_fc normalization flags to Eagle3 and PEagle speculator models. Changes span Eagle3SpeculatorConfig field descriptions, Eagle3DraftModel and PEagleDraftModel forward logic, Eagle3Converter conversion API, training CLI argument parsing, and corresponding integration and unit tests.

Changes

Eagle3/PEagle norm_output and norm_before_fc

Layer / File(s)	Summary
Eagle3SpeculatorConfig field descriptions src/speculators/models/eagle3/config.py	Rewrites description strings for norm_before_fc and norm_output fields to reflect their actual behavior in Eagle3/PEagle architectures.
Eagle3DraftModel forward and from_training_args wiring src/speculators/models/eagle3/core.py	Branches forward logits computation on config.norm_output to normalize hidden states before lm_head when enabled; adds norm_output to Eagle3SpeculatorConfig constructed in from_training_args; removes a stale comment and expands the noqa annotation.
PEagleDraftModel forward and from_training_args wiring src/speculators/models/peagle/core.py	Conditionally applies self.input_norm to sampled_hidden before self.fc in forward; forwards norm_before_fc and norm_output kwargs into PEagleSpeculatorConfig in from_training_args.
Eagle3Converter norm parameters src/speculators/convert/eagle/eagle3_converter.py	Adds norm_before_fc and norm_output parameters to convert and _build_eagle3_speculator_config, threading them into the constructed Eagle3SpeculatorConfig with fallback to eagle_config values.
Training CLI and auto-defaulting logic scripts/train.py, docs/cli/train.md	Defines --norm-output CLI argument and updates --norm-before-fc help text in parse_args; adds a main() branch that auto-defaults both flags to True for eagle3/peagle with llama draft architecture; documents --norm-output in the CLI reference.
Integration and unit tests tests/integration/conftest.py, tests/integration/models/test_model_forward.py, tests/unit/test_config.py	Extends make_eagle3_model and make_peagle_model fixtures with norm_before_fc/norm_output parameters; adds TestNormOutputParams CUDA integration tests for Eagle3 and PEagle forward/backward; adds unit tests for Eagle3SpeculatorConfig.norm_output round-trip serialization and default value.

Possibly related PRs

• vllm-project/speculators#550: Introduced the make_eagle3_model and make_peagle_model integration test fixtures that this PR extends with norm_before_fc and norm_output parameters.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'feat(eagle3): add Eagle 3.1 training support' accurately summarizes the main objective of the PR, which is to add Eagle 3.1 training features to address attention drift through normalization flags.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The description matches the changeset, covering the new Eagle3/P-EAGLE normalization flags, docs, wiring, and tests.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mergify]

The quality checks have failed. Please run make style and make quality under
the root directory to address the lint failures. You will need to install the
dev optional install to get the required linting packages:
https://github.com/vllm-project/speculators/blob/main/CONTRIBUTING.md

[orestis-z]

Eagle 3.1 Evaluation Results

Trained Eagle 3.1 drafters for Qwen/Qwen3-8B on 508k samples (magpie + ultrachat from Qwen3-8B-Regenerated-Collection), 2 epochs, online mode — same setup as the Eagle 3 baselines from #563.

Acceptance Lengths

Evaluated on RedHatAI/speculator_benchmarks (9 subsets), k=5 draft tokens, temp=0.6, top_p=0.95, top_k=20, 80 requests per subset.

Subset	Eagle 3 (Llama)	Eagle 3 (Qwen3)	Eagle 3.1 (Llama)	Eagle 3.1 (Qwen3) 5e-4	Eagle 3.1 (Qwen3) 3e-4
HumanEval	3.242	3.368	3.484	3.432	3.396
Math Reasoning	3.479	3.626	3.691	3.613	3.617
QA	2.681	2.776	2.820	2.754	2.710
Question	2.838	3.035	3.112	3.014	2.988
RAG	2.707	2.838	2.960	2.789	2.771
Summarization	2.284	2.407	2.480	2.426	2.435
Tool Call	2.707	2.790	2.840	2.805	2.773
Translation	2.559	2.642	2.705	2.659	2.690
Writing	2.860	3.002	3.081	3.013	3.024
Average	2.817	2.943	3.019	2.945	2.934

Eagle 3.1 (Llama-arch) is the best overall at +7.2% over Eagle 3 (Llama). Combining Qwen3-arch with Eagle 3.1 does not compound — gains are redundant since both norm strategies address the same underlying magnitude drift problem. Lowering LR from 5e-4 to 3e-4 for Eagle 3.1 (Qwen3) does not help, confirming the redundancy is architectural, not LR-related.

Per-Position Acceptance Rates (averaged across subsets)

Position	Eagle 3 (Llama)	Eagle 3 (Qwen3)	Eagle 3.1 (Llama)	Eagle 3.1 (Qwen3) 5e-4	Eagle 3.1 (Qwen3) 3e-4
0	0.7120	0.7302	0.7192	0.7312	0.7324
1	0.4748	0.5024	0.5053	0.5040	0.5023
2	0.3107	0.3401	0.3561	0.3410	0.3376
3	0.1968	0.2239	0.2545	0.2240	0.2208
4	0.1231	0.1460	0.1842	0.1449	0.1408

Eagle 3.1 (Llama) improvement grows with draft depth, while Qwen3-arch and Eagle 3.1 (Qwen3) plateau.

fc_norm Ablation (per-hidden-state norm, paper approach)

Ran the same Llama-arch, LR=3e-4 config but with --fc-norm (per-hidden-state RMSNorm, as in the Eagle 3.1 paper) instead of --norm-before-fc (single RMSNorm over the concatenated vector).

Subset	Eagle 3 (baseline)	3.1 norm_before_fc	3.1 fc_norm	Δ norm_before_fc	Δ fc_norm
HumanEval	3.242	3.484	3.262	+7.5%	+0.6%
Math Reasoning	3.479	3.691	3.492	+6.1%	+0.4%
QA	2.681	2.820	2.579	+5.2%	-3.8%
Question	2.838	3.112	2.875	+9.6%	+1.3%
RAG	2.707	2.960	2.676	+9.4%	-1.1%
Summarization	2.284	2.480	2.328	+8.6%	+1.9%
Tool Call	2.707	2.840	2.686	+4.9%	-0.8%
Translation	2.559	2.705	2.569	+5.7%	+0.4%
Writing	2.860	3.081	2.870	+7.7%	+0.3%
Average	2.817	3.019	2.815	+7.2%	-0.1%

norm_before_fc gives +7.2%, while fc_norm shows no improvement. fc_norm checkpoint: eagle31-fcnorm-ckpt1.

Note on bf16 training and norm weights

All runs used --hidden-states-dtype bfloat16 (the default), which prevents RMSNorm weights from learning: Adam updates converge to ~±lr (~3e-4), but bf16 epsilon at 1.0 is ~8e-3, so norm weight updates round to zero. All norm weights across all checkpoints are exactly 1.0. The fix is --hidden-states-dtype float32.

However, this does not invalidate the results above. Any norm succeeded by a linear layer (FC, lm_head) is mathematically unaffected — the linear layer can absorb the diagonal scaling that the norm weights would have provided (W @ diag(w) is just another linear transform). Since input_norm (norm_before_fc) and fc_norm both feed exclusively into the FC layer, and FC weights (initialized near zero) have no bf16 precision issues, the FC learned the necessary scaling during training. The only norm potentially affected is self.norm when norm_output=True, since its output feeds both lm_head and the next TTT step's concatenation — but this is a shared factor across all Eagle 3.1 configs.

Training Config

	Eagle 3 (Llama)	Eagle 3 (Qwen3)	Eagle 3.1 (Llama)	Eagle 3.1 (Qwen3) 5e-4	Eagle 3.1 (Qwen3) 3e-4
Architecture	--draft-arch llama	--draft-arch qwen3	--draft-arch llama --norm-before-fc --norm-output	--draft-arch qwen3 --norm-before-fc --norm-output	--draft-arch qwen3 --norm-before-fc --norm-output
LR	3e-4	5e-4	3e-4	5e-4	3e-4
Model	llamaarch-ckpt1	qwen3arch-ckpt1	eagle31-llamaarch-ckpt1	eagle31-qwen3arch-ckpt1	eagle31-qwen3arch-3e4-ckpt1

• Dataset: Qwen3-8B-Regenerated-Collection (magpie + ultrachat), 508k samples
• 2 epochs, --draft-vocab-size 32000, cosine schedule
• 7 GPUs: 2 vLLM (DP=2) + 5 training (FSDP), online mode, ~26 hours per run
• Qwen3-arch serving requires vllm#43132
• Eagle 3.1 serving requires vllm#46488
• Eval reproduction script: eval_eagle31_comparison.sh

[orestis-z]

@fynnsu @shanjiaz — should we keep the single --eagle31 flag (simpler, but requires a vLLM-side mapping to norm_output/norm_before_fc), or split into --norm-before-fc + --norm-output to match vLLM's config fields directly?

[shanjiaz]

@fynnsu @shanjiaz — should we keep the single --eagle31 flag (simpler, but requires a vLLM-side mapping to norm_output/norm_before_fc), or split into --norm-before-fc + --norm-output to match vLLM's config fields directly?

I think I would prefer --norm-before-fc and --norm-output. Although if we're using llama arch should we default to eagle3.1?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @orestis-z.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[orestis-z]

Addressed the feedback — split --eagle31 into two independent flags matching vLLM's config fields:

• --norm-before-fc / --no-norm-before-fc: RMSNorm before FC projection
• --norm-output / --no-norm-output: post-norm hidden state feedback across TTT steps

Both default to `True` when `--draft-arch llama` is used with eagle3/peagle, since the benchmarks show Eagle 3.1 (Llama) is uniformly better (+7.9% avg acceptance length). For qwen3 arch, both default to `False` (plain Eagle 3). Can be explicitly overridden with --no-* flags.

Also rebased onto latest main.

vLLM side: norm_before_fc works out of the box. norm_output needs a one-liner in algos.py — the model code already reads it via getattr(self.config, "norm_output", False), just the config propagation is missing: vllm-project/vllm#46488

[coderabbitai]

Caution

Review failed

An error occurred during the review process. Please try again later.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[orestis-z]

Blocked by vllm-project/vllm#46488

[orestis-z]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/speculators/convert/eagle/eagle3_converter.py`:
- Line 139: The `norm_before_fc` parameter on line 139 uses OR fallback logic
that prevents explicit False values from overriding checkpoint settings, while
`norm_before_residual` on line 138 and `norm_output` on line 140 use parameters
directly without fallback. Either align all three fields by removing the OR
fallback from the `norm_before_fc` assignment to match the other fields pattern,
or if checkpoint preservation is intentionally required only for
`norm_before_fc`, add a clear comment explaining why this field requires
different handling than `norm_before_residual` and `norm_output`.

In `@tests/integration/models/test_model_forward.py`:
- Around line 359-398: In the TestNormOutputParams class, the test methods
test_norm_output, test_norm_output_without_norm_before_fc, and
test_peagle_norm_before_fc all unpack variables from model calls but do not use
them. Prefix the unused variables draft_tokens and metrics with an underscore in
all three methods where they are unpacked (changing them to _draft_tokens and
_metrics) to follow Python conventions for intentionally unused variables and
silence the Ruff linter warnings.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 43f11dbb-a456-4322-8d73-ab5cf4a45d1c

📥 Commits

Reviewing files that changed from the base of the PR and between 9700c0f and ee1af19.

📒 Files selected for processing (9)

• docs/cli/train.md
• scripts/train.py
• src/speculators/convert/eagle/eagle3_converter.py
• src/speculators/models/eagle3/config.py
• src/speculators/models/eagle3/core.py
• src/speculators/models/peagle/core.py
• tests/integration/conftest.py
• tests/integration/models/test_model_forward.py
• tests/unit/test_config.py

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @orestis-z.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[fynnsu]

Generally looks good but I think the default chose logic feels a little complicated. Let's maybe wait for the Eagle3.1 + Qwen run (with the changed LR) to see if we can simplify this.

[orestis-z]

Generally looks good but I think the default chose logic feels a little complicated. Let's maybe wait for the Eagle3.1 + Qwen run (with the changed LR) to see if we can simplify this.

What do you think about I update it to remove the auto-defaulting logic, make both flags default to False regardless of arch, and merge it as such? We can follow up with defaults if necessary after my experiments conclude.

[mergify]

The quality checks have failed. Please run make style and make quality under
the root directory to address the lint failures. You will need to install the
dev optional install to get the required linting packages:
https://github.com/vllm-project/speculators/blob/main/CONTRIBUTING.md

[orestis-z]

Eagle 3.1 (fc-norm) Evaluation Results — Qwen3-30B-A3B

Trained Eagle 3.1 drafter for Qwen/Qwen3-30B-A3B using the per-layer FC normalization from orestis-z#1 (--fc-norm --norm-output --draft-arch llama). Same eval setup as the Qwen3-8B results above.

Acceptance Lengths

Evaluated on RedHatAI/speculator_benchmarks (9 subsets), k=5 draft tokens, temp=0.6, top_p=0.95, top_k=20, 80 requests per subset.

Subset	Eagle 3 (Llama)	Eagle 3.1 fc-norm (Llama)	Delta
HumanEval	2.300	2.373	+3.2%
Math Reasoning	2.460	2.479	+0.8%
QA	2.000	2.013	+0.7%
Question	2.130	2.179	+2.3%
RAG	2.140	2.162	+1.0%
Summarization	1.940	1.944	+0.2%
Tool Call	2.160	2.167	+0.3%
Translation	2.120	2.153	+1.6%
Writing	2.150	2.152	+0.1%
Average	2.156	2.180	+1.1%

Per-Position Acceptance Rates (averaged across subsets)

Position	Eagle 3 (Llama)	Eagle 3.1 fc-norm (Llama)
0	0.5858	0.5890
1	0.2991	0.3028
2	0.1491	0.1558
3	0.0790	0.0846
4	0.0423	0.0481

Improvement grows with draft depth (pos 4: +13.7% relative).

Plots

Training Config

	Eagle 3 (Llama)	Eagle 3.1 fc-norm (Llama)
Architecture	--draft-arch llama	--draft-arch llama --fc-norm --norm-output
Model	eagle3-ckpt1	eagle31-fcnorm-ckpt1

• Dataset: sharegpt + ultrachat (~328k samples), 2 epochs, --draft-vocab-size 32000, cosine schedule, LR=3e-4
• 6x H100 80GB: 1 vLLM (TP=1) + 5 training (FSDP), online mode, ~3.5h per run
• Eagle 3.1 fc-norm uses per-layer RMSNorm from orestis-z/speculators#1
• vLLM serving required patching speculators/algos.py to forward fc_norm and norm_output config fields, and renaming fc_norms → fc_norm in checkpoint weights (naming mismatch between speculators and vLLM)

[rahul-tuli]

LGTM!