feat(train): default to Eagle 3.1 settings (llama arch + norms)

[orestis-z]

Summary

PR #610 benchmarks show that Eagle 3.1 with --draft-arch llama --norm-before-fc --norm-output gives +7.2% average acceptance length over Eagle 3 baselines (Qwen3-8B). This PR flips the CLI defaults so new training runs use these settings automatically.

Changes:

• --draft-arch: default "qwen3" → "llama"
• --norm-before-fc: default False → True, switched to BooleanOptionalAction (adds --no-norm-before-fc)
• --norm-output: default False → True, switched to BooleanOptionalAction (adds --no-norm-output)

Backward compatibility: Config class defaults (Eagle3SpeculatorConfig) remain False — old checkpoints deserialize correctly. Users can restore Eagle 3 + qwen3 behavior with --draft-arch qwen3 --no-norm-before-fc --no-norm-output.

Test plan

• ruff check passes on changed files
• pytest tests/unit/train/test_draft_config_init.py — all 36 tests pass
• pytest tests/unit/test_config.py -k norm — all 4 config roundtrip tests pass (config defaults unchanged)

🤖 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: 7536130a-a23f-4ada-ab54-93187784915c

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

This PR changes CLI defaults in scripts/train.py: --draft-arch now defaults to "llama" instead of "qwen3", and --norm-before-fc/--norm-output switch from store_true (default False) to BooleanOptionalAction (default True). Documentation and a test helper are updated accordingly.

Changes

CLI defaults and documentation update

Layer / File(s)	Summary
Draft architecture default change scripts/train.py, tests/unit/train/test_draft_config_init.py	--draft-arch default changes from "qwen3" to "llama"; test helper's create_transformer_layer_config call updated to pass draft_arch="llama".
Normalization flag semantics and docs scripts/train.py, docs/cli/train.md	--norm-before-fc and --norm-output change from store_true (default False) to BooleanOptionalAction (default True) with updated help text; documentation updated to describe Eagle 3.1 default behavior and --no-* variants for original Eagle 3 behavior.

Possibly related PRs

• vllm-project/speculators#610: Both PRs update the same Eagle 3.1 normalization CLI flags (--norm-before-fc / --norm-output) in scripts/train.py and corresponding docs/cli/train.md changes.
• vllm-project/speculators#620: Both PRs modify the --draft-arch default in scripts/train.py, impacting the same CLI behavior.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: switching training defaults to Eagle 3.1 settings.
Description check	✅ Passed	The description matches the changeset and accurately explains the CLI default flips and compatibility notes.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch eagle31-defaults

---

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.}

[fynnsu]

Looks good. Do we already have support for Eagle3.1 in vllm? Specifically loading from the speculators config?

[orestis-z]

Looks good. Do we already have support for Eagle3.1 in vllm? Specifically loading from the speculators config?

Yes, see https://github.com/vllm-project/vllm/blob/main/vllm/transformers_utils/configs/speculators/algos.py#L38-L40

[coderabbitai]

🧹 Nitpick comments (1)

tests/unit/train/test_draft_config_init.py (1)

411-420: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Consider adding a parse_args() default-value test.

This file exercises create_transformer_layer_config directly but nothing here (or shown elsewhere) asserts that parse_args() actually yields draft_arch="llama", norm_before_fc=True, and norm_output=True by default, or that --no-norm-before-fc/--no-norm-output correctly flip them. A small test would guard against future regressions in these argparse defaults.

As per path instructions, "Check that new code paths introduced in the PR are covered."

🤖 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/unit/train/test_draft_config_init.py` around lines 411 - 420, Add a
parse_args default-value test to cover the new argparse code path: verify the
parser in the training draft config flow returns draft_arch="llama",
norm_before_fc=True, and norm_output=True by default, and that the
--no-norm-before-fc and --no-norm-output flags correctly flip those values.
Place the test alongside the existing create_transformer_layer_config coverage
in test_draft_config_init.py so it directly exercises parse_args and guards the
CLI defaults from regressions.

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.

Nitpick comments:
In `@tests/unit/train/test_draft_config_init.py`:
- Around line 411-420: Add a parse_args default-value test to cover the new
argparse code path: verify the parser in the training draft config flow returns
draft_arch="llama", norm_before_fc=True, and norm_output=True by default, and
that the --no-norm-before-fc and --no-norm-output flags correctly flip those
values. Place the test alongside the existing create_transformer_layer_config
coverage in test_draft_config_init.py so it directly exercises parse_args and
guards the CLI defaults from regressions.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: b89bd0a9-adba-4b2b-acaf-045d77afed1e

📥 Commits

Reviewing files that changed from the base of the PR and between d4fd1d6 and 2823373.

📒 Files selected for processing (3)

• docs/cli/train.md
• scripts/train.py
• tests/unit/train/test_draft_config_init.py

[orestis-z]

Defaults affect all speculator types, not just Eagle3

--draft-arch, --norm-before-fc, and --norm-output are global CLI args shared by all speculator types. Changing their defaults here silently changes behavior for DFlash/DSpark users who relied on the old defaults (e.g. --draft-arch qwen3).

Suggested fix: keep the argparse defaults as None and resolve per-speculator-type after parsing:

# After parsing, fill in speculator-specific defaults
if args.draft_arch is None:
    args.draft_arch = "llama" if args.speculator_type == "eagle3" else "qwen3"
if args.norm_before_fc is None:
    args.norm_before_fc = args.speculator_type == "eagle3"
if args.norm_output is None:
    args.norm_output = args.speculator_type == "eagle3"

This way Eagle 3.1 gets the new defaults, DFlash/DSpark keep the old ones, and explicit flags still override everything.

---

@fynnsu what do you think? The global effect of this slipped through, my bad.

I remember we previously discussed not complicating the defaults.

I see 3 options:

• we use defaults per speculator type (my recommendation)
• we close this PR and leave the defaults as they are
• we extrapolate that llama + norms is better for all spec types not just eagle3 (I believe this is risky and/or not scientific)

[shanjiaz]

@orestis-z I think Fynn and I agree with you we should go with option 1, defaults for per speculator type.

[orestis-z]

Pushed the per-speculator-type defaults fix in 09bdc78:

• Argparse defaults are now None for --draft-arch, --norm-before-fc, --norm-output
• After parsing, defaults resolve based on --speculator-type:
  • eagle3 → llama arch, norms enabled
  • everything else → qwen3 arch, norms disabled
• Explicit flags still override in all cases
• Added tests for dflash defaults alongside the existing eagle3 ones