Support fine-tuning released DFlash/DSpark drafters (causal SWA, attention sink, warm start) by h-guo18 · Pull Request #2149 · NVIDIA/Model-Optimizer

h-guo18 · 2026-08-11T14:32:07Z

Support fine-tuning released DFlash/DSpark drafters (causal SWA, attention sink, warm start)

What does this PR do?

Type of change: New feature + bug fix

Adds what ModelOpt was missing to fine-tune an already-published DFlash/DSpark draft
model. The concrete target is
nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark
on its hybrid Mamba/attention/MoE base, but every change is generic.

Before this PR that checkpoint could not be trained faithfully — or even loaded: its
attention-sink tensors were dropped as unexpected keys, its block-causal attention had no
implementation, and its capture layers were silently overwritten with ModelOpt's defaults.

New user-facing options (all default to today's behavior, so existing runs are unchanged):

Option	Values	Purpose
`dflash_draft_attention`	`bidirectional` (default) / `causal`	Block-internal attention pattern. `causal` restricts a query at block position `i` to draft positions `<= i`.
`dflash_attention_sink`	`false` (default) / `true`	Learnable per-head `attention_sink_bias [num_heads]` on every draft layer — one extra logit appended before the softmax and dropped after, so a head can put probability mass nowhere instead of being forced to attend inside its window (the GPT-OSS formulation).
`dflash_init_checkpoint`	path	Warm-start the draft from an exported checkpoint instead of a random init. Any missing/unexpected/wrong-shaped tensor raises rather than warns.
`dflash_architecture_config.target_layer_ids`	list	Which base layers feed the draft's `fc`. Previously recomputed unconditionally with no override.

Bugs fixed along the way (each one silently corrupts training rather than failing):

The exporter hard-coded dflash_config.causal: False and only wrote it under SWA, so even
a correctly-trained causal draft would be served non-causally. It now reflects the trained
setting, and emits attention_sink_bias when enabled.
_build_generate_swa_mask returned None whenever swa_window_size was unset, which
would have dropped the causal structure at generation time while training used it.
target_layer_ids was recomputed from the uniform default on every convert. The released
drafter uses [1,5,19,29,41,51]; the default for a 52-layer base is [1,11,20,30,39,49]
— different layers. Here it surfaced as a matmul shape error only because the plane
counts disagreed; with a matching count it would have trained on the wrong features
silently.
The streaming dataset assumed the draft's aux layers all sit below the base's final layer
(aux = planes[:-1], target = planes[-1]). A draft whose top aux id is the final layer
cannot get an extra plane — vLLM captures each layer once — so final_aux_is_base_hidden
now lets the last plane serve both roles. It is derived from the model, not configured by
hand.
DSpark head weights load from either the flat layout ModelOpt exports (upstream DeepSpec
convention) or the nested markov_head. layout the NVIDIA release uses. Without the remap
the two [131072, 512] Markov tables — ~14% of the draft's parameters — stay randomly
initialized while everything else warm-starts, with no error.
nemotron_h is enabled in _FINAL_NORM_TYPE_BY_MODEL_TYPE: despite the hybrid stack,
NemotronHModel.norm_f is a plain RMSNorm, and without the entry the offline/streaming
fake base raises instead of reconstructing the distillation target.

Usage

dflash:
  dflash_init_checkpoint: /path/to/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark
  dflash_draft_attention: causal
  dflash_attention_sink: true
  dflash_swa_window_size: 1024
  dflash_block_size: 8
  dflash_mask_token_id: 990
  dflash_architecture_config:
    target_layer_ids: [1, 5, 19, 29, 41, 51]

A full worked example is at
modelopt_recipes/general/speculative_decoding/dspark_nemotron35_warmstart.yaml.

Testing

Unit tests — 124 pass (test_hf_dflash.py, test_hf_dspark.py, test_hf_domino.py,
test_hf_dflash_offline.py, test_modeling_final_norm.py), 32 of them new: causal mask
structure (lower-triangular per block, no cross-block leakage, context visibility
unchanged), the sink math (degenerates to plain attention at -inf, absorbs mass
monotonically, receives gradient), warm-start load/reject paths, Markov key remapping, and
explicit target_layer_ids.

Checkpoint compatibility — the released drafter loads with zero missing/unexpected keys
and zero shape mismatches; all 77 tensors (6 attention sinks and both Markov tables
included) match bit-exactly, and a training step runs with gradients reaching the sink and
Markov parameters.

End-to-end streaming training — Nemotron-3.5 base served by vLLM (1 node, TP8) feeding
8 trainer GPUs over NIXL; the draft warm-starts from the released checkpoint and trains with
causal + sink + SWA 1024. 128 Daring-Anteater conversations, 20 epochs (the plot shows the
first 5, where the trend is clearest — the curves flatten after that):

Over the first 5 epochs loss falls 1.85 → 1.36 and train accuracy rises
0.25 → 0.49; across the full 20 epochs they reach 1.21 and 0.48 (peak 0.54)
before flattening. This validates the pipeline end-to-end — capture layers, plane split,
mask direction, sink loading and warm-start weights all have to be right for this curve to
appear. It is not a model-quality result: 128 samples over 20 epochs overfits by
construction, and the corpus is not generated by the base model, so the absolute numbers are
not meaningful.

TODO (follow-up)

A complete, robust checkpoint/config converter. Both conversions are handled ad hoc here:

Draft config → training config. The recipe transcribes ~15 fields by hand from the
drafter's config.json. Only the shape-bearing ones (num_hidden_layers,
num_attention_heads, intermediate_size, markov_rank) fail loudly when mistyped; the
rest — mask_token_id, causal, swa_window_size, block_size — train "successfully" on
a wrong value and only surface later as a mysteriously low acceptance length. A converter
should derive the whole block from the checkpoint, including its aliases (pard_token,
dspark_markov_rank, dflash_query_causal, top-level sliding_window /
attention_sink_bias) and duplicated fields.
Weight layout. The markov_head. remap is a load-time hook. A converter should normalize
layouts explicitly, and decide whether export should also emit the release's aliases so a
round-trip reproduces the original format (today it renames architectures to
DFlashDraftModel).
Base config. Serving this base on vLLM needs its config.json layer-type vocabulary
updated for the transformers-5 path (mamba → linear_attention, attention →
full_attention, plus a matching hybrid_override_pattern). That is done by hand today and
is not covered by this PR.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed.
Is this change backward compatible?: Yes
Did you write any new necessary tests?: Yes
Did you add or update any necessary documentation?: Yes
Did you update Changelog?: No

### What does this PR do? Type of change: New feature The DFlash/DSpark draft could only be trained with one attention pattern: non-causal (MiMo-style) blocks, optionally sliding-windowed, with no attention sink. Drafters that declare `dflash_config.causal=true` and carry `self_attn.attention_sink_bias` weights — e.g. `nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark` — could not be trained faithfully or even loaded: the sink tensors were dropped as unexpected keys, and the exporter hardcoded `causal: False`, so training and vLLM inference disagreed. Add two opt-in `DFlashConfig` options, both defaulting to today's behavior: - `dflash_draft_attention: "bidirectional" | "causal"` selects the block-internal attention pattern. `"causal"` restricts a query at block position `i` to draft positions `<= i`. - `dflash_attention_sink: bool` adds a learnable per-head `attention_sink_bias [num_attention_heads]` to every draft layer: one extra logit appended to the attention logits before the softmax and dropped after, so a head can place probability mass nowhere instead of being forced to attend inside its window (the GPT-OSS formulation). Fused SDPA cannot express that extra column, so this runs an eager attention path, used only when the option is enabled. Three latent issues surfaced while wiring this up and are fixed here: - `_build_generate_swa_mask` returned `None` whenever `dflash_swa_window_size` was unset, which would have silently dropped the causal structure at generation time while training used it. - The exporter only wrote `dflash_config.causal` under SWA, hardcoded to `False`. It is now written unconditionally from the trained setting, since vLLM's `_dflash_layer_causal` treats it as an all-layer override whose default varies per layer type. - DSpark head weights load from either the flat layout ModelOpt exports (matching upstream DeepSpec) or the nested `markov_head.` layout used by released NVIDIA drafters, via a load pre-hook. The export format is unchanged, so drafters already trained and deployed with flat names keep working. Also enable `nemotron_h` in `_FINAL_NORM_TYPE_BY_MODEL_TYPE`: despite the hybrid Mamba/attention/MoE stack, `NemotronHModel.norm_f` is a plain RMSNorm, and without the entry the offline/streaming fake base raises rather than reconstructing the distillation target. ### Usage ```yaml dflash: dflash_draft_attention: causal # default: bidirectional dflash_attention_sink: true # default: false dflash_swa_window_size: 1024 ``` ### Testing - 111 unit tests pass (`test_hf_dflash.py`, `test_hf_dspark.py`, `test_hf_domino.py`, `test_hf_dflash_offline.py`, `test_modeling_final_norm.py`), including 19 new ones covering the causal mask structure, the sink math and its gradient, key remapping, and export round-trips. - Verified `nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark` loads with zero missing/unexpected keys and zero shape mismatches, and that all 77 tensors — the 6 sink biases and both Markov tables included — match the checkpoint bit-exactly. - Default (`bidirectional`, no sink) mask output is unchanged element-wise, so existing drafters train exactly as before. ### Before your PR is "*Ready for review*" - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes - **Did you write any new necessary tests?**: Yes - **Did you add or update any necessary documentation?**: Yes - **Did you update [Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**: No Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

### What does this PR do? Type of change: New feature Training always started the draft module from a random init. Continuing from a published drafter — the usual way to fine-tune a release such as `nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark` — had no entry point. Add `dflash_init_checkpoint`: a path to an exported draft checkpoint, loaded into the draft module right after it is built (before the dtype/device move, so the weights are cast with the rest of the module). It accepts the export directory or the `model.safetensors` file itself, tolerates a `dflash_module.` prefix so a raw training checkpoint works too, and ignores `embed_tokens`/`lm_head`, which the draft takes from the base model. Any missing, unexpected, or wrong-shaped tensor raises instead of warning. Loading only part of a draft and leaving the rest randomly initialized looks like a warm start but trains from a corrupted starting point, and nothing downstream would flag it. The shape check resolves the module's load pre-hooks first, so a tensor arriving under DSpark's nested `markov_head.` layout is validated too rather than slipping through to a less obvious failure later. ### Usage ```yaml dflash: dflash_init_checkpoint: /path/to/exported/drafter ``` ### Testing - 119 unit tests pass, 8 new ones here: weights actually loaded, file vs directory paths, round-trip with sink + causal attention, the default path still random-inits, and four rejection cases (missing path, draft depth mismatch, sink on/off mismatch, wrong shape under the nested layout). - End-to-end against the released Nemotron-3.5 DSpark drafter: `convert()` with `dflash_init_checkpoint` reproduces all 77 draft tensors bit-exactly (6 attention sinks and both Markov tables included), applies `causal` + sink + SWA 1024 from the checkpoint's own config, and runs a training step with a finite loss and gradients reaching the sink and Markov parameters. ### Before your PR is "*Ready for review*" - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes - **Did you write any new necessary tests?**: Yes - **Did you add or update any necessary documentation?**: Yes - **Did you update [Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**: No Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

### What does this PR do? Type of change: Bug fix `HFDFlashModel.modify` recomputed `target_layer_ids` from the uniform default on every convert, with no way to override it. A published draft is trained against specific capture points — the released `nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark` uses `[1,5,19,29,41,51]`, while the default for a 52-layer base is `[1,11,20,30,39,49]` — so fine-tuning one fed the draft features from layers it had never seen. Here that surfaced as a matmul shape error, but only because the plane counts happened to disagree; with a matching count it would have trained silently on the wrong features. An explicit `dflash_architecture_config.target_layer_ids` is now used as given, validated for length (one per draft layer) and range. That fix exposes a second case. The streaming dataset splits captured planes as `aux = planes[:-1]`, `base/KD target = planes[-1]`, which assumes the draft's aux layers all sit below the base's last layer — true for ModelOpt's default ids (they stop at `num_target_layers - 3`) and for the M3 example, but not for a draft whose top aux id *is* the final layer. vLLM captures each layer at most once, so no extra plane can be requested. `final_aux_is_base_hidden` makes the last plane serve both roles; it is derived from the model in `main.py` rather than configured by hand, since a wrong manual value fails as an opaque matmul error deep in the draft's `fc`. ### Usage ```yaml dflash: dflash_architecture_config: target_layer_ids: [1, 5, 19, 29, 41, 51] # match the checkpoint ``` ### Testing - 124 unit tests pass, 5 new here: explicit ids are used, the default still applies when unset, wrong count and out-of-range both raise, and the ids round-trip through export. - End-to-end streaming run against the released drafter (vLLM-served Nemotron-3.5 base, 8 trainer GPUs): the trainer now logs `target_layer_ids=[1,5,19,29,41,51], final_aux_is_base_hidden=True`, and training converges — loss 1.85 → 1.21, train_acc 0.25 → 0.48 over 20 epochs. ### Before your PR is "*Ready for review*" - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes - **Did you write any new necessary tests?**: Yes - **Did you add or update any necessary documentation?**: Yes - **Did you update [Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**: No Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

### What does this PR do? Type of change: Documentation Worked example for continuing training from a published DFlash/DSpark draft: it wires up `dflash_init_checkpoint`, the checkpoint's own `target_layer_ids`, and the causal + attention-sink settings the released Nemotron-3.5 drafter was trained with. The architecture block is transcribed by hand from the drafter's `config.json`, which is the weak point and is flagged as a TODO in the file: only the shape-bearing fields fail loudly when mistyped, while the mask token, causal flag, window size and block size all train "successfully" on a wrong value and surface later as a mysteriously low acceptance length. A converter should derive the block from the checkpoint's config (including its aliases) instead. ### Testing Used as-is for the end-to-end streaming run reported in this PR. ### Before your PR is "*Ready for review*" - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes - **Did you write any new necessary tests?**: N/A (example config) - **Did you add or update any necessary documentation?**: Yes - **Did you update [Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**: No Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

Loss and train-accuracy over the 20-epoch streaming run used to validate the end-to-end pipeline (vLLM-served Nemotron-3.5 base, draft warm-started from the released DSpark checkpoint). Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

copy-pr-bot · 2026-08-11T14:32:11Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

coderabbitai · 2026-08-11T14:33:08Z

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 56c99b85-2253-4276-a1b4-b6307464a9f1

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

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

github-actions · 2026-08-11T14:35:50Z

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

The full 20-epoch run flattens after epoch ~5, which obscures the trend the plot is meant to show. Same run, first 80 steps. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

codecov · 2026-08-11T14:48:00Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.76%. Comparing base (e4fe1e5) to head (f723144).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2149      +/-   ##
==========================================
+ Coverage   78.73%   78.76%   +0.03%     
==========================================
  Files         522      522              
  Lines       60357    60439      +82     
==========================================
+ Hits        47523    47606      +83     
+ Misses      12834    12833       -1

Flag	Coverage Δ
unit	`55.35% <100.00%> (+0.07%)`	⬆️

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.

h-guo18 added 5 commits August 11, 2026 10:55

Show the first 5 epochs in the warm-start curves

f723144

The full 20-epoch run flattens after epoch ~5, which obscures the trend the plot is meant to show. Same run, first 80 steps. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

Navigation Menu

Uh oh!

Support fine-tuning released DFlash/DSpark drafters (causal SWA, attention sink, warm start) - #2149

Support fine-tuning released DFlash/DSpark drafters (causal SWA, attention sink, warm start)#2149
h-guo18 wants to merge 6 commits into
mainfrom
feat/dspark-causal-swa-sink

h-guo18 commented Aug 11, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Aug 11, 2026

Uh oh!

coderabbitai Bot commented Aug 11, 2026 •

edited

Loading

Review skipped

Uh oh!

github-actions Bot commented Aug 11, 2026 •

edited

Loading

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

Uh oh!

codecov Bot commented Aug 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

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

Uh oh!

Support fine-tuning released DFlash/DSpark drafters (causal SWA, attention sink, warm start)

What does this PR do?

Usage

Testing

TODO (follow-up)

Before your PR is "Ready for review"

Uh oh!

copy-pr-bot Bot commented Aug 11, 2026

Uh oh!

coderabbitai Bot commented Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Uh oh!

github-actions Bot commented Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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

Uh oh!

codecov Bot commented Aug 11, 2026

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

h-guo18 commented Aug 11, 2026 •

edited

Loading

coderabbitai Bot commented Aug 11, 2026 •

edited

Loading

github-actions Bot commented Aug 11, 2026 •

edited

Loading

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