[h-guo18]

What does this PR do?

Type of change: Refactor. Training behavior is unchanged. One backward-breaking API change: DFlash2's and LiLiCorr's training metrics move under a single key (see Metrics).

This PR is stacked on #2762, which is stacked on #2757. #2762's branch is its base. It is the first of two PRs: merge it first, then #2990, which builds the shared loss terms on this seam and moves DSpark and Domino onto the shared forward.

• One loss context per step. HFDFlashModel.forward built the batch, ran the draft and computed the loss inline. The variants built on it reached the loss's pieces piecemeal. DFlash2 asked _compute_loss for its unreduced terms (return_terms, _DFlashLossTerms) to recompose the block objective. DFlash2 and LiLiCorr overrode forward() only to carry out the metrics their _compute_loss had left on the model. The forward now hands the loss one object per step:

  • DraftBatch: the batch, its blocks and the draft's hidden states.
  • LossContext: the labels, the draft's logits and the position weights. DraftLogits holds the distribution the draft serves and the backbone's logits before any head. The per-token cross-entropy and the teacher's logits are computed lazily, each at most once. The cross-entropy keeps its gradient only when the objective trains through it, not under KD, even though D-PACE reads it first under no_grad.
  • _compute_loss(ctx) returns a LossOutput: the loss, the accuracy train_acc reports, and detached 0-dim metrics, which the forward returns as draft_metrics.

• Hooks. What a variant may change becomes a hook on HFDFlashModel:

  • _online_base_outputs: what the base model provides online.
  • _draft_logits: a head on the backbone's logits.
  • _position_weighting: D-PACE or decay.
  • _compute_loss: the loss on the context.

  DFlash2's tv/lambda blend reads the cross-entropy and the weights from the context. DFlash2 and LiLiCorr return their metrics instead of overriding forward(). DSpark and Domino keep their own forwards in this PR. The forward's loss mask now comes from _build_loss_mask. The separate padding PR against main adds the same helper, with the same text, at the same place (right after get_exporter), so a merge of the two keeps one copy of it.

• A byte-identical move. The block-supervision helpers, the position weighting, _masked_top1, the loss context and DSpark's per-token TVD move from hf_dflash.py and hf_dspark.py into the new dflash_losses.py. The new module names no variant. Every moved definition is byte-identical to its previous text, and only import lines change elsewhere. The move is its own commit, separate from the reshaping.

• Metrics (backward breaking). DFlash2's selector_metrics and LiLiCorr's lilicorr_metrics are replaced by one key, draft_metrics, as detached 0-dim tensors with no host sync.

  • Inner key names stay as they were.
  • LiLiCorr's values change from Python floats to tensors, so a reader calls .item() where it needs a float. DFlash2's were already tensors.
  • Both keys ship in 0.48.0. Nothing in the repo reads them. Code outside it that does must switch to draft_metrics. The CHANGELOG lists this under Backward Breaking Changes.
  • DSpark and Domino keep dspark_metrics and domino_metrics until [2/2] refactor(speculative): train the DFlash family through one forward with shared loss terms #2990, which also reports DFlash's objective terms. Until then DFlash's own draft_metrics is empty.

• Removed: return_terms, _DFlashLossTerms and its TODO. Also removed: DFlash2's and LiLiCorr's forward() overrides, with the _selector_metrics/_lilicorr_metrics side channels.

The 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.
• The forward's early return, where no draft has run, keeps _draft_zero_loss.

Usage

Configs are unchanged. A variant changes its head and its loss through hooks:

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

    def _compute_loss(self, ctx):
        out = super()._compute_loss(ctx)  # the block objective on ctx.logits.final
        return out._replace(loss=out.loss + self.my_head_loss(ctx))

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

Testing

• tests/unit/torch/speculative plus tests/unit/torch/export/test_hf_spec_rope_export.py: 513 passed on transformers 5.5.4, 5.12.1 and 5.18.0 (509 on fix(speculative): DSpark and Domino AR validation and label alignment #2762's head; the other 4 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 3.4 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 fix(speculative): DSpark and Domino AR validation and label alignment #2762's head. It compares loss, every gradient, train_acc, the metrics with DFlash2's and LiLiCorr's keys mapped onto draft_metrics, and the generated draft tokens.
  • One training step for each of 23 cases, plus one generation step for the 7 online ones, in fp32 and bf16: 46/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. Together the cases run offline and online, with D-PACE and with decay, including zero-loss steps.
  • 15 further cases, 30/30 identical: the DSpark gated head, the rnn head online, the KD flag and D-PACE on DSpark and Domino, steps whose position weights sum to <= 1 under DFlash2 (ce and lambda) and LiLiCorr KD, DFlash2 KD online, and six cases whose labels are -100 on some tokens.
  • The first commit alone against fix(speculative): DSpark and Domino AR validation and label alignment #2762's head, on the DFlash, DFlash2 and LiLiCorr cases: 32/32 identical.
• The second commit (the move), checked mechanically: the 13 definitions that leave hf_dflash.py and hf_dspark.py reappear byte-identically in dflash_losses.py. No other definition changes, and the other files change only in import lines.
• 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 fix(speculative): DSpark and Domino AR validation and label alignment #2762's head. DFlash2 and LiLiCorr under the same wrapping stop with the same mixed Tensor/DTensor error on fix(speculative): DSpark and Domino AR validation and label alignment #2762's head and on this one; this PR does not change that.
• Merged with the separate padding PR's branch in both orders: no conflict, the same tree either way, and one _build_loss_mask. On that tree, the padding PR's regression test and the DSpark and Domino test files pass.
• New test: under D-PACE, the context that a real DFlash2 training forward passes to the loss keeps the cross-entropy's gradient for ce, tv and lambda but not for KD.
  • Mutation check: if _loss_context always asks for the gradient, LossContext ignores ce_grad, or LossContext.ce always keeps its graph, the rest of the suite still passes (512 tests) and only the KD case fails.
  • The bitwise A/B cannot see that change: on the DFlash, DFlash2 and LiLiCorr KD cases the loss and every gradient stay identical (6/6). Each KD step under D-PACE, the default configuration, would just hold a vocab-wide graph that it never backpropagates.
  • Dropping the gradient instead already failed three existing tests, and now also fails the ce, tv and lambda cases.
• Updated tests: the _compute_loss tests go through _loss_context and _compute_loss, and the selector and LiLiCorr metrics are read from draft_metrics. "KD without base_outputs raises" becomes "KD without a teacher raises": the context always carries the base outputs, so the remaining failure is a batch with neither the base logits nor its final hidden.

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 selector_metrics and lilicorr_metrics, and LiLiCorr's 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: ✅ dflash_losses.py carries the SpecForge and DeepSpec notices of the two files its code moved from. Like them, it is excluded from the license-insertion hook.
• Did you write any new necessary tests?: ✅ One new test pins which objectives keep the cross-entropy's gradient, which the bitwise A/B cannot see. The existing tests now drive the new seam (_loss_context, _compute_loss, draft_metrics), and the A/B covers the rest.
• Did you update Changelog?: ✅ One Backward Breaking Changes entry for the two metrics keys.
• Did you get Claude approval on this PR?: ❌ Not run yet.

Additional Information

Against #2762's head, the source changes are +489 / −478 lines. Of the added lines, 132 are ten definitions moved byte-identically into dflash_losses.py, and 37 are its license header. DraftLogits, LossOutput and LossContext (68 lines) are new in the first commit and move with them. Reshaping and moving stay in separate commits.

Merge order: this PR, then #2990.

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