Add peagle test

[shanjiaz]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

Added Peagle smoke/regression training

Description

Added training only regression test for p-eagle and offline training smoke test.

Related Issue

Tests

• There's a regression on vLLM main right now. This test would fail with low acceptance rate until Add parallel drafting to v2 model runner unsupported features vllm#43010 is merged.
• Smoke test would only pass for nightly since previous vLLM releases don't support short-form serving command for p-eagle speculators models.
• Tested locally with the vLLM fix branch and it passed.

I have filled in:

• 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.
• (Optional) 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: 1472fb01-ce30-4a1b-ac57-837c12b1ba24

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

This PR introduces PEagle as a new speculator type with corresponding test coverage, adjusts FSDP mixed-precision policy application at the model level, constrains model artifact resolution to specific file patterns, and updates offline training error handling from raise to skip on missing data.

Changes

Training infrastructure and PEagle speculator support

Layer / File(s)	Summary
FSDP mixed-precision policy configuration src/speculators/train/utils.py	Top-level model FSDP wrapping no longer passes explicit mixed-precision policy, changing from fully_shard(model, mp_policy=mp_policy) to fully_shard(model) while per-layer wrapping retains the policy argument.
Model artifact resolution filtering tests/e2e/regression/test_training_only_acceptance.py	_resolve_repo now applies file-pattern constraints (*.arrow, *.json, *.pt, hidden_states/*) to both cached and network-fallback snapshot_download calls, narrowing downloaded artifacts.
PEagle smoke test parametrization tests/e2e/smoke/test_offline_training.py	Extended offline training smoke test with PEagle speculator type case, including custom layer configuration parameters and target_layer_ids=None.
PEagle regression test and vocabulary configuration tests/e2e/regression/test_training_only_acceptance.py	New test_peagle_qwen3_8b_sharegpt regression test trains PEagle model with draft_vocab_size=32000 and acceptance threshold validation; existing Eagle3 and DFlash tests updated to use draft_vocab_size=32000 instead of 8192.
Offline training error handling behavior tests/e2e/utils.py	In run_training, missing/generated data handling changed from "raise" to "skip" when training offline, altering error recovery behavior in non-online scenarios.

Possibly related PRs

• vllm-project/speculators#389: Both PRs modify the same _resolve_repo helper in regression tests, adjusting how model artifact downloads are constrained.

Suggested labels

training, two-reviews

Suggested reviewers

• fynnsu
• rahul-tuli

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 28.57% 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 'Add peagle test' accurately describes the main purpose of the changeset, which is adding new peagle-related tests across multiple test files.
Description check	✅ Passed	The description is related to the changeset, explaining the addition of p-eagle training tests (regression and smoke tests) and relevant context about testing conditions.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch add-peagle-test

---

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 and usage tips.}

[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/train/utils.py`:
- Line 154: The root call to fully_shard is missing the mixed-precision policy
so the top-level parameters inherit default dtypes; update the call to
fully_shard at the end of the sharding sequence to pass the mp_policy argument
(i.e., replace the bare fully_shard(model) with a call that passes the existing
mp_policy variable to fully_shard) so the root module uses the same
MixedPrecisionPolicy as the child modules.

In `@tests/e2e/utils.py`:
- Around line 256-259: The test command currently appends "--on-missing", "skip"
to train_cmd which allows missing hidden states to be silently ignored; change
this to fail fast by replacing the "skip" value (or removing the flag) so
missing hidden states cause a test failure — locate the train_cmd construction
in tests/e2e/utils.py and update the "--on-missing" argument value to an
explicit failure mode (e.g., "error" or "fail") or remove the flag entirely to
ensure the test fails on missing data.

🪄 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: 8487a20f-9a39-4742-bac5-61642cce590a

📥 Commits

Reviewing files that changed from the base of the PR and between 4f80f5d and dee4e14.

📒 Files selected for processing (4)

• src/speculators/train/utils.py
• tests/e2e/regression/test_training_only_acceptance.py
• tests/e2e/smoke/test_offline_training.py
• tests/e2e/utils.py