> This page is for version Nightly (default).
> 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.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.

# nemo_automodel.recipes.llm.train_dflash2

DFlash 2 draft-model training recipe (Qwen3-style targets).

DFlash 2 ([https://inco.ai/blog/dflash2/](https://inco.ai/blog/dflash2/)) keeps DFlash's one-pass parallel block
draft and adds two cheap modules: a two-tap dynamic convolution around every
draft sublayer, which carries the short-range within-block work and removes most
of DFlash's suffix decay, and a pairwise path selector that walks one coherent
path through each position's top-k candidates instead of keeping every top-1 pick
independently. See `nemo_automodel.components.speculative.dflash.dflash2_core`.

This recipe reuses every piece of the DFlash recipe -- online target hidden-state
capture, anchor sampling, the block attention mask, gradient accumulation, and
checkpointing -- and only swaps in the DFlash 2 draft class, stamps the extra
`dflash_config` fields a serving engine needs to rebuild it, and swaps in the
DFlash 2 trainer wrapper.

IMPORTANT: as with DFlash, regenerate the training responses with the target
model first -- training teacher-forces ground-truth tokens while inference is
autoregressive, and the distribution mismatch hurts acceptance length otherwise.

## Module Contents

### Classes

| Name                                                                                 | Description                                                                      |
| ------------------------------------------------------------------------------------ | -------------------------------------------------------------------------------- |
| [`TrainDFlash2Recipe`](#nemo_automodel-recipes-llm-train_dflash2-TrainDFlash2Recipe) | Recipe for DFlash 2 draft-model training: in-block convolutions + path selector. |

### Functions

| Name                                                     | Description                          |
| -------------------------------------------------------- | ------------------------------------ |
| [`main`](#nemo_automodel-recipes-llm-train_dflash2-main) | Entrypoint for `TrainDFlash2Recipe`. |

### Data

[`logger`](#nemo_automodel-recipes-llm-train_dflash2-logger)

### API

```python
class nemo_automodel.recipes.llm.train_dflash2.TrainDFlash2Recipe()
```

**Bases:** [TrainDFlashRecipe](/nemo-automodel/nemo_automodel/recipes/llm/train_dflash#nemo_automodel-recipes-llm-train_dflash-TrainDFlashRecipe)

Recipe for DFlash 2 draft-model training: in-block convolutions + path selector.

```python
nemo_automodel.recipes.llm.train_dflash2.TrainDFlash2Recipe._build_dflash_config(
    recipe_cfg,
    target_layer_ids: list[int]
) -> dict
```

Extend the DFlash draft config with the convolution and selector shapes.

```python
nemo_automodel.recipes.llm.train_dflash2.TrainDFlash2Recipe._build_trainer_module(
    attention_backend: str,
    recipe_cfg
)
```

Build the DFlash 2 trainer wrapper (block CE + candidate-selection CE).

```python
nemo_automodel.recipes.llm.train_dflash2.TrainDFlash2Recipe._draft_cls(
    draft_spec: nemo_automodel.components.speculative.dflash.registry.DFlashDraftSpec
) -> type[torch.nn.Module]
```

Build the DFlash 2 draft instead of the plain DFlash one.

```python
nemo_automodel.recipes.llm.train_dflash2.TrainDFlash2Recipe._empty_extra_eval_metric_sums() -> dict[str, list[torch.Tensor]]
```

Create rank-symmetric DFlash 2 validation accumulators.

```python
nemo_automodel.recipes.llm.train_dflash2.TrainDFlash2Recipe._extra_eval_metric_sums(
    metrics
) -> dict[str, tuple[torch.Tensor, torch.Tensor]]
```

Return additive DFlash 2 validation statistics.

Every returned value is a pair of scalar tensors on the trainer device.
The shared validation loop SUM-reduces each pair before division.

```python
nemo_automodel.recipes.llm.train_dflash2.TrainDFlash2Recipe._extra_train_metric_sums(
    metrics
) -> dict[str, tuple[float, float]]
```

Return the selector and backbone diagnostics as window sums.

The two loss terms are per-micro-batch means, so their denominator is the
micro-batch count; the accuracy-like curves are token- or block-weighted,
matching how `train/loss` and `train/accuracy` are averaged.

```python
nemo_automodel.recipes.llm.train_dflash2.TrainDFlash2Recipe._log_extra_train_metrics(
    epoch_idx: int
) -> None
```

Log the DFlash 2 diagnostics for the most recent step (rank-0 local).

```python
nemo_automodel.recipes.llm.train_dflash2.TrainDFlash2Recipe._run_trainer_step(
    target_batch
)
```

Forward through the DFlash 2 wrapper and cache the step's diagnostics.

```python
nemo_automodel.recipes.llm.train_dflash2.TrainDFlash2Recipe.setup()
```

Build everything via the DFlash recipe, then reset the per-step metric cache.

```python
nemo_automodel.recipes.llm.train_dflash2.main(
    config_path: str | None = None
)
```

Entrypoint for `TrainDFlash2Recipe`.

```python
nemo_automodel.recipes.llm.train_dflash2.logger = logging.getLogger(__name__)
```