Speculative Decoding Draft Co-Training on Automodel Backend#

This guide covers draft-model co-training on the Automodel backend: training a speculative-decoding drafter alongside the policy so it tracks the policy’s drift during RL, then refitting both into vLLM for faster rollout generation.

This is the Automodel counterpart to the Eagle3 Speculative Decoding guide, which covers the Megatron backend. Concepts (teacher signal, refit, why co-train at all) are shared between the two guides; this one is specific to the Automodel drafters and their config surface. For GRPO fundamentals, see the GRPO guide.

Supported Drafters#

`policy.draft.speculator_type`	Drafter family	Checkpoint examples	Requires
`"dspark"`	DSpark block drafter (multi-token block proposal)	`deepseek-ai/dspark_qwen3_8b_block7`	`policy.automodel_cfg.enabled=true`
`"dflash"`	DFlash block drafter (markov-free, confidence-free variant of DSpark)	`RedHatAI/*-speculator.dflash`	same as `dspark`
`"eagle3"`	EAGLE3 TTT (test-time-training) drafter	`RedHatAI/-speculator.eagle3`, SGLang SpecForge native-flat checkpoints (e.g. `lmsys/SGLang-EAGLE3-`)	Automodel (this guide) or Megatron (see the Eagle3 guide)

All three algorithms train alongside the policy on the same rollout batch, using the policy’s own in-flight logits/hidden states as the distillation teacher (not a separate frozen verifier), and stream draft.* weights into vLLM on every refit so the next rollout speculates with the freshest drafter.

Get Started#

Usage#

Draft co-training requires the Automodel backend:

policy:
  automodel_cfg:
    enabled: true

DSpark#

policy:
  draft:
    enabled: true
    model_name: deepseek-ai/dspark_qwen3_8b_block7
    speculator_type: "dspark"
    loss_weight: 1.0
    num_anchors: 32
    learning_rate: 1.0e-04
    ce_loss_alpha: 0.1
    l1_loss_alpha: 0.9
    confidence_loss_alpha: 1.0
    loss_decay_gamma: 4.0
    train_embed_and_head: true
  generation:
    vllm_kwargs:
      speculative_config:
        method: "dspark"
        model: ${policy.draft.model_name}
        num_speculative_tokens: 7
        attention_backend: "FLASH_ATTN"
        draft_sample_method: "probabilistic"

DFlash#

Same fields as DSpark, but DFlash checkpoints have no confidence head, so confidence_loss_alpha defaults to 0.0 for speculator_type: "dflash":

policy:
  draft:
    enabled: true
    model_name: RedHatAI/Qwen3-8B-speculator.dflash
    speculator_type: "dflash"
    loss_weight: 1.0
    num_anchors: 32
    learning_rate: 1.0e-04
    ce_loss_alpha: 0.1
    l1_loss_alpha: 0.9
    loss_decay_gamma: 4.0
    train_embed_and_head: true
  generation:
    vllm_kwargs:
      speculative_config:
        method: "dflash"
        model: ${policy.draft.model_name}
        num_speculative_tokens: 7
        attention_backend: "FLASH_ATTN"
        draft_sample_method: "probabilistic"

EAGLE3#

policy:
  draft:
    enabled: true
    model_name: RedHatAI/Qwen3-8B-speculator.eagle3
    loss_weight: 1.0
    speculator_type: "eagle3"
    learning_rate: 1.0e-04
    ttt_steps: 3
    ttt_step_loss_decay: 1.0
    train_embed_and_head: true
  generation:
    vllm_kwargs:
      speculative_config:
        method: "eagle3"
        model: ${policy.draft.model_name}
        num_speculative_tokens: 3

speculator_type defaults to "eagle3" when omitted, so it can be left out (as in the snippet above) unless you want it explicit for readability.

Config Reference#

policy.draft is a discriminated union keyed on speculator_type, with each family’s fields declared directly on policy.draft (no nested sub-blocks).

Shared across all three:

Field	Description
`enabled`	Attach and train a draft model alongside the policy.
`model_name`	Pretrained draft checkpoint to start from. Required – from-scratch draft init is not supported.
`speculator_type`	`"eagle3"` (default), `"dspark"`, or `"dflash"`.
`loss_weight`	Weight on the auxiliary draft loss term.

speculator_type: "dspark" / "dflash":

Field	Default	Description
`num_anchors`	`64`	Anchor blocks sampled per sequence each training forward. Draft-side transient memory scales with `num_anchors * block_size * vocab`; co-training shares the GPU with the full policy, so raise this only after confirming headroom (the shipped Qwen3-8B dspark/dflash recipes use `32`).
`learning_rate`	`1e-4`	Draft param-group learning rate – needs to be well above the policy’s RL learning rate to track policy drift.
`ce_loss_alpha`	`0.1`	Cross-entropy weight against the rollout tokens.
`l1_loss_alpha`	`0.9`	Total-variation distillation weight against the policy’s raw logits.
`confidence_loss_alpha`	`1.0` (dspark) / `0.0` (dflash)	Confidence-head BCE weight; `dflash` has no confidence head, so its default is `0.0`.
`loss_decay_gamma`	`4.0`	Exponential per-block-position decay on the loss mask.
`train_embed_and_head`	`true`	Train the draft’s `embed_tokens`/`lm_head` instead of keeping the checkpoint copies frozen. Megatron backends handle embedding/lm_head ownership differently; see NVIDIA-NeMo/RL#4691.

speculator_type: "eagle3":

Field	Default	Description
`learning_rate`	`1e-4`	Draft param-group learning rate.
`ttt_steps`	`3`	Test-time-training unroll depth.
`ttt_step_loss_decay`	`1.0`	Per-unroll-step loss decay factor.
`train_embed_and_head`	`true`	Train the draft’s `embed_tokens`/`lm_head`.

Architecture fields (block size, target layer ids, mask token id, aux capture layers, draft_vocab_size, d2t/t2d vocab maps, markov/confidence head layout, norm_before_residual) are always read from the draft checkpoint’s config.json – only training options are set in the recipe.

Limitations#

Not compatible with policy.automodel_cfg.sequence_parallel, policy.automodel_cfg.lora_cfg.enabled, or policy.sequence_packing.enabled.

How It Works#

Speculative decoding needs a small drafter that stays well-aligned with the serving policy to get useful acceptance rates. A drafter trained once and never updated drifts out of sync as RL moves the policy’s distribution, degrading acceptance length over training. Co-training closes that loop: the drafter is trained every step on the same data the policy trains on, and refit alongside the policy so rollout always speculates with an up-to-date drafter.

rollout (policy samples responses)
  -> policy training forward pass
       -> hidden states / logits captured mid-forward (detached)
  -> draft forward: cross-attends captured hidden states,
     distills against the captured policy logits
  -> combined loss: policy_loss + loss_weight * draft_loss
  -> backward, optimizer.step (policy and draft param groups,
     each with its own learning rate)
  -> refit: policy.* AND draft.* weights sync to the vLLM
     generation workers together
  -> next rollout speculates with the updated drafter

The policy loss path is unaffected – draft co-training only adds a second term and captures already-computed intermediate policy tensors along the way.

Automodel Reuse Boundary#

nemo_rl/models/automodel/draft/ is a thin extension layer over nemo_automodel.components.speculative.* (Automodel r0.6.0). Attention masks and the unchanged helpers are imported from there; the DSpark/DFlash model and loss are local copies whose RL deltas are listed in each file’s header, and EAGLE3’s per-position KL, loss mask and accuracy counting follow speculators’ kl_div_loss / align_for_step. EAGLE3’s network (attention, decoder layer, embedding, fc, lm_head) is a direct subclass of Automodel’s LlamaEagle3DraftModel; NeMo RL does not reimplement the attention math.

What stays local to NeMo RL, and why:

RL-shaped teacher signal and normalization. The teacher is the policy’s own live logits/hidden states (not a frozen verifier), anchor sampling respects RL rollout response boundaries instead of pretraining-corpus document boundaries, and loss normalization matches NeMo RL’s gradient-accumulation/context-parallel conventions.
Training-loop and distributed-topology glue (draft/{setup,hidden_capture, runtime,checkpoint}.py): hooking into the policy’s own forward pass to capture teacher hidden states/logits, keeping gradient scaling consistent across microbatches and context parallelism, deterministic anchor sampling across TP replicas, coupled policy+draft checkpointing, and adapting third-party checkpoint formats (speculators, SGLang SpecForge).
Wiring the RL main loop (nemo_rl/models/automodel/{setup,train}.py, nemo_rl/algorithms/loss/wrapper.py, nemo_rl/models/policy/{lm_policy,workers/automodel_policy_worker}.py): building the draft optimizer param group, capturing teacher logits at the right point in the forward pass, combining the loss terms, config validation, and streaming draft.* weights through the same refit path as the policy weights.

None of this is available in Automodel’s own Eagle3TrainerModule, which is built for standalone offline SFT of a drafter, not RL co-training.

Evidence#

Recipe Coverage#

Recipes for the Qwen3-8B control arm plus its three drafter siblings, and a Qwen3-30B-A3B-2507-Instruct (EP32) EAGLE3 example, live under examples/configs/recipes/llm/:

Recipe	Drafter
`grpo-qwen3-8b-4n8g-automodel.yaml`	none
`grpo-qwen3-8b-4n8g-automodel-dspark.yaml`	DSpark
`grpo-qwen3-8b-4n8g-automodel-dflash.yaml`	DFlash
`grpo-qwen3-8b-4n8g-automodel-eagle3.yaml`	EAGLE3
`grpo-qwen3-30ba3b-2507-instruct-4n8g-automodel.yaml`	none
`grpo-qwen3-30ba3b-2507-instruct-4n8g-automodel-eagle3.yaml`	EAGLE3

The Qwen3-8B dspark/dflash/eagle3 recipes run nightly (tests/test_suites/llm/grpo-qwen3-8b-4n8g-automodel-{dspark,dflash,eagle3}.sh, registered in tests/test_suites/nightly.txt).

Results#

Qwen3-8B: dspark vs. dflash vs. eagle3 vs. no-draft baseline#

Qwen3-8B draft co-training curves

train/reward and validation/accuracy are basically the same across all four runs, so training a draft alongside the policy isn’t costing any policy quality here. The baseline has no drafter, so its spec_acceptance_length just sits at 1; dspark and dflash land around 3.5-4, eagle3 around 2. The speedup shows up in timing/train/generation: the more of the end-to-end step time that rollout generation accounts for, the bigger the win from speculative decoding.

Qwen3-30B-A3B-2507-Instruct (MoE, EP32): eagle3 vs. no-draft baseline#

Qwen3-30B-A3B-2507-Instruct draft co-training curves

Same story on the bigger MoE model. Reward and validation accuracy track the baseline, eagle3’s acceptance length settles around 2.3-2.7 against the baseline’s flat 1, and generation time is lower for the eagle3 run – same pattern, the speedup comes from generation, and matters more the longer the response gets.