> This page is for version 0.7.0 · 26.10.
> 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.

# nemo_automodel.recipes.speculative_decoding.train_dflash

DFlash draft-model training recipe (Qwen3-style and Kimi K3 targets).

DFlash drafts a whole block of tokens in parallel via MASK-token denoising
conditioned on the frozen target's hidden states (see
`nemo_automodel.components.speculative.dflash`). This recipe mirrors the EAGLE
recipes' scaffolding -- online target hidden-state capture, gradient
accumulation with a trailing-window flush, and the same checkpointer plumbing --
but trains the DFlash draft with its block-wise cross-entropy objective.

## Module Contents

### Classes

| Name                                                                                               | Description                                                                            |
| -------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------- |
| [`TrainDFlashRecipe`](#nemo_automodel-recipes-speculative_decoding-train_dflash-TrainDFlashRecipe) | Recipe for DFlash draft-model training on Qwen3-style dense / MoE and Kimi K3 targets. |

### Functions

| Name                                                                                                                           | Description                                                                                        |
| ------------------------------------------------------------------------------------------------------------------------------ | -------------------------------------------------------------------------------------------------- |
| [`_all_ranks_have_valid`](#nemo_automodel-recipes-speculative_decoding-train_dflash-_all_ranks_have_valid)                     | Min-reduce a per-rank "this micro-batch has valid anchors" flag.                                   |
| [`_average_unsynced_ddp_grads`](#nemo_automodel-recipes-speculative_decoding-train_dflash-_average_unsynced_ddp_grads)         | Average locally accumulated grads over the DDP group, as DDP's synced backward would (sum / size). |
| [`_project_onto_qwen3_config_keys`](#nemo_automodel-recipes-speculative_decoding-train_dflash-_project_onto_qwen3_config_keys) | Project the target's decoder config onto the keys a plain Qwen3 config declares.                   |
| [`main`](#nemo_automodel-recipes-speculative_decoding-train_dflash-main)                                                       | Entrypoint for `TrainDFlashRecipe`.                                                                |

### Data

[`logger`](#nemo_automodel-recipes-speculative_decoding-train_dflash-logger)

### API

```python
class nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe(
    cfg
)
```

**Bases:** [SpecDecodeRecipeBase](/nemo/automodel/nemo-automodel/nemo_automodel/recipes/speculative_decoding/_spec_recipe_base#nemo_automodel-recipes-speculative_decoding-_spec_recipe_base-SpecDecodeRecipeBase)

Recipe for DFlash draft-model training on Qwen3-style dense / MoE and Kimi K3 targets.

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe._build_dflash_config(
    recipe_cfg,
    target_layer_ids: list[int]
) -> dict
```

Build the draft `dflash_config` block. Subclasses extend it (e.g. Domino).

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe._build_qwen3_draft_config(
    recipe_cfg,
    target_text_config,
    draft_cls,
    draft_num_hidden_layers: int,
    num_target_layers: int,
    target_layer_ids: list[int],
    attention_backend: str
) -> transformers.models.qwen3.configuration_qwen3.Qwen3Config
```

Derive the draft config for a Qwen3-shaped target.

A small non-causal Qwen3 stack that reuses the target's architecture
defaults (head_dim, rope_theta, rms_norm_eps, ...). Targets whose draft is
not Qwen3-shaped register their own builder on the spec instead and never
reach this.

**Parameters:**

**`recipe_cfg`**

The recipe's `recipe_args` mapping.

---

**`target_text_config`**

The target's decoder config.

---

**`draft_cls`**

The draft class being built, stamped into `architectures`.

---

**`draft_num_hidden_layers`** `int`

Depth of the draft stack.

---

**`num_target_layers`** `int`

Depth of the target, for the `fc` input width.

---

**`target_layer_ids`** `list[int]`

Target layers captured as draft context.

---

**`attention_backend`** `str`

The draft's attention implementation.

---

**Returns:** `Qwen3Config`

The draft `Qwen3Config`.

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe._build_target_model(
    recipe_cfg,
    target_path: str,
    draft_spec: nemo_automodel.components.speculative.dflash.registry.DFlashDraftSpec
) -> torch.nn.Module
```

Load the frozen (optionally tensor-parallel) target model.

`draft_spec.build_target_kwargs` supplies any architecture-specific
`from_pretrained` arguments (Kimi K3 pins the text-only architecture and
an expert-parallel backend); it is empty for a Qwen3-shaped target.

With a `distributed:` section and `tp_size>1` the target is sharded
in place by `from_pretrained` (its FSDP2 parallelize plan); the small
draft stays replicated and runs DDP over the "dp" axis (which excludes
"tp"), and the trainer module gathers the target's vocab-sharded lm_head
/ embed_tokens outputs. Absent, the original single-GPU-per-rank DP path
is used. Sets `self.dist_setup` / `self.device_mesh` / `self.dp_mesh`
as a side effect and returns the (grad-disabled) target.

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe._build_target_wrapper(
    target_layer_ids: list[int]
) -> nemo_automodel.components.speculative.dflash.target.HFDFlashTargetModel
```

Build the frozen-target hidden-state capture wrapper.

Subclasses override to capture extra teacher signals (e.g. JetSpec also
captures the target logits for its forward-KL distillation).

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe._build_trainer_module(
    attention_backend: str,
    recipe_cfg
)
```

Build the trainer wrapper. Subclasses override to swap the wrapper (e.g. Domino).

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe._draft_cls(
    draft_spec: nemo_automodel.components.speculative.dflash.registry.DFlashDraftSpec
) -> type[torch.nn.Module]
```

Pick the draft class from the resolved spec.

Subclasses override to select a different draft of the same family; the
DFlash 2 recipe returns `draft_spec.draft2_cls`. The returned class name
is also what lands in the saved config's `architectures`, which is how a
serving engine tells the two drafts apart.

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe._draft_ddp_process_group()
```

Process group for the draft's gradient all-reduce.

With tensor parallelism the draft is replicated across tp ranks, so a
full-world all-reduce would average duplicate gradients; restrict it to
the "dp" sub-axis (which excludes tp) so it reduces only across real data
replicas. Without a mesh (tp_size=1) `dp_mesh` is None -> return None ->
the default full-world group, unchanged.

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe._empty_extra_eval_metric_sums() -> dict[str, list[torch.Tensor]]
```

Create zeroed subclass validation accumulators on the trainer device.

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe._extra_eval_metric_sums(
    metrics
) -> dict[str, tuple[torch.Tensor, torch.Tensor]]
```

Return additional validation numerator and denominator pairs.

The base DFlash metrics are accumulated directly by `_run_eval`.
Subclasses use this hook for extra scalar statistics, with both tensors
on the same device as `metrics.loss` so they can participate in the
same ordered distributed SUM reductions.

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe._extra_train_metric_sums(
    metrics
) -> dict[str, tuple[float, float]]
```

Return algorithm-specific training numerator and denominator pairs.

These are accumulated over the micro-batches between two log points and
divided at the log point, the same way `train/loss` and
`train/accuracy` are, so every curve on the dashboard covers the same
window. Returning the per-micro-batch mean instead would report a single
micro-batch out of `log_every_steps * grad_accumulation_steps`.

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe._load_extra_state(
    ckpt_dir: str
) -> None
```

Restore DFlash meta: global_step and epoch, and validate mask_token_id.

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe._log_extra_train_metrics(
    epoch_idx: int
) -> None
```

Hook for subclasses to log extra per-step metrics at a log point (no-op here).

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe._resolve_mask_token_id(
    recipe_cfg,
    vocab_size: int
) -> int
```

staticmethod

Resolve and validate the MASK token id that fills non-anchor block positions.

DFlash fills every non-anchor slot of a `[anchor, MASK, MASK, ...]` block
with this id, and the draft's `embed_tokens` row at that id becomes the
learned "predict here" signal. It must be chosen deliberately (a reserved /
unused token), exactly like P-EAGLE's `mask_token_id`: the previous silent
fallback to `tokenizer.pad_token_id` was unsafe because `pad` is commonly
aliased to `eos` (or another meaningful token), which conflates the mask
signal with real content and quietly degrades acceptance without erroring.
Require it explicitly and range-check it; the inference runtime must fill the
block slots with the same id.

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe._run_eval()
```

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe._run_trainer_step(
    target_batch
)
```

Run one trainer-module forward. Subclasses override to inject extra inputs (e.g. lambda_base).

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe._save_extra_state(
    path: str,
    epoch: int
) -> None
```

Persist DFlash meta: global_step, epoch, block_size, and target layers.

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe.run_train_validation_loop()
```

Run the DFlash training loop.

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.TrainDFlashRecipe.setup()
```

Build the target model, DFlash draft, data, optimizer, and trainer module.

```python
nemo_automodel.recipes.speculative_decoding.train_dflash._all_ranks_have_valid(
    local_has_valid: int,
    is_ddp: bool,
    device
) -> bool
```

Min-reduce a per-rank "this micro-batch has valid anchors" flag.

Under DDP a data-dependent `NoValidAnchorsError` skip is per-rank: if one
rank skips its backward (and its gradient all-reduce) while another runs its,
the collective mismatches (hang) and the accumulation windows desync. Taking
the MIN across ranks makes the skip decision unanimous -- every rank skips the
micro-batch unless all of them have something to learn from it. The reduce is
a tiny independent collective, safe inside `no_sync` (which only gates the
DDP backward all-reduce). Single-process runs return the local flag unchanged.

```python
nemo_automodel.recipes.speculative_decoding.train_dflash._average_unsynced_ddp_grads(
    ddp_module: torch.nn.parallel.DistributedDataParallel
) -> None
```

Average locally accumulated grads over the DDP group, as DDP's synced backward would (sum / size).

```python
nemo_automodel.recipes.speculative_decoding.train_dflash._project_onto_qwen3_config_keys(
    target_text_config: dict
) -> dict
```

Project the target's decoder config onto the keys a plain Qwen3 config declares.

The draft is always a Qwen3-shaped stack, but its config starts from the
target's, and a Qwen3.5 text config carries fields the draft has no use for:
linear-attention shapes, MTP, output gating, and `partial_rotary_factor`
(top-level and inside `rope_parameters`, on its own key or as mRoPE
sections). None of them may reach the saved draft config -- the published
drafters ship without them, and they are not inert there: the HF Qwen3 stack
the draft trains with applies full rotary regardless, so a serving runtime
that honours a leaked `partial_rotary_factor: 0.25` would rebuild the
rotary table at a quarter width and silently mismatch the trained weights.

**Parameters:**

**`target_text_config`** `dict`

`to_dict()` of the target's decoder config.

---

**Returns:** `dict`

The subset of `target_text_config` a `Qwen3Config` declares, with

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.main(
    config_path: str | None = None
)
```

Entrypoint for `TrainDFlashRecipe`.

```python
nemo_automodel.recipes.speculative_decoding.train_dflash.logger = logging.getLogger(__name__)
```