[h-guo18]

What does this PR do?

Type of change: New feature. It changes the default weighting for DSpark and Domino configs that leave dflash_loss_objective unset, and for runs resumed from checkpoints saved before it.

This PR is stacked on #2990 [2/2] (branch haoguo/dflash-shared-forward, which gives the DFlash family one training forward), whose branch is its base; review and merge it after that one. The stack below is #2757 → #2762 → #2989 [1/2] → #2990 [2/2]. Three other PRs also sit on #2990 and are independent of this one: logging draft_metrics in the trainer (#2992), loss terms in config (#2993), and the shared AR-decode helper (#2994).

Since they were added in 0.46.0, DSpark and Domino have ignored dflash_loss_objective. They always weighted block positions with the static dflash_loss_decay_factor decay, while DFlash, DFlash2 and LiLiCorr follow the objective, whose default is dpace. On the shared forward this came down to one override per variant: _position_weighting returning "decay". Meanwhile DFlashModel.modify warned that the decay factor is ignored under D-PACE, which was false for these two.

• DSpark and Domino follow the objective.
  • The two overrides are removed. So is the _position_weighting hook, since nothing else overrode it; _loss_context now reads dflash_loss_objective directly.
  • Both variants are weighted from their first predicted block position, 0 (1 - _label_shift).
  • D-PACE reads the draft's confidence off the served (final) logits: DSpark's Markov-biased logits and Domino's GRU-corrected ones.
  • One set of weights per step reduces every term: DSpark's CE, TVD and confidence BCE, and both of Domino's objectives, final and backbone. Domino's backbone objective in the λ-curriculum is therefore weighted by the served draft's confidence, not its own.
  • The decay path is unchanged.
• New runs of the shipped recipes keep their tuned weighting. dspark.yaml and domino.yaml set dflash_loss_decay_factor: 7.0 without an objective, so they would have switched to D-PACE.
  • Both now pin dflash_loss_objective: decay. That commit lands first, so every commit trains them the same.
  • The Kimi-K2.6 and MiniMax-M3 DSpark launchers inherit the pin from dspark.yaml.
  • The Qwen3-8B Domino launcher repeats the recipe's knobs for visibility, so it now lists the pin too, as the LiLiCorr launcher does.
• Docs.
  • A Backward Breaking Changes entry in the CHANGELOG, under 0.49.0. It also tells users to finish runs started on an earlier version before upgrading.
  • The DSpark and Domino skill sheets list the pinned decay among the recipes' differences from DFlash.
  • Each sheet also gets a known-failures row for a resumed run, keyed on the warning the restore logs.

Behavior change:

• A DSpark or Domino config that leaves dflash_loss_objective unset now trains with D-PACE and ignores dflash_loss_decay_factor, as the existing warning says.
• A config that sets dpace now gets it.
• A run resumed from a checkpoint saved before this change continues with D-PACE, unless its config set decay.
  • The checkpoint's modelopt_state holds the default dpace, which the old code ignored.
  • main.py resumes an HF-format checkpoint, the format DDP writes and the shipped DSpark and Domino recipes use, by restoring that saved config. It never re-reads the recipe, so the new pin does not reach the run.
  • The checkpoint itself loads unchanged. The only sign is the restore's "dflash_loss_decay_factor=7.0 is ignored" warning, which these recipes already printed, wrongly, before this change.
  • Finish such runs on the version that started them.
• As for DFlash under D-PACE, a step whose only supervised label has weight below 1 takes the zero-loss path. Under decay, a lone label at block position 0 had weight exactly 1 and trained.
• Generation and export are unaffected.

Usage

dflash:
  dflash_loss_objective: dpace   # `decay` keeps the static schedule, as the shipped recipes do
  dflash_dpace_alpha: 0.5
  dflash_architecture_config:
    projector_type: dspark       # or domino

On top of the shipped recipes: dflash.dflash_loss_objective=dpace.

Testing

• New test. test_next_token_variants_weight_positions_by_the_objective[dspark|domino-dpace|decay] runs the real training forward and checks, bitwise:

  • the position weights handed to the loss equal _block_position_weights(mask, objective, 0, ...) computed from the final logits' cross-entropy;
  • DSpark's loss_ce and Domino's loss_backbone equal their cross-entropy reduced under those weights.

  It fails on the base for both D-PACE cases, which got decay weights. In scratch runs it also fails when the confidence is read off the backbone, when the first weighted position is 1, and when Domino's backbone objective gets D-PACE weights of its own. Each case takes at most 0.13 s on 4 CPUs (taskset -c 0-3).

• Unit tests. The test files this PR changes or exercises, on CPU: test_hf_dflash.py, test_hf_dflash2.py, test_hf_dspark.py, test_hf_domino.py, test_hf_lilicorr.py and test_hf_dflash_offline.py under tests/unit/torch/speculative/plugins, and tests/unit/recipe/test_loader.py. 806 passed on each of transformers 5.5.4 and 5.18.0, CI's tf_min and tf_latest pins (802 on the base; the other 4 are the new test's cases).

• GPU tests. tests/gpu/torch/speculative/plugins on 1 GPU: 20 passed, 1 skipped (that test needs 2 GPUs), as on the base.

• Bitwise A/B against the base. One training step per case, plus generation at steps 0, 2 and 4 for the online cases, in fp32 and bf16. Compared: loss, every gradient, train_acc, every draft_metrics value and the generated tokens.

  • The 23 existing cases (46 runs) are identical. They cover DFlash (KD and CE), DFlash2 (ce, kd, tv and lambda), LiLiCorr, DSpark (vanilla and rnn heads, with the confidence head) and Domino, offline and online, with D-PACE and decay, including zero-loss steps. All 15 cases with an explicit decay objective are among them.
  • The 5 new DSpark/Domino D-PACE training cases differ, as intended. For example, the DSpark offline loss goes from 0.674028516 to 0.671546638 in fp32, and the Domino online loss from 3.45916319 to 3.43934727. Every gradient changes (28 of 29 in one bf16 run). DSpark's loss_ce, loss_tvd and confidence_loss change, as do Domino's loss_final and loss_backbone. The generated tokens are identical.
  • Each of those 5 cases also equals, bitwise, the base with the two overrides deleted at runtime.
  • The 2 new zero-loss D-PACE cases are identical: both trees take the zero path.
  • Further cases, as expected. 15 are identical: DSpark's gated head, its rnn head online, DSpark and Domino under decay with the KD flag set (which they ignore), zero-loss steps of DFlash2 and LiLiCorr, DFlash2 KD online, and six cases whose labels are -100 on some tokens. 6 more DSpark/Domino D-PACE cases (the gated head, the KD flag set, labels with -100) differ as intended and equal, bitwise, the base with the overrides deleted.

• Resume across the change, on transformers 5.5.4 and 5.18.0. Tiny DSpark and Domino models were converted from each tree's shipped recipe and saved with save_pretrained. Each was then loaded the way main.py resumes an HF-format checkpoint (load_vlm_or_llm under patch_transformers5_params_loading, then restore_draft_precision) and run for one training step.

  • Saved on the base, both store dflash_loss_objective: dpace with decay factor 7.0. Resumed on the base, the step's position weights equal the static decay. Resumed here, they equal the D-PACE weights, and the restore logs the decay-factor warning.
  • Saved here from the pinned recipes, both store and restore decay, and the step uses the static decay.

• FSDP2 probe (2 gloo ranks, the draft as its own fully_shard unit, one rank on the zero-loss branch).

  • DFlash with CE and with KD finishes on both ranks, with the same losses as on the base: 3.47548 and 3.48123 under CE, 3.47471 and 3.47468 under KD.
  • DSpark and Domino fail with "mixed torch.Tensor and DTensor" on the base and here alike, under either objective. Their heads call draft submodules outside the draft unit's forward, before any weighting runs.

• Recipes and pre-commit. Both recipes load through ModelOptDFlashRecipe with objective decay and decay factor 7.0. The launcher's example-resolution tests pass (75). pre-commit run --files <changed files> is clean.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: ❌ DSpark and Domino configs that leave dflash_loss_objective unset now train with D-PACE. Set dflash_loss_objective: decay to keep the previous weighting, as the shipped recipes now do. A run resumed from a checkpoint saved before this change restores that checkpoint's objective (dpace unless it was set) instead of the recipe's, so it switches to D-PACE; finish such runs on the version that started them. Checkpoints still load, and export and generation are unaffected.
• 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?: ✅ One Backward Breaking Changes entry under 0.49.0, since DSpark and Domino have shipped since 0.46.0. It covers runs resumed across the change.
• Did you get Claude approval on this PR?: ❌ Not run yet.

Additional Information

• Merge order: refactor(speculative): one home each for the DFlash family's shared rules #2757 → fix(speculative): DSpark and Domino AR validation and label alignment #2762 → [1/2] refactor(speculative): hand the DFlash-family loss one context per training step #2989 [1/2] → [2/2] refactor(speculative): train the DFlash family through one forward with shared loss terms #2990 [2/2] → this PR. The other PRs on [2/2] refactor(speculative): train the DFlash family through one forward with shared loss terms #2990 (feat(speculative): log the DFlash family's draft_metrics from the example trainer #2992 logging draft_metrics, feat(speculative): DFlash-family loss terms in config; one teacher with an opt-in softcap #2993 loss terms in config, refactor(speculative): draft DSpark and Domino through one AR-validation loop #2994 the shared AR-decode helper) do not depend on this one.
• This PR and the loss-terms-in-config PR (feat(speculative): DFlash-family loss terms in config; one teacher with an opt-in softcap #2993) both change how the DFlash family configures its objective, and they conflict in CHANGELOG.rst, hf_dflash.py, hf_dspark.py, hf_domino.py, dspark.yaml and domino.yaml. Whichever merges second has to be rebased on the other. It merges cleanly with the draft_metrics logging PR (feat(speculative): log the DFlash family's draft_metrics from the example trainer #2992) and the AR-decode helper (refactor(speculative): draft DSpark and Domino through one AR-validation loop #2994).
• The resumed-run change is documented, not shimmed. A restore-time rule that maps older DSpark/Domino checkpoints to decay would need to know which checkpoints predate this change, and nothing marks them reliably:
  • the saved config cannot tell a defaulted dpace from a chosen one;
  • release/0.48.0 is already cut without this change;
  • dev installs record other version strings.
• The Nemotron-3.5-Lightning DSpark warm-start recipe sets dflash_loss_objective: dpace and no decay factor. Until now it trained with uniform position weights. It now trains with D-PACE, as it states, and a run of it resumed across the change switches too. That switch logs no warning, since the recipe sets no decay factor. The recipe is left unpinned because it names an objective. Pin decay there if its fine-tune should keep the old weighting.

🤖 Generated with Claude Code

[copy-pr-bot]

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]

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