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)

### 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`](https://huggingface.co/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

```yaml
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):

![warm-start training
curves](https://raw.githubusercontent.com/h-guo18/Model-Optimizer/pr-assets/dspark_nemotron35_warmstart_curves.png)

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](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


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added configurable causal or bidirectional attention for DFlash
models.
* Added optional attention sinks, checkpoint warm starts, and explicit
target-layer selection.
* Improved streaming data handling for shared auxiliary and base hidden
states.
* Added Nemotron-3.5 Lightning DSpark warm-start training and serving
recipes.

* **Bug Fixes**
  * Preserved configured attention behavior during model export.
* Prevented warm-start checkpoints from being reapplied during
restoration.
* Improved checkpoint compatibility, validation, and attention-mask
handling.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

Author: h-guo18

examples/speculative_decoding/eagle_utils.py

@@ -59,12 +59,16 @@ def make_speculative_data_module(
     train_len=None,
     answer_only_loss=False,
     shift_labels=True,
+    final_aux_is_base_hidden=False,
 ) -> dict:
     """Create data module for speculative decoding training.
 
     Args:
         shift_labels: If True, labels are shifted by 1 for autoregressive training (EAGLE3).
             If False, labels are unshifted for diffusion-style training (DFlash).
+        final_aux_is_base_hidden: Streaming only. True when the draft's top aux layer is the
+            base's final layer, so the last captured plane is both the final aux feature and
+            the base (KD-target) hidden instead of an extra dedicated plane.
     """
     # Load chat template from file if provided
     chat_template = None
@@ -115,6 +119,7 @@ def make_speculative_data_module(
             model=data_args.streaming_model_name,
             max_seq_len=train_len,
             answer_only_loss=answer_only_loss,
+            final_aux_is_base_hidden=final_aux_is_base_hidden,
         )
         train_dataset = EagleVllmStreamingDataset(
             entries=ds,

examples/speculative_decoding/main.py

@@ -291,6 +291,7 @@ def train():
         train_len=training_args.training_seq_len,
         answer_only_loss=training_args.answer_only_loss,
         shift_labels=not is_dflash,
+        final_aux_is_base_hidden=recipe.data.final_aux_is_base_hidden,
     )
 
     callbacks = [EagleTrainingPlot(training_args.ar_validate_steps, training_args.estimate_ar)]

modelopt/torch/export/plugins/hf_spec_export.py

@@ -412,21 +412,34 @@ def _export_config(self):
         else:
             config["layer_types"] = ["full_attention"] * draft_config.num_hidden_layers
 
-        # Sliding-window attention: all draft layers use non-causal SWA (MiMo-style). vLLM's
+        # Sliding-window attention: all draft layers use SWA. vLLM's
         # _resolve_layer_attention reads dflash_config.use_swa + swa_window_size; with
-        # layer_types left all "full_attention" it applies a non-causal sliding window to
-        # every draft layer (window from swa_window_size / top-level sliding_window).
+        # layer_types left all "full_attention" it applies a sliding window to every draft
+        # layer (window from swa_window_size / top-level sliding_window).
         swa_window = getattr(self.model, "dflash_swa_window_size", None)
         if swa_window is not None:
             config["sliding_window"] = swa_window
             config["dflash_config"].update(
                 {
                     "use_swa": True,
                     "swa_window_size": swa_window,
-                    "causal": False,
                 }
             )
 
+        # Block-internal attention pattern. Emitted unconditionally (not just under SWA):
+        # vLLM's _dflash_layer_causal treats dflash_config.causal as an all-layer override,
+        # and its default differs per layer type, so writing it explicitly is what keeps
+        # inference consistent with how the draft was actually trained.
+        config["dflash_config"]["causal"] = (
+            getattr(self.model, "dflash_draft_attention", "bidirectional") == "causal"
+        )
+
+        # Learnable per-head attention sink. vLLM reads dflash_config.attention_sink_bias to
+        # decide whether to build the sink parameter and pass it to its attention kernel.
+        if getattr(self.model, "dflash_attention_sink", False):
+            config["dflash_config"]["attention_sink_bias"] = True
+            config["attention_sink_bias"] = True
+
         # Inject the export-time YaRN rope_scaling from the dflash_export_rope_scaling
         # config field (empty dict disables). Mirrors eagle's eagle_export_rope_scaling.
         export_rope_scaling = getattr(self.model, "dflash_export_rope_scaling", None)

modelopt/torch/speculative/config.py

@@ -153,12 +153,64 @@ class DFlashConfig(ModeloptBaseConfig):
         default=None,
         description=(
             "Sliding-window attention (SWA) window size for the DFlash draft. When set, ALL "
-            "draft layers use non-causal sliding-window attention (MiMo-style): each draft "
-            "query attends only to context positions within `dflash_swa_window_size` tokens "
-            "before it, while block-internal attention stays bidirectional. None (default) "
-            "keeps full attention over all context. Must be >= dflash_block_size. Exported to "
-            "the draft config as dflash_config.use_swa/swa_window_size (+ top-level "
-            "sliding_window) so vLLM applies the same window at inference."
+            "draft layers use sliding-window attention: each draft query attends only to "
+            "context positions within `dflash_swa_window_size` tokens before it. None "
+            "(default) keeps full attention over all context. Must be >= dflash_block_size. "
+            "Exported to the draft config as dflash_config.use_swa/swa_window_size (+ "
+            "top-level sliding_window) so vLLM applies the same window at inference. Whether "
+            "block-internal attention is bidirectional or causal is controlled separately by "
+            "`dflash_draft_attention`."
+        ),
+    )
+
+    dflash_draft_attention: Literal["bidirectional", "causal"] = ModeloptField(
+        default="bidirectional",
+        description=(
+            "Attention pattern *inside* each draft block (context attention is always "
+            "restricted to positions before the block's anchor, and additionally windowed "
+            "when dflash_swa_window_size is set).\n"
+            "- 'bidirectional' (default): every query in a block sees all block_size draft "
+            "positions, including ones after it (MiMo-style). This is what ModelOpt has "
+            "always trained and matches drafts such as XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash "
+            "and z-lab/Qwen3.5-9B-DFlash.\n"
+            "- 'causal': a query at block position i only sees draft positions <= i, so the "
+            "block is predicted autoregressively. Required to faithfully train drafts whose "
+            "config declares dflash_config.causal=true, e.g. "
+            "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark.\n"
+            "Exported verbatim to dflash_config.causal, which vLLM's "
+            "qwen3_dflash._dflash_layer_causal reads as a per-model override."
+        ),
+    )
+
+    dflash_attention_sink: bool = ModeloptField(
+        default=False,
+        description=(
+            "Add a learnable per-head attention sink to every draft attention layer. The "
+            "sink is one extra logit per head appended to the attention logits before the "
+            "softmax and dropped afterwards, letting a head place probability mass nowhere "
+            "instead of being forced to attend within a (possibly short) window — the "
+            "GPT-OSS/Nemotron formulation. Adds one `self_attn.attention_sink_bias` "
+            "parameter of shape [num_attention_heads] per layer. Required to load and "
+            "continue training drafts whose checkpoint carries those weights, e.g. "
+            "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16-DSpark. Exported to "
+            "dflash_config.attention_sink_bias for vLLM."
+        ),
+    )
+
+    dflash_init_checkpoint: str | None = ModeloptField(
+        default=None,
+        description=(
+            "Path to an exported draft checkpoint to warm-start from, so training continues "
+            "from published weights instead of a fresh random init. Accepts either a "
+            "directory in the deployment layout this repo exports (``model.safetensors`` "
+            "with no ``dflash_module.`` prefix, alongside ``config.json``) or the "
+            "``model.safetensors`` file itself. Weights are loaded into the draft module "
+            "after it is built, so the architecture still comes from "
+            "``dflash_architecture_config`` — the checkpoint must match it. Any mismatch "
+            "(missing, unexpected, or wrong-shaped tensors) raises rather than silently "
+            "leaving part of the draft randomly initialized. ``embed_tokens``/``lm_head`` "
+            "entries are ignored: the draft takes those from the base model. None "
+            "(default) trains from scratch."
         ),
     )
 
@@ -235,8 +287,8 @@ def _check_dpace_alpha(self) -> "DFlashConfig":
         if not 0.0 < self.dflash_dpace_alpha <= 1.0:
             raise ValueError(f"dflash_dpace_alpha must be in (0, 1], got {self.dflash_dpace_alpha}")
         if self.dflash_swa_window_size is not None:
-            # Block-internal attention is left un-windowed (bidirectional), so the window must
-            # cover a full block; otherwise the effective inference window would differ.
+            # Block-internal attention is left un-windowed, so the window must cover a full
+            # block; otherwise the effective inference window would differ.
             if self.dflash_swa_window_size < self.dflash_block_size:
                 raise ValueError(
                     f"dflash_swa_window_size ({self.dflash_swa_window_size}) must be >= "

modelopt/torch/speculative/dflash/conversion.py

@@ -80,4 +80,7 @@ def restore_dflash_model(
 ) -> nn.Module:
     """Function for restoring a previously converted model to a DFlash model."""
     assert not metadata, "No metadata expected!"
+    # Warm start is a one-time training-time init; the restored weights would overwrite it
+    # anyway, and replaying it on the meta device would raise.
+    config = config.model_copy(update={"dflash_init_checkpoint": None})
     return convert_to_dflash_model(model, config)[0]

modelopt/torch/speculative/dflash/dflash_model.py

@@ -49,4 +49,7 @@ def modify(self, config):
         self.dflash_report_acc = config.dflash_report_acc
         self.dflash_use_torch_compile = config.dflash_use_torch_compile
         self.dflash_swa_window_size = config.dflash_swa_window_size
+        self.dflash_draft_attention = config.dflash_draft_attention
+        self.dflash_attention_sink = config.dflash_attention_sink
+        self.dflash_init_checkpoint = config.dflash_init_checkpoint
         self.dflash_export_rope_scaling = config.dflash_export_rope_scaling