Add reverse KL divergence (RKL) loss for speculative decoding

[guan404ming]

Purpose

Draft training only had forward KL, which is mass-covering; reverse KL KL(draft || target) is mode-seeking and can favour acceptance. Add it as a selectable loss.

Tests

• Two-node-independent unit tests in test_metrics.py: identical dists -> ~0 and random -> non-negative, equals forward KL with args swapped, and resolve_loss_fn("rkl") wiring. 17 passed.

Checklist

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: 47e29211-f023-4faa-8c25-38cf82fedaab

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

Adds a reverse KL loss variant, wires rkl into loss-function resolution and the training CLI, and adds unit tests for the new loss and resolver mapping.

Changes

Reverse KL loss support

Layer / File(s)	Summary
Loss implementation and tests src/speculators/models/metrics.py, tests/unit/models/test_metrics.py	Adds reverse_kl_div_loss and unit tests for identity, swapped-input equivalence, and resolve_loss_fn("rkl") lookup.
Resolver and CLI wiring src/speculators/models/metrics.py, scripts/train.py	Adds "rkl" to resolve_loss_fn, updates its accepted-name docstring, and extends the --loss-fn choices and help text.

Possibly related PRs

• vllm-project/speculators#542: Extends the same --loss-fn and resolve_loss_fn plumbing pattern with another named loss option.
• vllm-project/speculators#648: Updates resolve_loss_fn and tests/unit/models/test_metrics.py with another added loss variant.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly describes the added reverse KL loss for speculative decoding.
Description check	✅ Passed	The description is directly related to the PR and summarizes the reverse KL loss and tests.
Docstring Coverage	✅ Passed	Docstring coverage is 90.00% which is sufficient. The required threshold is 80.00%.
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

---

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

[coderabbitai]

🧹 Nitpick comments (1)

tests/unit/models/test_metrics.py (1)

118-120: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add a parser-level test for --loss-fn rkl.

These assertions stop at resolve_loss_fn("rkl"), so a typo or regression in scripts/train.py's choices/CLI wiring would still pass. Please add one test that parses --loss-fn rkl end to end.

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/models/test_metrics.py` around lines 118 - 120, Add an end-to-end
parser-level test for the new `--loss-fn rkl` path so CLI wiring in
`scripts/train.py` is exercised, not just `resolve_loss_fn`. Extend the existing
metrics test coverage by parsing the training arguments with `--loss-fn rkl` and
asserting the parsed loss function resolves to `reverse_kl_div_loss`, using the
relevant parser/CLI entrypoint from `scripts/train.py` and the `resolve_loss_fn`
symbol to locate the flow. Keep the existing unit test for
`resolve_loss_fn("rkl")`, but add the new parser assertion to catch regressions
in `choices` or argument parsing.

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/models/test_metrics.py`:
- Around line 118-120: Add an end-to-end parser-level test for the new
`--loss-fn rkl` path so CLI wiring in `scripts/train.py` is exercised, not just
`resolve_loss_fn`. Extend the existing metrics test coverage by parsing the
training arguments with `--loss-fn rkl` and asserting the parsed loss function
resolves to `reverse_kl_div_loss`, using the relevant parser/CLI entrypoint from
`scripts/train.py` and the `resolve_loss_fn` symbol to locate the flow. Keep the
existing unit test for `resolve_loss_fn("rkl")`, but add the new parser
assertion to catch regressions in `choices` or argument parsing.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 8593e77b-371d-48a4-ada9-c91ccf623826

📥 Commits

Reviewing files that changed from the base of the PR and between 669b878 and 3605f23.

📒 Files selected for processing (3)

• scripts/train.py
• src/speculators/models/metrics.py
• tests/unit/models/test_metrics.py

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @guan404ming.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @guan404ming.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork