[4/7] fix(speculative): weigh DFlash2's selector loss like the backbone loss by h-guo18 · Pull Request #3143 · NVIDIA/Model-Optimizer

h-guo18 · 2026-10-30T13:07:41Z

Stack (merge in this order; each PR's base is the branch of the one before it). [7/7] reproduces nvidia/Muse-Spark-Open-816B-A42B-DFlash2 and needs all of them.

[1/4] fix(speculative): four training fixes behind the Muse-Spark-Open DFlash2 recipe #2617: draft RoPE frequencies stay in fp32
[2/7] fix(speculative): make streaming hidden-state capture work on vLLM 0.29 #2522: streaming hidden-state capture on vLLM >= 0.29
[3/7] fix(speculative): keep the streaming corpus memory-mapped instead of copying it per rank #3142: streaming corpus stays memory-mapped
[4/7] fix(speculative): weigh DFlash2's selector loss like the backbone loss #3143: DFlash2 selector loss weighted like the backbone loss <- this PR
[2/4] perf(speculative): deliver training batches out of order; compile + FlexAttention in the DFlash2 recipe #3001: out-of-order training batches; compile + FlexAttention in the DFlash2 recipe
[3/4] perf(speculative): fuse the DFlash-family loss with the lm_head projection #3010: fused DFlash-family loss
[4/4] feat(speculative): Muse-Spark-Open DFlash2 recipes that reproduce the released drafter #3144: Muse-Spark-Open DFlash2 recipes

What does this PR do?

Type of change: Bug fix

DFlash2's selector cross-entropy now weighs block positions the way the backbone loss does, and shares the backbone's denominator:

Weights. Each covered position's selector cross-entropy is multiplied by the backbone's per-position loss weight (decay or D-PACE), read from _DFlashLossTerms.weights. Before, every covered position counted equally.
Denominator. The sum is divided by the backbone loss's weight sum, _DFlashLossTerms.weight_sum, which also counts the positions whose gold token is outside the top-k. Before, it was divided by the covered positions alone, so the term grew whenever the backbone's top-k recall dropped.
The selector now reads the backbone's supervision mask from _DFlashLossTerms.supervised_mask instead of rebuilding it, so the two cannot drift apart.

Why. A block position's choice only matters if every earlier position was accepted, and the backbone's weighting models exactly that. With this change the selector spends its capacity where acceptance is decided. This is how the SGLang/SpecForge DFlash2 reference trains its selector, and it is the objective that the released nvidia/Muse-Spark-Open-816B-A42B-DFlash2 drafter was trained with.

What I have not measured. I have no A/B of this change on its own. main already trains the selector on the strict top-k (no gold-token substitution) and zero-initializes its successor codebook. Together with this change, that is the objective that beat the old gold-substituting one by +7.8% (greedy) and +10.1% (temperature 1) MT-Bench acceptance length, in from-scratch Muse-Spark-Open runs of 550K samples each. How that gain splits between the strict top-k and the weighting here, I did not measure.

The module docstring still described the gold-token substitution that the strict top-k replaced. It now describes the current objective.

No CHANGELOG entry: DFlash2 first ships in 0.48.0, which is still at a release candidate.

Usage

No API or config change. The selector term's value changes, so a run's logged loss will not match a run on the previous objective.

Testing

New unit test TestCandidateSelectorAlignment::test_positions_weigh_like_the_backbone_and_misses_stay_in_the_denominator in tests/unit/torch/speculative/plugins/test_hf_dflash2.py. It checks that the term is the backbone-weighted mean over every supervised position, a miss included. With the previous objective restored, it fails and the other 33 tests in the file pass.
pytest of the DFlash-family unit tests (test_hf_dflash.py, test_hf_dflash2.py, test_fakebase.py, test_dflash_fused_head.py, test_hf_lilicorr.py, test_modeling_final_norm.py) on the top of the stack: passed with transformers 5.5.4 and 5.18.0 (CPU).
End to end, see [7/7], which trains with this objective.

Before your PR is "Ready for review"

Is this change backward compatible?: ✅ (no API change; the selector term's value changes)
If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
Did you write any new necessary tests?: ✅
Did you update Changelog?: N/A (DFlash2 has not shipped yet)
Did you get Claude approval on this PR?: ❌

Additional Information

🤖 Generated with Claude Code

The selector cross-entropy now uses the backbone loss's own per-position weights (decay or D-PACE) and divides by the backbone's weight sum, misses included. Before, every covered position counted equally and the term was divided by the covered positions alone. A block position's choice only matters if every earlier position was accepted, which is what the backbone's weighting models, so the selector now spends its capacity the same way. Dividing by the covered positions alone also scaled the term up whenever the backbone's top-k recall dropped. This matches the SGLang/SpecForge DFlash2 reference, and it is the objective the released Muse-Spark-Open DFlash2 drafter was trained with. The selector reads the backbone's mask and weights from _DFlashLossTerms instead of rebuilding the mask, so the two cannot drift apart. The module docstring still described the gold-token substitution that the strict top-k replaced; it now describes the current objective. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

copy-pr-bot · 2026-10-30T13:07:46Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

coderabbitai · 2026-10-30T13:07:50Z

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

h-guo18 · 2026-10-31T08:43:38Z

Folded into #2617 together with #2617's own fix and the other two small fixes, rebased onto main after #2757. It is now [1/4] of the Muse-Spark-Open DFlash2 recipe stack. The commits are the same, re-signed, with the selector one ported onto #2757's helpers; see #2617 for the combined description and testing.

…n DFlash2 recipe (#2617)  **Stack** (merge in this order; each PR's base is the branch of the one before it). [4/4] reproduces [nvidia/Muse-Spark-Open-816B-A42B-DFlash2](https://huggingface.co/nvidia/Muse-Spark-Open-816B-A42B-DFlash2) and needs all of them. 1. #2617: four training fixes: draft RoPE in fp32, streaming capture on vLLM >= 0.29, a memory-mapped streaming corpus, DFlash2's selector loss weighted like the backbone's **<- this PR** 2. #3001: out-of-order training batches; compile + FlexAttention in the DFlash2 recipe 3. #3010: fused DFlash-family loss 4. #3144: Muse-Spark-Open DFlash2 recipes ---  ### What does this PR do? Type of change: Bug fix Four independent fixes to speculative-decoding training, all needed to reproduce the Muse-Spark-Open DFlash2 drafter. They used to be four PRs (#2617, #2522, #3142, #3143). The other three are closed and point here. Each fix is its own commit, and the streaming connector takes two. Review follow-ups come after them. All four touch only `modelopt/torch/speculative`. The whole PR is rebased onto `main` after #2757, the DFlash-family refactor. #### 1. The draft's RoPE frequencies stay in fp32 on a fresh convert On a freshly converted DFlash-family model, `restore_draft_precision()` repeated the `.to(device, dtype)` that `modify()` had already applied. `Module.to()` casts floating buffers too, so this rounded the rotary `inv_freq` that `modify()` had just built in fp32 down to bf16. `examples/speculative_decoding/main.py` calls `restore_draft_precision()` unconditionally, trusting its docstring's "no-op in that case". - **Why it matters.** The rounding hits the frequencies, not the cos/sin, so the phase error grows with position: about 0.4% per frequency, and several radians for the highest-frequency pairs by position 4096. vLLM computes the angles from fp32 frequencies, so the draft trained under different RoPE phases than it was served with. - **Who is affected.** Every DFlash, DSpark, Domino, DFlash2 or LiLiCorr draft trained from scratch through `main.py` since #2483. - **The fix.** `modify()` records whether it placed the draft, and `restore_draft_precision()` returns early when it did. On the `from_pretrained` path, where the base is on the meta device, it still places the draft. - **Measured impact.** An 80-step Gemma-4-E4B DSpark A/B, fp32 against bf16 frequencies, stayed within run-to-run noise. The effect on acceptance length was not measured. #### 2. Streaming hidden-state capture on vLLM >= 0.29 Streaming training, where vLLM serves produce hidden states and the trainer pulls them over NIXL RDMA, did not work on current vLLM. Two of the defects produce wrong data without an error. | Commit | Defect | Symptom | | --- | --- | --- | | 1 | vLLM #51718 changed the KV-cache layout to `[B, H, N, C]` | **Silent.** The trainer sees one plane and hands the draft an empty aux tensor. `planes_dim()` now finds the plane axis from the configured capture count. | | 1 | Captures were addressed by their position in the batch | **Silent.** A capture read another request's hidden states. With concurrency, the hidden/token match was 0.00-0.02 instead of 0.96-1.00. Each request's slots now come from its own block ids. | | 1 | Capture index tensors were built on the wrong stream | A non-deterministic out-of-bounds assert, minutes in. | | 2 | The capture layer is merged into the attention group on uniform-attention models | Loud: serve does not start on Qwen3, Llama and most dense models. The lookup now recognises the merged group in both the worker's and the scheduler's view. | Because two of these are silent, the new checks are loud: - The resolved group and cache view are logged at startup. - Requests whose block ids cannot cover their tokens raise, naming the request. - The first few requests verify every derived slot. The module's logger now sits under `vllm.*`. vLLM configures only its own logger tree, so until now every `logger.info` in this module was dropped, the startup diagnostics included. The connector needs vLLM >= 0.22, for `HiddenStateCacheSpec`, and is tested on vLLM >= 0.29. It reads the KV-cache layout from both before and after vLLM #51718. The `prefer_cross_layer_blocks` override is removed: #51718 deleted the base property, and before that its default was already `False`, the value the override returned. #### 3. The streaming corpus stays memory-mapped `StreamingDataset.__init__` called `list()` on an Arrow-backed `datasets.Dataset`. That decoded the whole corpus into Python objects in every rank, and each forked DataLoader worker then copy-on-wrote the pages it touched. On a 3.1M-row corpus this measured 80 GB per rank, and with 4 ranks × 8 workers it OOM-killed a 950 GB trainer node. An indexable source is now kept as is, so its rows stay memory-mapped and shared across workers. Only an iterator is still materialized. One changelog entry. #### 4. DFlash2's selector loss is weighted like the backbone loss The selector cross-entropy now matches the backbone loss in three ways: - **Weights.** It multiplies each covered position by the backbone's per-position weight (decay or D-PACE). - **Denominator.** It divides by the backbone's weight sum, which counts the positions whose gold token is outside the top-k. Before, every covered position counted equally, and the term was divided by the covered positions alone, so it grew whenever backbone recall dropped. - **Mask.** It reads the backbone's mask, weights and weight sum from `_DFlashLossTerms` instead of rebuilding them, so the two losses cannot drift apart. A block position's choice only matters if every earlier one was accepted, which is what the backbone's weighting models. The SGLang/SpecForge DFlash2 reference trains its selector this way, and the released drafter was trained with this objective. On top of #2757, the selector takes its gold tokens from `_block_tokens` at the variant's label shift. One CHANGELOG entry, because the 0.48.0 release branch already ships DFlash2 with the old objective. ### Usage No API or config change. The selector term's value changes, so a DFlash2 run's logged loss will not match one on the previous objective. ### Testing New tests, each failing on the code it fixes: - `TestDFlashLazyRotaryEmb` gains two tests: a fresh convert keeps `inv_freq` in fp32, and the meta-device `from_pretrained` path still places the draft. - `test_an_indexable_corpus_is_kept_and_an_iterator_materialized` covers the streaming dataset. - `TestCandidateSelectorAlignment::test_positions_weigh_like_the_backbone_and_misses_stay_in_the_denominator` covers the selector. - `tests/gpu_vllm/torch/speculative/plugins/test_rdma_hidden_states_connector.py` covers the connector. One test runs `build_connector_meta` and `save_kv_layer` on CUDA for two prompts prefilled behind a decoding request, whose capture blocks are neither contiguous nor in batch order, and checks that each pool slot holds its own request's hidden states, on both KV-cache layouts. The others check that an ambiguous cache view is refused and that the capture group is found in each of the three shapes vLLM gives it. The file sits in `gpu_vllm` because the connector imports vLLM at module level. Reverting `save_kv_layer` to reading `slot_mapping` by position, or `build_connector_meta` to the attention group's block ids, fails the first test on both layouts. On this branch, rebased onto `main` (e352cac) after #2757, with Python 3.12, torch 2.10 and transformers 5.18.0 on CPU: - `test_hf_{dflash,dflash2,lilicorr,streaming_dataset}.py` and `tests/unit/torch/export/test_hf_spec_rope_export.py`: 244 passed. - The bitwise A/B of the whole stack before and after the rebase is in [3/4]. After the review follow-ups, in a venv with vLLM 0.30.0 (the gpu_vllm CI image's version), torch 2.13 and transformers 5.18.0, on one RTX PRO 6000: the connector tests and `test_hf_streaming_dataset.py`, 27 passed. End to end: - Qwen3-8B DFlash2 streaming on GB200 (1 serve node, TP4, plus 1 trainer node, 4-rank DDP): 600 steps, loss 17.0 -> 3.4. Before the connector fixes, this run failed at three different points. - Muse-Spark-Open DFlash2 streaming on GB300, on the full 3.48M-row corpus, is in [4/4]. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ❌ The streaming connector now needs vLLM >= 0.22 (tested on >= 0.29), listed under Backward Breaking Changes. No API or config change. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ A Bug Fixes entry for each of the four fixes, and a Backward Breaking Changes entry for the connector's vLLM floor. The 0.48.0 release branch already carries the RoPE regression (#2483) and DFlash2, so those two fixes are to released code as well. - Did you get Claude approval on this PR?: ❌ The Claude review commented rather than approved. Its findings are addressed in the follow-up commits. ### Additional Information The four fixes were found bringing up DFlash2 streaming and the Muse-Spark-Open release. 🤖 Generated with [Claude Code](https://claude.com/claude-code)   ## Summary by CodeRabbit * **Bug Fixes** * Draft restoration preserves rotary embedding precision and avoids unnecessary device placement, including for models initially loaded on a meta device. * Selector loss now respects per-position weights and includes missed candidates in its denominator. * Hidden-state transfer handles supported cache layouts and request block mappings more reliably. * Fixed memory-mapped training-corpus access and concurrent vLLM hidden-state capture. * **Performance** * Indexable streaming datasets are retained without copying; other iterable inputs are supported. * **Compatibility** * The streaming speculative-decoding connector requires vLLM 0.22 or later.  --------- Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

h-guo18 closed this Oct 31, 2026

Navigation Menu

Repository navigation

[4/7] fix(speculative): weigh DFlash2's selector loss like the backbone loss - #3143

[4/7] fix(speculative): weigh DFlash2's selector loss like the backbone loss#3143
h-guo18 wants to merge 1 commit into
haoguo/streaming-corpus-mmapfrom
haoguo/dflash2-selector-loss-weights

h-guo18 commented Oct 30, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Oct 30, 2026

Uh oh!

coderabbitai Bot commented Oct 30, 2026

Draft PR not reviewed

Uh oh!

h-guo18 commented Oct 31, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

h-guo18 commented Oct 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Usage

Testing

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Oct 30, 2026

Uh oh!

coderabbitai Bot commented Oct 30, 2026

Draft PR not reviewed

Uh oh!

h-guo18 commented Oct 31, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

h-guo18 commented Oct 30, 2026 •

edited

Loading