> This page is for version Latest.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# Train a DFlash Drafter for Speculative Decoding

> Train a DFlash block-parallel speculative-decoding drafter with NeMo AutoModel, from data through multi-GPU training, including the DFlash 2, Domino, and JetSpec variants.

**A guide for training a DFlash speculative-decoding drafter to accelerate LLM
inference with [NeMo AutoModel](https://github.com/NVIDIA-NeMo/Automodel).**

---

## What is DFlash?

DFlash is a *block-parallel* drafter. Where an autoregressive drafter proposes
one token at a time, DFlash proposes an entire block of `block_size` tokens in a
single **non-causal** forward pass: the block's first position holds the real
anchor token and the rest are `MASK`, and the draft "denoises" the whole block at
once conditioned on hidden states captured from the frozen target at
`target_layer_ids`. The draft shares and freezes the target's `embed_tokens` and
`lm_head`, training only the draft stack, its feature projection, and its norms.

It follows the same scaffolding as the EAGLE and DSpark recipes: online target
hidden-state capture, gradient accumulation, and consolidated-safetensors
checkpointing.

## Variants

All four share the DFlash backbone, anchor sampling, and block attention mask;
only the head and the objective differ.

| Variant  | What it adds                                                                                                                                              | Recipe               | Example config       |
| -------- | --------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | -------------------- |
| DFlash   | the block-parallel baseline                                                                                                                               | `TrainDFlashRecipe`  | `qwen3_dflash.yaml`  |
| DFlash 2 | a two-tap in-block convolution around every sublayer, plus a pairwise path selector that walks one coherent path through each position's top-k candidates | `TrainDFlash2Recipe` | `qwen3_dflash2.yaml` |
| Domino   | a serial GRU correction head that refines each block position on the previous ones                                                                        | `TrainDominoRecipe`  | `qwen3_domino.yaml`  |
| JetSpec  | causal in-block attention plus forward-KL distillation against the target distribution                                                                    | `TrainJetSpecRecipe` | `jetspec/`           |

DFlash 2 keeps drafting in one pass. The convolution and the selector add
parameters inside the stack rather than an autoregressive correction, so it has
its own draft class whose checkpoint layout matches the published DFlash 2
drafters.

## Objective

Only block positions `1..block_size-1` are supervised (position 0 is the clean
anchor). Their cross-entropy is weighted by a position decay
`w_k = exp(-(k-1)/loss_decay_gamma)`, so early positions (the ones a verifier
reaches first) dominate the loss.

Paper `loss_decay_gamma` values track the block size: 7 for `block_size` 16, 5
for 10, and 4 for 8.

Set `loss_type: variable_prefix` to train the D2SD VP-Drafter objective
([arXiv:2606.04446](https://arxiv.org/abs/2606.04446)) instead: each block draws
a visible-prefix length from a truncated geometric prior
(`prefix_weight_base`, default 0.9), only the masked suffix is supervised, and
the decay restarts at the prefix boundary. This matches the regime a drafter
sees when it re-drafts behind a partially accepted block.

DFlash 2 adds one term on top of the block CE, weighted by
`selector_loss_weight`: a cross-entropy over the position's top-k candidates,
scored against the ground-truth predecessor. Positions whose true token missed
the candidate list carry no selector signal and are reported as
`candidate_recall`.

## Data

Use a chat dataset of OpenAI-format `messages` rows. As with the other
speculative recipes, use Open-PerfectBlend prompts with the **responses
regenerated by the target model** (training is teacher-forced; regenerate before
training to avoid a train/inference distribution mismatch). Point
`recipe_args.train_data_path` at the regenerated JSONL or a Hugging Face dataset
id.

> **Note**
>
> `mask_token_id` has no default and must name a reserved, rarely-used token. Do
> not reuse `pad`: it is frequently aliased to `eos`, which conflates the mask
> signal with real content and quietly erodes acceptance. The inference runtime
> must fill block slots with the same id.

## Run It

Example configs live under `examples/speculative/dflash/`.

```bash
torchrun --standalone --nproc_per_node=2 \
  -m nemo_automodel.recipes.llm.train_dflash \
  -c examples/speculative/dflash/qwen3_dflash.yaml
```

Swap the module and config to train a variant. For example,
`-m nemo_automodel.recipes.llm.train_dflash2 -c examples/speculative/dflash/qwen3_dflash2.yaml`.

Multi-GPU wraps the draft in DDP and replicates the frozen target. Set
`distributed.tp_size` to shard a large target across ranks
(`qwen3_dflash_tp.yaml`); `distributed.cp_size` enables context parallelism,
which cannot be combined with sequence packing or with `tp_size > 1`.

The published drafters diverge from their target in two ways the recipe
reproduces: a sliding window over the context (`draft_sliding_window`, which also
keeps the flex `BlockMask` sparse on long sequences) and a different attention
shape (`draft_num_attention_heads` and friends). See `qwen3_8_27b_dflash2.yaml`.

Training logs `train/loss`, `train/accuracy`, and `train/accept_len`. Setting
`recipe_args.val_data_path` adds globally reduced `val_loss`, `val_accuracy`,
and `val_accept_len`; adding a top-level `wandb:` block uploads all of them.

## Key Config Fields

| Field                                                                        | Meaning                                                                                  |
| ---------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------- |
| `target_model_name_or_path`                                                  | frozen target (e.g., `Qwen/Qwen3-8B`)                                                    |
| `draft_num_hidden_layers`                                                    | draft stack depth (paper: 5)                                                             |
| `block_size`                                                                 | tokens drafted in parallel per block (paper: 16)                                         |
| `num_anchors`                                                                | blocks sampled per sequence per step                                                     |
| `loss_decay_gamma`                                                           | block-position decay; `null` disables it                                                 |
| `loss_type`                                                                  | `dflash` (fixed anchor) or `variable_prefix`                                             |
| `target_layer_ids`                                                           | target feature layers fed to the draft (defaults to an even spread)                      |
| `mask_token_id`                                                              | reserved token id filling non-anchor block positions (required)                          |
| `attention_backend`                                                          | `flex_attention` (GPU main path) or `sdpa` (portable fallback)                           |
| `packed_sequence_size`                                                       | enables sequence packing when greater than 0                                             |
| `draft_sliding_window`                                                       | bounds how far back a block reads the context; unset attends over the whole prefix       |
| `draft_num_attention_heads` / `draft_num_key_value_heads` / `draft_head_dim` | size the draft's attention independently of the target's                                 |
| `output_multiplier` / `final_logit_softcapping`                              | the target's logit transform, applied in both training and decoding; identity unless set |
| `input_embedding_scale`                                                      | scales the embedded draft block; identity unless set                                     |
| `conv_kernel_size` / `conv_group_size`                                       | DFlash 2 convolution taps and channels per dynamic correction                            |
| `selector_rank` / `selector_top_k` / `selector_loss_weight`                  | DFlash 2 path-selector width, candidates scored per position, and objective weight       |

`spec_generate` accepts `temperature`, `top_p`, and `top_k`. The truncation
applies to the *target's* distribution (the one rejection sampling verifies
against), so sampled decoding stays lossless with respect to whatever sampling
the target would have used on its own. The DFlash 2 evaluations quote acceptance
lengths at temperature 1.0, top-p 0.95, top-k 20.

Supported targets: Qwen3 and Qwen3.5 (dense and MoE), including the multimodal
`*ForConditionalGeneration` variants such as `Qwen/Qwen3.8-27B`. DFlash and
JetSpec drafts serve on vLLM through its `dflash` method.