DSpark Speculative Decoding Training

This example shows how to train a DSpark (semi-autoregressive speculative decoding) draft model alongside the policy model during RL. DSpark drafts multiple tokens in parallel from intermediate policy hidden states, then vLLM verifies them in a single forward pass — accelerating rollouts without a separate draft training pipeline.

Key Features

• Online draft training: The draft model is trained jointly with the policy during RL, so it stays in distribution as the policy updates.

• Weight sync to vLLM: Draft weights are synced to vLLM every rollout step via direct IPC (colocate) or packed tensor transfer (non-colocate), enabling immediate speculative decoding in the next rollout.

• Freeze-policy mode: Optionally freeze the policy and train only the draft model, useful when the RL signal is weak or policy degradation is a concern.

Prerequisites

1. Pre-trained DSpark draft checkpoint (recommended). Training from random init converges slowly. Pre-train the draft backbone on a supervised corpus first, then pass the checkpoint via --dspark-pretrained-model.

2. Policy model in both HF and torch_dist formats:

cd /root/vime
source scripts/models/qwen3-4B.sh

PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
    ${MODEL_ARGS[@]} \
    --hf-checkpoint /root/Qwen3-4B \
    --save /root/Qwen3-4B_torch_dist

3. Dataset (e.g., dapo-math-17k):

hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/dapo-math-17k

Key Arguments

Argument	Description
--dspark-pretrained-model	Path to pre-trained DSpark safetensors. Strongly recommended.
--dspark-block-size	Number of draft tokens per block (default: 7).
--dspark-num-draft-layers	Number of decoder layers in the draft backbone (default: 5).
--dspark-target-layer-ids	Comma-separated policy layer indices to capture hidden states from (default: “1,9,17,25,33”).
--dspark-freeze-policy	Freeze policy and train only the draft model.
--dspark-ce-loss-alpha	Weight for cross-entropy loss (default: 0.1).
--dspark-l1-loss-alpha	Weight for L1/TV loss (default: 0.9).
--dspark-draft-loss-weight	Weight multiplying draft loss added to policy loss (default: 1.0).
--vllm-speculative-config	JSON config passed to vLLM. Setting method: "dspark" also enables DSpark draft training.

Mode Comparison

Mode	Flag	GPU Layout	Weight Sync
Colocate	--colocate	Train + rollout share the same GPUs	Direct IPC (fastest)
Non-colocate	(default)	Train and rollout on separate GPU sets	Packed tensor transfer

Running the Example

Colocate Mode (Recommended)

Train and rollout share the same 8 GPUs. The policy model is offloaded to CPU during rollout, then restored for the next training step.

bash examples/dspark/run-qwen3-4B-dspark-colocate.sh

GPU layout:

GPUs	Role
0–7	Policy Megatron train + vLLM rollout (colocate)

Non-Colocate Mode

Train and rollout run on separate GPU groups. This avoids the offload overhead but requires more GPUs.

bash examples/dspark/run-qwen3-4B-dspark-non-colocate.sh

GPU layout:

GPUs	Role
0–3	Policy Megatron train
4–7	vLLM rollout (with DSpark speculative decoding)

What to Expect

On Qwen3-4B with 8x A800 GPUs and a pre-trained DSpark draft checkpoint:

Metric	Typical Value
Draft acceptance rate	30–40%
Mean acceptance length	3.2–3.8
Rollout speedup	~2x vs no speculative decoding
Weight sync time	~10s per step

FAQ

1. Do I need a pre-trained draft model? Strongly recommended. Training from random init requires many more steps to converge. Pre-train the draft backbone on a supervised corpus, then pass the checkpoint via --dspark-pretrained-model.

2. What does --dspark-freeze-policy do? It freezes the policy model and trains only the draft model. The policy logits are detached so gradients only flow to the draft. Use this when the RL signal is weak or when you want to improve speculative decoding without affecting the policy.

3. How are draft weights synced to vLLM? Through vLLM’s standard draft weight-update session: colocated engines use IPC and non-colocated engines use NCCL.

4. What is --dspark-block-size? The number of tokens the draft model predicts in parallel per block. Larger values increase potential speedup but may reduce acceptance rate. The default (7) works well for most models.

5. How do I choose --dspark-target-layer-ids? These are the policy layer indices from which hidden states are captured as input to the draft model. For a 36-layer model, "1,9,17,25,33" samples every 8th layer. More target layers = richer draft input but higher cost.

References

1. DSpark Paper — Semi-autoregressive speculative decoding.

2. vLLM Speculative Decoding Docs

3. vime Speculative Decoding Docs