> 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. `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 | | `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 | 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.