feat(eagle3): add per-layer FC normalization (--fc-norm)

[orestis-z]

Summary

• Add --fc-norm flag for per-layer RMSNorm on each auxiliary hidden state before concatenation and FC projection: concat(Norm(h_a), Norm(h_b), Norm(h_c))
• This matches the Eagle 3.1 paper specification and vLLM's fc_norm config field
• Differs from --norm-before-fc which applies a single norm to the full concatenation: Norm(concat(h_a, h_b, h_c))
• norm_before_fc and fc_norm are validated as mutually exclusive in the config (enabling both would cause double-norming)

Test plan

pytest tests/unit/test_config.py -k "fc_norm" -v
pytest tests/integration/models/test_model_forward.py -k "fc_norm" -v

• test_fc_norm — Eagle3 with fc_norm + norm_output, verifies fc_norms module created and forward + backward pass
• test_peagle_fc_norm — P-Eagle with fc_norm
• test_eagle3_config_fc_norm_roundtrip — config serialization roundtrip
• test_eagle3_config_rejects_both_norm_before_fc_and_fc_norm — mutual exclusivity validation

🤖 Generated with Claude Code

[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: 489d39b7-011f-4277-83bf-54dda04a9460

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

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 29.41% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: adding per-layer FC normalization via --fc-norm.
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 by describing the new fc-norm flag, its behavior, and related 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

[coderabbitai]

Actionable comments posted: 4

🧹 Nitpick comments (1)

tests/integration/models/test_model_forward.py (1)

385-408: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win

Assert the checkpoint key contract for fc_norm.

The PR objective depends on vLLM-compatible parameter names; add a direct state_dict assertion so this test catches accidental key drift.

✅ Proposed test hardening

         assert model.fc_norm is not None
         assert len(model.fc_norm) == 3
+        assert all(f"fc_norm.{idx}.weight" in model.state_dict() for idx in range(3))
         assert model.input_norm is None
@@
         assert model.fc_norm is not None
         assert len(model.fc_norm) == 3
+        assert all(f"fc_norm.{idx}.weight" in model.state_dict() for idx in range(3))
         samples = _make_samples([128])

As per path instructions, tests should cover new code paths introduced in the PR.

🤖 Prompt for 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.

In `@tests/integration/models/test_model_forward.py` around lines 385 - 408, The
fc_norm tests only validate runtime behavior and miss the parameter-name
contract required for vLLM compatibility. Update test_fc_norm and/or
test_peagle_fc_norm to assert the model.state_dict contains the expected fc_norm
checkpoint key(s) for the fc_norm path, using make_eagle3_model,
make_peagle_model, and fc_norm to locate the affected code. Keep the existing
shape/loss checks, but add a direct state_dict key assertion so accidental key
drift is caught when new code paths are introduced.

Source: Path instructions

🤖 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 `@docs/cli/train.md`:
- Around line 139-141: The CLI docs for the train flags are missing the
llama-specific default for `--norm-output`, so the current description still
implies it is always `False`. Update the flag documentation near `--fc-norm` to
explicitly note that `scripts/train.py` applies the same `is_llama_eagle`
conditional defaulting to `--norm-output`, and make sure the wording matches the
current behavior for llama vs non-llama runs. Use the existing flag descriptions
in this section as the reference point and keep the defaults consistent with the
implementation.

In `@scripts/train.py`:
- Around line 1037-1049: Update the help text for the argparse options in
train.py so it matches the real defaulting logic: both --fc-norm and
--norm-output are only implicitly true when the draft architecture condition is
satisfied together with --speculator-type being eagle3 or peagle, not for llama
draft arch alone. Adjust the help strings near the parser.add_argument calls for
these flags to reflect the combined condition used elsewhere in the argument
parsing logic.

In `@src/speculators/convert/eagle/eagle3_converter.py`:
- Line 43: The fc_norm handling in Eagle3 conversion is treating an explicit
False the same as “not provided,” so the caller cannot override a checkpoint’s
truthy value. Update the config merge logic in the Eagle3 converter’s conversion
path (the code that builds the speculator config from eagle_config and the
fc_norm parameter) to distinguish None/default from an explicit boolean, so
fc_norm=False is written through instead of falling back to the checkpoint
value. Apply the same fix anywhere the converted config is assembled from legacy
checkpoint settings in this converter.

In `@src/speculators/models/eagle3/core.py`:
- Around line 104-115: Reject invalid Eagle 3 configs by validating the
normalization flags in the model setup: in Eagle3 core initialization, ensure
`config.fc_norm` cannot be enabled at the same time as
`norm_before_fc`/`input_norm` creation, since that would make `forward` apply
both normalization paths. Add a config check near the `self.fc_norm` /
`self.input_norm` setup in `src/speculators/models/eagle3/core.py` and raise a
clear error when both FC normalization modes are requested, so only one
documented Eagle 3.1 path is allowed.

---

Nitpick comments:
In `@tests/integration/models/test_model_forward.py`:
- Around line 385-408: The fc_norm tests only validate runtime behavior and miss
the parameter-name contract required for vLLM compatibility. Update test_fc_norm
and/or test_peagle_fc_norm to assert the model.state_dict contains the expected
fc_norm checkpoint key(s) for the fc_norm path, using make_eagle3_model,
make_peagle_model, and fc_norm to locate the affected code. Keep the existing
shape/loss checks, but add a direct state_dict key assertion so accidental key
drift is caught when new code paths are introduced.

🪄 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: f48a6149-82b5-47f7-b673-4178bb64cd9c

📥 Commits

Reviewing files that changed from the base of the PR and between aba50b0 and 86507a9.

📒 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

[shanjiaz]

Looks good just one small non-blocking concern

[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