perf(speculative): fuse the DFlash-family loss with the lm_head projection by h-guo18 · Pull Request #3010 · NVIDIA/Model-Optimizer

h-guo18 · 2026-10-25T14:02:00Z

What does this PR do?

Type of change: new feature (training performance)

DFlash-family training projects the draft's rows through the frozen lm_head into [rows, vocab] logits, but the loss reads those logits only through per-row statistics:

KD cross-entropy against the teacher;
label cross-entropy, for the CE objective, D-PACE confidences and DFlash2's tv / lambda;
the log-normalizer, for LiLiCorr;
argmax for accuracy, and top-k for the DFlash2 selector and the LiLiCorr lattice.

New dflash_fused_loss (default false) computes exactly those statistics in row chunks fused with the projection. The logits are never materialized whole.

Primitive (plugins/dflash_fused_head.py, fused_head_stats):
- For each chunk of rows it projects the draft rows, and the teacher's final hidden through the same head (or takes the teacher's logits as given).
- It reduces them to the requested statistics: on CUDA through torch.compile, on CPU eagerly.
- For every differentiable statistic it writes the hidden-space Jacobian [rows, H] in the same pass. Backward is then sum_t grad_t[r] * J_t[r], so any downstream combination of the statistics is exact, and no vocab-wide tensor outlives its chunk.
- The top-k logits are a differentiable output whose backward gathers the head rows one chunk at a time.
- The head's logit transform (multiplier, c * tanh(x / c) cap) is supported as explicit arguments. No head on main applies one, so the integration passes none.
Wiring:
- HFDFlashModel._compute_loss: KD with the teacher's logits (online) or its final hidden (offline/streaming), label CE, D-PACE confidences and accuracy.
- The DFlash2 selector's candidates and their logits.
- The LiLiCorr lattice's candidates, their logits and its log-normalizer.
- The statistics reach the variants through _DFlashLossTerms.head_stats. With the flag off, every eager line runs as before.
Refused at convert, rather than silently training something else:
- a class whose forward or _compute_loss does not declare support in its own body. DSpark and Domino add a trainable [r -> V] projection to the logits, which this PR does not cover. They override forward, not _compute_loss, so both owners are checked.
- an lm_head that is not a plain, bias-free nn.Linear with no forward hooks (e.g. quantized or accelerate-hooked);
- a trainable lm_head (dflash_freeze_base_model=false): the fused pass computes no head gradient.

Results re-verified on this branch. One RTX PRO 6000 (250 W cap), torch 2.11, bf16 autocast as in the recipes. Shapes are those of the measured target: hidden 8192, vocab 202048, 2 x 256 blocks x 8 = 4096 draft rows. Timings are the loss path's forward + backward from the draft hidden, median of 5; peak is memory allocated above the step's inputs.

objective	eager (this tree)	fused	peak extra memory
DFlash KD, hidden teacher, D-PACE	226 ms	184 ms (1.23x)	12.3 -> 3.2 GiB
DFlash KD, logits teacher, D-PACE	177 ms	130 ms (1.36x)	12.3 -> 3.9 GiB
DFlash CE, D-PACE	133 ms	122 ms (1.09x)	7.7 -> 2.4 GiB
DFlash2 CE + lambda + selector (recipe)	147 ms	137 ms (1.08x)	9.5 -> 2.4 GiB
DFlash2 KD + selector	247 ms	201 ms (1.23x)	12.3 -> 3.2 GiB
LiLiCorr (recipe objective)	242 ms	260 ms (0.93x)	17.2 -> 5.1 GiB

Equivalence, eager vs fused, same inputs:
- loss rel <= 1.7e-5;
- accuracy and selector metrics identical, lattice metrics within 1.7e-4;
- draft-hidden gradient rel L2 3.5-4.4e-3 (cos >= 0.99999).
Against an fp32 run of the same step, both bf16 paths are equally close: gradient rel L2 6.1e-3 eager vs 6.3e-3 fused (KD), 3.4e-3 vs 3.2e-3 (CE).
The primitive at the measured target's head (x0.177, 20·tanh(x/20), 2048 x 8192 x 202048, bf16) vs autograd on materialized logits, all four statistics and top-16 with random per-row upstream weights: gradient rel L2 4.2e-3, top-1 and top-16 identical.
Why the speedup is smaller here:
- At these shapes the loss path is GEMM-bound on this card. Eager spends 144 of 202 ms kernel time in the three head GEMMs; fused spends 176 of 184 ms on the same FLOPs plus 7.6 ms of compiled statistics.
- Back-to-back GEMMs run at lower clocks under the power cap: one 4096-row head GEMM takes 48 ms after idle and 56 ms back to back. The 2048-row chunking itself costs about 3%.
- Each differentiable statistic costs one Jacobian GEMM, while eager's backward does one GEMM for the summed gradient. LiLiCorr reads two (label CE and the log-normalizer), hence 0.93x here for 3.4x less memory.
Liger 0.8.5 LigerFusedLinearJSDLoss (beta=0, no cap, no per-row weights, scalar output) at 4096 x 8192 x 202048: 278 ms, vs 194 ms for fused_head_stats (KD CE with per-row weights + TV), and 206 ms with the top-16 and the soft cap as well.

Results from a fork of this repo (not main), where this was first built. The fork's eager path did more redundant work than main's:

it projected the teacher inside a checkpoint (re-projected in backward);
in streaming KD it ran the label CE with gradient, although only its detached value (the D-PACE confidences) was used;
its per-chunk slicing of the full logits zero-filled full-size gradients;
it projected the teacher a fifth time for teacher metrics;
it synced the host about six times per step for metrics (.item()), against one on main.

The setup was a DFlash2 drafter (5 layers, hidden 8192, 24K MLP, block 8, 256 anchors) on a 202K-vocab MoE target whose head applies x0.177 and 20·tanh(x/20). Streaming training, per-device batch 2, seq 4096, DDP, compile + FlexAttention on.

One GB300 trainer node, rank-0 medians:
- in-step 0.351 -> 0.257 s (-27%): forward 0.123 -> 0.085, backward 0.193 -> 0.138, optimizer unchanged at 0.020;
- loss-path kernel time ~135-140 -> ~37 ms: head GEMMs 37 -> 21 ms (five projections -> three) plus ~13 ms of compiled statistics;
- over 150 steps, the per-20-step distillation CE, acceptance, selector and teacher metrics tracked the eager run within window noise.
One RTX PRO 6000, same shapes:
- the fork's loss path 627-649 -> 184-197 ms (3.2-3.4x), peak extra memory +11-17 -> +2-4 GiB;
- eager vs fused over 9 loss configurations: loss rel <= 1e-4, accuracy identical, metrics <= 3e-4, hidden-gradient rel 4-10e-3;
- TV-based gradients rel 2.5e-2, because sign(q - p) flips on bf16 rounding of the teacher logits; 4e-3 with identical teacher logits;
- against fp32, the fork's fused gradient was twice as close as its eager one (1.0% vs 2.1%).

What changed from the fork's version:

a config field instead of environment variables;
a per-class support check instead of a module-level registry;
the head transform as explicit arguments instead of attributes read off the head;
differentiable top-k logits instead of a second gather of [rows, k, H] head rows held until backward (1 GiB at the shape above);
reverse KL and teacher top-1 metrics left out (nothing on main reads them);
automatic dynamic shapes in torch.compile, so eager-attention runs with a varying row count do not recompile per shape.

Online KD (teacher = the base's own logits) still costs what it costs today: the base builds [B, seq, vocab] logits and the teacher rows are gathered as [rows, vocab] (1.5 GiB in bf16 at the shape above). The fused path removes only the draft's side.

Usage

python main.py --config modelopt_recipes/general/speculative_decoding/dflash2.yaml ... dflash.dflash_fused_loss=true

mtsp.convert(model, [("dflash", {**dflash_cfg, "dflash_fused_loss": True})])

Testing

New unit tests, on transformers 5.5.4 and 5.18.0 (CI's tf_min / tf_latest). Each takes <= 0.3 s on 4 cores.
- test_dflash_fused_head.py: the primitive vs autograd on materialized logits, in fp32 with 5 chunks. It covers every statistic and its gradient under random per-row upstream weights, the top-k ids and their logits' gradient, no transform / multiplier / multiplier + cap, a hidden and a logits teacher, and values under no_grad.
- test_hf_dflash.py::TestDFlashFusedLoss: one training step, eager vs fused, comparing loss, accuracy and every draft gradient. Cases: online KD (logits teacher, D-PACE), offline KD (hidden teacher, decay) and CE.
- It also checks that convert refuses DSpark, Domino, a trainable lm_head and a hooked lm_head.
- test_hf_dflash2.py: the same for DFlash2 with the selector live, ce and lambda.
- test_hf_lilicorr.py: the same for LiLiCorr with every lattice term on. Its gradient tolerance is scaled to each tensor, because the normalizer's share of the gradient is below a fixed atol.
- Mutation checks catch each of these: perturbed statistics, a dropped top-k gradient, and a non-differentiable normalizer.
test_dflash_fused_head.py, test_hf_dflash.py, test_hf_dflash2.py, test_hf_lilicorr.py, test_hf_dspark.py, test_hf_domino.py and test_hf_dflash_offline.py: 312 passed on 5.5.4 and on 5.18.0.
GPU: tests/gpu/torch/speculative/plugins/test_hf_dflash.py: 16 passed, 1 skipped (needs 2 GPUs). The new TestDFlashFusedLossGPU runs the compiled path for DFlash KD and DFlash2 + selector.
The real-shape re-check above (one RTX PRO 6000).
Not run on this branch: multi-GPU DDP or FSDP2, an end-to-end training run, GB300. The GB300 and convergence numbers above come from the fork.

Dependencies / related PRs

This PR is against main and needs no other PR.

Same code, re-check on merge:
- 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 -> [2/2] refactor(speculative): train the DFlash family through one forward with shared loss terms #2990, and feat(speculative): D-PACE position weighting for DSpark and Domino #2991-refactor(speculative): draft DSpark and Domino through one AR-validation loop #2994 on top. They move _compute_loss into dflash_losses.LossContext and conflict with this PR in hf_dflash.py, hf_dflash2.py and hf_lilicorr.py. Whichever lands second re-ports.
- On top of [2/2] refactor(speculative): train the DFlash family through one forward with shared loss terms #2990 the seam is a LossContext whose ce, teacher, tvd and argmax come from one fused pass:
  - kd -> kd_ce;
  - tvd -> 2 * tv;
  - ce / tv / lambda -> hard_ce with gradient;
  - only the backbone logits of DSpark and Domino could be fused.
- refactor(speculative): one home each for the DFlash family's shared rules #2757 moves the teacher's one-position shift into _teacher_logits, which _fused_head_stats currently applies itself.
Optional, lights up when it lands: feat(speculative): DFlash-family loss terms in config; one teacher with an opt-in softcap #2993.
- Its tvd term would read the primitive's tv.
- Its single teacher (the base's final hidden, projected per row) makes the cheaper hidden-teacher mode the only online one, dropping the [B, seq, vocab] base logits and the gathered teacher rows.
- Its teacher-only softcap needs a teacher-side transform argument, which the primitive does not take yet.
Same code, trivial conflict:
- fix(speculative): weigh DFlash-family losses per anchor position, not per sequence #2986: adjacent class attribute in hf_dflash.py; its per-row weights compose exactly with the fused statistics.
- perf(speculative): deliver training batches out of order; compile + FlexAttention in the DFlash2 recipe #3001 and Add offline AR validation for DFlash drafts on dumped target hidden states #2976: CHANGELOG.rst only. perf(speculative): deliver training batches out of order; compile + FlexAttention in the DFlash2 recipe #3001's compile + FlexAttention + out-of-order loading were on in the fork measurements.
- fix(speculative): keep padding out of DSpark and Domino loss masks #2761: adjacent tests in test_hf_dflash.py only.
Semantic interaction: Honour dflash_freeze_base_model=False in the DFlash base forward #2692, Add an optional language-modelling anchor for a trainable DFlash base #2824, [3/3] Add a packed pretraining data mode for speculative-decoding training #2826, [4/4] Add a draft-loss weight so a co-training control arm is expressible #2827, [4/4] Support Falconer packed inputs and loss masks #2880 train the base. A trainable lm_head is refused here, because the fused pass computes no head gradient.
No open PR brings the fork's soft-capped fake-base lm_head, teacher top-1 metrics or extra lk modes to main.

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?: ✅ Off by default.
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?: ✅
Did you get Claude approval on this PR?: ❌ Draft

Additional Information

🤖 Generated with Claude Code

…ction The DFlash family reads the draft's [rows, vocab] logits only through per-row statistics: KD cross-entropy against the teacher, label cross-entropy, the log-normalizer, argmax and top-k. New DFlashConfig.dflash_fused_loss (default off) computes them in row chunks fused with the frozen lm_head projection (plugins/dflash_fused_head.py). Each chunk is projected and reduced to its statistics, through torch.compile on CUDA, and every differentiable statistic writes its hidden-space Jacobian [rows, H] in the same pass. Backward is sum_t grad_t[r] * J_t[r], so any downstream combination is exact and no vocab-wide tensor outlives its chunk. The top-k logits are a differentiable output whose backward gathers the head rows one chunk at a time. Wired into HFDFlashModel._compute_loss (KD on the teacher's logits or its final hidden, label CE, D-PACE confidences, accuracy), the DFlash2 selector's candidates and the LiLiCorr lattice's candidates and normalizer. With the flag off every eager line runs as before. Convert refuses the flag when the class owning forward or _compute_loss does not declare support (DSpark, Domino), when lm_head is not a plain bias-free nn.Linear, and when lm_head is trainable. Loss path forward + backward at hidden 8192, vocab 202048, 4096 draft rows, one RTX PRO 6000, bf16 autocast, this tree's eager path vs fused: DFlash KD, hidden teacher 226 -> 184 ms, peak +12.3 -> +3.2 GiB DFlash KD, logits teacher 177 -> 130 ms, peak +12.3 -> +3.9 GiB DFlash CE 133 -> 122 ms, peak +7.7 -> +2.4 GiB DFlash2 CE+lambda + selector 147 -> 137 ms, peak +9.5 -> +2.4 GiB DFlash2 KD + selector 247 -> 201 ms, peak +12.3 -> +3.2 GiB LiLiCorr 242 -> 260 ms, peak +17.2 -> +5.1 GiB Loss rel <= 1.7e-5, accuracy and metrics equal, hidden-grad rel L2 3.5-4.4e-3. 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-25T14:02:04Z

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-25T14:02:09Z

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

github-actions · 2026-10-25T14:09:19Z

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-3010/
Built to branch `gh-pages` at 2026-10-25 14:08 UTC. Preview will be ready when the GitHub Pages deployment is complete.

codecov · 2026-10-25T14:10:21Z

Codecov Report

❌ Patch coverage is 99.00990% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.85%. Comparing base (3c423a9) to head (c2cb885).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...opt/torch/speculative/plugins/dflash_fused_head.py	99.20%	1 Missing ⚠️
modelopt/torch/speculative/plugins/hf_dflash.py	97.95%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff            @@
##             main    #3010    +/-   ##
========================================
  Coverage   79.85%   79.85%            
========================================
  Files         661      662     +1     
  Lines       74935    75363   +428     
========================================
+ Hits        59836    60181   +345     
- Misses      15099    15182    +83

Flag	Coverage Δ
unit	`61.22% <99.00%> (+0.09%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Navigation Menu

Repository navigation

perf(speculative): fuse the DFlash-family loss with the lm_head projection - #3010

perf(speculative): fuse the DFlash-family loss with the lm_head projection#3010
h-guo18 wants to merge 1 commit into
mainfrom
haoguo/dflash-fused-loss

h-guo18 commented Oct 25, 2026

Uh oh!

copy-pr-bot Bot commented Oct 25, 2026

Uh oh!

coderabbitai Bot commented Oct 25, 2026

Draft PR not reviewed

Uh oh!

github-actions Bot commented Oct 25, 2026

Built to branch `gh-pages` at 2026-10-25 14:08 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

Uh oh!

codecov Bot commented Oct 25, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

h-guo18 commented Oct 25, 2026

What does this PR do?

Usage

Testing

Dependencies / related PRs

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Oct 25, 2026

Uh oh!

coderabbitai Bot commented Oct 25, 2026

Draft PR not reviewed

Uh oh!

github-actions Bot commented Oct 25, 2026

Built to branch gh-pages at 2026-10-25 14:08 UTC. Preview will be ready when the GitHub Pages deployment is complete.

Uh oh!

codecov Bot commented Oct 25, 2026

Codecov Report

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Built to branch `gh-pages` at 2026-10-25 14:08 UTC.
Preview will be ready when the GitHub Pages deployment is complete.