[h-guo18]

What does this PR do?

Type of change: Refactor. Training behavior is unchanged except for the two disclosed cases below. One backward-breaking API change: DSpark's and Domino's training metrics move under draft_metrics (see Metrics).

This PR is stacked on #2989, whose branch is its base; that PR is stacked on #2762. It is the second of two PRs: merge #2989 first, then this one. A new DFlash-family variant is now added one way: it reuses DFlash's training forward and changes only its loss, plus its head if it has one.

• One training forward. DSpark and Domino each carried a copy of DFlash's training forward. Each copy ran the base model, built the loss mask, sampled anchors and built the draft inputs. Family fixes went missing in those copies; the padding-aware loss mask, for example, never reached them. HFDFlashModel.forward is now the only training forward, and no variant overrides forward(). DSpark and Domino change only the hooks [1/2] refactor(speculative): hand the DFlash-family loss one context per training step #2989 introduced, plus _build_objective:
  • _online_base_outputs: DSpark uses the inner model's final hidden state, which the loss projects only at the rows it reads. Domino takes only the target hidden states.
  • _draft_logits: DSpark's Markov bias or Domino's GRU correction on the backbone's logits. The AR-validation decodes from fix(speculative): DSpark and Domino AR validation and label alignment #2762 are unchanged. Domino's generation test now checks its drafts against _draft_logits.
  • _build_objective: the block objective, ce and tvd for DSpark and ce for Domino.
  • _position_weighting: DSpark and Domino still always decay.
  • _compute_loss: the objective plus the terms the variant owns. DSpark adds its confidence head, whose target reuses the TVD term's per-token values. Domino blends (λ) the objective on the corrected logits with the objective on the backbone logits.
• Loss terms as components. The divergences join dflash_losses.py, which names no variant.
  • The divergences (ce, kd, tvd, tv and lambda) form a closed registry. Each one reduces under the position weights the way the old code did.
  • Each divergence also declares whether it reads the teacher and whether it backpropagates through the cross-entropy. These declarations decide when offline training reconstructs the teacher's logits and when the cross-entropy keeps its gradient.
  • An objective is an ordered [(name, weight, term)] list, summed left to right. modify() builds it through _build_objective: DFlash and LiLiCorr keep kd or ce, and DFlash2 picks tv or lambda. DFlash2's tv and lambda are now objective terms, so they pass through the zero-loss gate instead of replacing the loss after it (change (b) below).
  • LossContext also computes the per-token TVD lazily, at most once. The TVD chunk compiles under dflash_use_torch_compile for every variant, no longer only in DSpark's modify().
• Metrics (backward breaking). Every variant now returns its metrics under draft_metrics, as detached 0-dim tensors. Each objective term is reported as loss_<term>, and backbone_accuracy is reported when a head sits on the backbone.
  • Removed keys: dspark_metrics and domino_metrics, which have shipped since 0.46.0. [1/2] refactor(speculative): hand the DFlash-family loss one context per training step #2989 already replaced selector_metrics and lilicorr_metrics.
  • Renamed inner keys: DSpark's ce_loss/l1_loss/base_accuracy become loss_ce/loss_tvd/backbone_accuracy. Domino's final_loss/base_loss/base_accuracy become loss_final/loss_backbone/backbone_accuracy. All other names stay as they were.
  • DSpark's and Domino's values change from Python floats to tensors, so a reader calls .item() where it needs a float.
  • Nothing in the repo reads these keys. Code outside it that does must switch to draft_metrics. The CHANGELOG entry from [1/2] refactor(speculative): hand the DFlash-family loss one context per training step #2989 now covers all four keys and the renamed inner keys.
• Removed: _compute_dspark_loss, _compute_domino_loss, _apply_markov_head and _apply_domino_head. Also removed: DSpark's and Domino's forward() overrides.

The three zero losses stay as they were:

• Inside the loss, the gate still routes through the draft's output (final.sum() * 0.0), so an FSDP2 unit around the draft still runs its backward. Domino's λ-blend sits behind the same gate.
• The forward's early return, where no draft has run, keeps _draft_zero_loss.
• DSpark's degenerate step keeps _draft_zero_loss, which also feeds its confidence head for DDP.

Disclosed changes:

• (a) DSpark and Domino adopt DFlash's loss mask, which intersects labels, attention_mask and loss_mask. This is the same fix as the separate padding PR, through the same _build_loss_mask helper. Merging the two conflicts only where the padding PR edits the DSpark and Domino forwards that this PR removes; keeping this PR's side resolves it (see Testing). The result is identical whenever attention_mask is all ones where labels are valid. That holds in every shipped recipe and in streaming and offline batches.
• (b) A DFlash2 step under tv or lambda whose total position weight is <= 1 now takes the zero-loss gate like every other objective. Before, the blend was computed over that tiny weight sum. The selector term is unchanged.
• Error path only: an offline DSpark batch that has neither base logits nor base_model_hidden_states now raises at the start of the forward, not after the draft has run.

Usage

Configs are unchanged. A variant is now written as hooks:

class HFMyDraftModel(HFDFlashModel):
    def _build_objective(self, config):
        return [objective_term("ce", 0.5), objective_term("tvd", 0.5)]

    def _draft_logits(self, batch):
        backbone = super()._draft_logits(batch).backbone
        return DraftLogits(final=backbone + self.my_head(batch), backbone=backbone)

outputs = model(**batch)
outputs.draft_metrics["loss_tvd"]  # detached 0-dim tensor; .item() for a float

Testing

• tests/unit/torch/speculative plus tests/unit/torch/export/test_hf_spec_rope_export.py: 520 passed on transformers 5.5.4, 5.12.1 and 5.18.0 (513 on [1/2] refactor(speculative): hand the DFlash-family loss one context per training step #2989; the other 7 are this PR's new cases). 5.5.4 and 5.18.0 fall within CI's tf_min and tf_latest pins. On 4 CPUs, each changed test takes at most 0.7 s and each new case under 0.1 s.
• tests/gpu/torch/speculative/plugins on one GPU: 20 passed, 1 skipped (it needs 2 GPUs).
• Bitwise A/B against [1/2] refactor(speculative): hand the DFlash-family loss one context per training step #2989. It compares loss, every gradient, train_acc, the metrics under the rename above, and the generated draft tokens.
  • One training step and one generation step for each of 23 cases, in fp32 and bf16: 44/46 identical. The cases cover DFlash (KD and CE), DFlash2 (ce, kd, tv and lambda, with the selector), LiLiCorr, DSpark (vanilla and rnn heads, with the confidence head) and Domino. Each runs offline and online, with D-PACE and with decay, including zero-loss steps. Domino's online cases generate with fix(speculative): DSpark and Domino AR validation and label alignment #2762's correction-head decode on both sides.
  • The other 2 runs are change (b), on a step that keeps one label of weight < 1. The loss goes from 0.96958977 to 0 in fp32 (0.968748331 to 0 in bf16). The backbone gradients lose the TV term, and the selector gradients and all metrics are unchanged. The same step under lambda goes from 2.86922026 to 0 in fp32 (2.87524128 to 0 in bf16).
  • Further cases, all identical:
    • the DSpark gated head, and the rnn head online;
    • DSpark and Domino with the KD flag and D-PACE set, both of which they still ignore;
    • DFlash2 ce on a zero-loss step, DFlash2 KD online, and LiLiCorr KD on a zero-loss step;
    • six cases whose labels are -100 on some tokens.
  • Each refactor commit against its parent, on the cases it affects: 30/32, 8/8 and 6/6. The 2 non-identical runs are change (b) in the commit that introduces it.
  • The new loss_<term> metrics equal the loss for DFlash and origin_loss for LiLiCorr, bitwise.
  • Against fix(speculative): DSpark and Domino AR validation and label alignment #2762's head, with [1/2] refactor(speculative): hand the DFlash-family loss one context per training step #2989, the results are the same: every run identical except the four runs of change (b).
• FSDP2 test: 2 ranks on gloo, the draft as its own fully_shard unit, and one rank on the in-loss zero branch. Both ranks finish, under CE and under KD, with the same losses as on [1/2] refactor(speculative): hand the DFlash-family loss one context per training step #2989 and on fix(speculative): DSpark and Domino AR validation and label alignment #2762's head. Domino under the same wrapping stops with the same mixed Tensor/DTensor error on fix(speculative): DSpark and Domino AR validation and label alignment #2762's head, on [1/2] refactor(speculative): hand the DFlash-family loss one context per training step #2989 and on this one; this PR does not change that.
• Merged with the separate padding PR's branch: Git reports conflicts only in hf_dspark.py and hf_domino.py, where the padding PR edits the forwards that this PR removes. Keeping this PR's side of both files gives this PR's tree plus the padding PR's regression test, with one _build_loss_mask. On that tree, the padding PR's regression test and the DSpark and Domino test files pass.
• New tests:
  • A DFlash2 step with nothing to train returns the zero loss under ce, tv and lambda. Before this PR, tv and lambda gave 0.968 and 3.352 on it.
  • Each shared divergence (ce, kd, tvd and tv) runs through the real _loss_context and _compute_loss under non-uniform decay weights. Each result equals, bitwise, a verbatim copy of the reduction it replaced. Before, the lambda test compared DFlash2's own copy of the CE reduction with DFlash's. That check went away when the blend started calling the shared term.
  • Mutation check: normalizing ce, kd, tvd or tv by the supervised count instead of the weight sum, or reducing ce under the mask, passes the rest of the suite (519 tests) and fails exactly one case of this test.
• The terms' needs_ce_grad declarations now decide whether the cross-entropy keeps its gradient. [1/2] refactor(speculative): hand the DFlash-family loss one context per training step #2989 added the test for that: under D-PACE, the context that a real training forward passes to the loss keeps the gradient for ce, tv and lambda but not for KD. Its tv, lambda and KD cases now also pin the declarations.
  • Mutation check: flipping needs_ce_grad on tv, lambda or kd passes the rest of the suite (519 tests) and fails exactly that case.
  • On a step configured like dflash2.yaml, flipping the lambda declaration leaves the loss bitwise unchanged, but 33 of 36 draft gradients change, because the backbone then trains only through the selector.
• Updated tests: DSpark's zero-weight DDP test runs the real _loss_context/_compute_loss. Domino's generation test from fix(speculative): DSpark and Domino AR validation and label alignment #2762 teacher-forces the drafts through _draft_logits; it still fails if drafting or training skips the correction head or feeds it the wrong prefix. DSpark's and Domino's metrics are read from draft_metrics.

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?: ❌ Configs and checkpoints are unchanged. However, draft_metrics replaces the forward output's dspark_metrics and domino_metrics. DSpark's and Domino's inner keys are renamed, and values that were Python floats are now 0-dim tensors (see Metrics and the CHANGELOG entry).
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: ✅ The shared terms live in dflash_losses.py, which carries the SpecForge and DeepSpec notices since [1/2] refactor(speculative): hand the DFlash-family loss one context per training step #2989.
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅ The Backward Breaking Changes entry from [1/2] refactor(speculative): hand the DFlash-family loss one context per training step #2989, under 0.49.0, now also names DSpark's and Domino's keys.
• Did you get Claude approval on this PR?: ❌ Not run yet.

Additional Information

Against #2989, the source changes are +297 / −507 lines. Follow-ups build on the hooks: DSpark and Domino honoring dflash_loss_objective, and logging draft_metrics in the trainer.

Merge order: #2989, then this PR.

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