feat(training): support configurable loss fn via training args

[christinaexyou]

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

Purpose

This PR addresses #541

Description

It adds a --loss-fn argument and input validation to scripts/train.py. If defined, it passes the resolved loss function to get_trainer_kwargs() so each model's compute_metrics call uses arg.loss_fn. If args.loss_fn=None, it falls back to the previously hardcoded loss fn.

Tests

I added unit tests here tests/unit/train/test_cli_args.py to test that arg.loss_fn is passed correctly to each model's get_trainer_kwargs() and that we maintain backward compatibility

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: afabf081-c732-4d5a-925a-fb5b548634a3

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 adds configurable loss function selection across draft model training via a new --loss-fn CLI argument. A resolve_loss_fn utility maps string names (kl, ce) to loss callables. The CLI parses and validates the argument, then each model family (DFlash, Eagle3, PEagle) threads the resolved loss function through training/validation kwargs and metric computations.

Changes

Loss function configurability

Layer / File(s)	Summary
Loss resolution foundation src/speculators/models/metrics.py	Introduces resolve_loss_fn(name: str) that maps "kl" and "ce" to kl_div_loss and ce_loss respectively, with validation for unknown names.
CLI argument and validation scripts/train.py	Adds --loss-fn CLI option (choices kl, ce), imports resolve_loss_fn, and validates the argument after parsing.
DFlash loss function threading src/speculators/models/dflash/core.py, src/speculators/models/dflash/metrics.py	DFlash model imports loss utilities, adds loss_fn parameter (default ce_loss) to forward(), updates get_trainer_kwargs() to resolve and return loss_fn in both train and validation kwargs, and passes loss_fn to compute_metrics().
Eagle3 loss function threading src/speculators/models/eagle3/core.py, src/speculators/models/eagle3/metrics.py	Eagle3 model imports loss utilities, adds loss_fn parameter (default kl_div_loss) to forward(), updates get_trainer_kwargs() to resolve and return loss_fn in both train and validation kwargs, and passes loss_fn to per-step compute_metrics().
PEagle loss function threading src/speculators/models/peagle/core.py, src/speculators/models/peagle/metrics.py	PEagle model imports loss utilities, adds loss_fn parameter (default kl_div_loss) to forward(), updates get_trainer_kwargs() to resolve and return loss_fn in both train and validation kwargs, and passes loss_fn to compute_metrics().
CLI argument validation tests tests/unit/train/test_cli_args.py	Unit tests validate --loss-fn propagation through each model's get_trainer_kwargs(), confirming model-specific defaults (ce for DFlash, kl for Eagle3/PEagle) and correct override behavior.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Suggested labels

enhancement, training, eagle3, dflash, two-reviews

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 48.00% 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 'feat(training): support configurable loss fn via training args' accurately summarizes the main change—adding configurable loss function support via CLI arguments for training.
Description check	✅ Passed	The description clearly explains the purpose (addressing #541), details the changes (--loss-fn argument, validation, passing to get_trainer_kwargs), mentions tests added, and notes backward compatibility maintenance.
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

---

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

[mergify]

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 Require two reviews

Wonderful, this rule succeeded.

PRs labelled "two-reviews" must have at least two approving reviews before merging.

• #approved-reviews-by >= 2

[christinaexyou]

I noticed that there were no tests to verify CLI args so I created a new file but lmk if you want me to rename/move the tests :)

[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

[shanjiaz]

Looks great! Pending quality check.

[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

[fynnsu]

Thanks! Should be almost good to merge

[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

[shanjiaz]

Great job! Thanks for contributing.

[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

[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

[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