Add shared resolve_target_layer_ids util for eagle3 and dflash

[fynnsu]

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

Purpose

Closes #454

Combines logic for target layer id resolution in both eagle3 and dflash.

Description

Creates a shared util fn resolve_target_layer_ids which both algorithms call in their from_training_args function.

Related Issue

#454

Tests

I locally ran tests/e2e/smoke/test_online_training.py::test_online_smoke which succeeded.

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: 31983cc9-765b-4fc5-bffd-c2c89d0ed135

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 change standardizes how target_layer_ids is resolved across Eagle3 and DFlash models by extracting shared resolution logic into a utility function and moving resolution from model __init__ to from_training_args for consistency. A new utility file provides verifier config loading and default layer ID computation.

Changes

Cohort / File(s)	Summary
Utility extraction src/speculators/models/utils.py	Added get_verifier_config() to load verifier configurations via AutoConfig, resolve_target_layer_ids() to centralize default computation and warning emission, and DEFAULT_TARGET_LAYER_IDS_WARNING constant for standardized messaging.
DFlash refactoring src/speculators/models/dflash/core.py	Moved target_layer_ids resolution from __init__ to from_training_args, removed AutoConfig usage from constructor, and now requires config.aux_hidden_state_layer_ids to be pre-set, raising ValueError if unset.
Eagle3 refactoring src/speculators/models/eagle3/core.py	Replaced inline target layer resolution logic in from_training_args with call to centralized resolve_target_layer_ids() utility function, consolidating default computation and warning handling.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• Add explicit target-layer-ids handling #379: Related changes to how target layer IDs are determined and configured in Eagle3; this PR refactors and centralizes the logic introduced in that PR.

Suggested labels

refactor, eagle3, dflash, two-reviews

Suggested reviewers

• shanjiaz
• dsikka

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: adding a shared utility function for resolving target layer IDs that both Eagle3 and DFlash models will use.
Description check	✅ Passed	The description is related to the changeset, explaining the purpose of combining target layer ID resolution logic across Eagle3 and DFlash, with references to the closed issue and test results provided.
Linked Issues check	✅ Passed	The PR fulfills the objectives outlined in issue #454: it extracts shared utility logic into resolve_target_layer_ids, moves DFlash's resolution into from_training_args, and refactors Eagle3 to use the centralized utility, matching both requirements.
Out of Scope Changes check	✅ Passed	All changes are directly scoped to the objectives of issue #454: extracting shared utility functions and refactoring both Eagle3 and DFlash to use them. No unrelated modifications are present.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

_{✏️ 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 resolve_target_layer

---

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: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@src/speculators/models/dflash/core.py`:
- Around line 64-69: The constructor DFlashDraftModel.__init__ currently
requires config.aux_hidden_state_layer_ids to be non-None (raising if None), but
the config schema still defaults that field to None which breaks
from_pretrained() loads; either make aux_hidden_state_layer_ids required at the
config schema level in DFlashSpeculatorConfig (remove the None default) or,
preferably, add a fallback in DFlashDraftModel.__init__ to call the existing
resolve_target_layer_ids() (the same logic used by from_training_args()) when
config.aux_hidden_state_layer_ids is None so models loaded via from_pretrained()
get resolved target_layer_ids instead of raising. Ensure you reference and
update DFlashDraftModel.__init__, DFlashSpeculatorConfig,
resolve_target_layer_ids(), and from_training_args()/from_pretrained()
accordingly.

🪄 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: 24b5c8f5-0952-4f0a-a14b-209d5574fa93

📥 Commits

Reviewing files that changed from the base of the PR and between 8fdee2d and aa6c8fa.

📒 Files selected for processing (3)

• src/speculators/models/dflash/core.py
• src/speculators/models/eagle3/core.py
• src/speculators/models/utils.py

[shanjiaz]

lol thanks for refactoring this.