source/docs/cli/train.md

train

Trains speculator models using either online or offline hidden states. Supports single-GPU and multi-GPU distributed training.

Basic Usage

Single-GPU:

speculators train \
  --verifier-name-or-path meta-llama/Llama-3.1-8B-Instruct \
  --data-path ./training_data \
  --save-path ./checkpoints \
  --draft-vocab-size 32000 \
  --epochs 10

Multi-GPU (DDP):

torchrun --standalone --nproc_per_node=4 -m speculators.train \
  --verifier-name-or-path meta-llama/Llama-3.1-8B-Instruct \
  --data-path ./training_data \
  --save-path ./checkpoints \
  --draft-vocab-size 32000 \
  --epochs 10

Multi-GPU (FSDP sharded):

torchrun --standalone --nproc_per_node=4 -m speculators.train \
  --verifier-name-or-path meta-llama/Llama-3.1-8B-Instruct \
  --data-path ./training_data \
  --save-path ./checkpoints \
  --draft-vocab-size 32000 \
  --epochs 10 \
  --fsdp-shard

Arguments

Model Arguments

• --verifier-name-or-path (str, required) HuggingFace model ID or local path for the verifier/target model.

• --trust-remote-code (flag) Allow executing code from HF Hub when loading the verifier's tokenizer.

• --speculator-type (str, default: "eagle3") Type of speculator model to train. Options: eagle3, dflash, dflash2, dspark, peagle, mtp

• --from-pretrained (str, default: "") Path or HF id of an existing draft checkpoint to load weights from and train — either a previously trained draft or the initialized-but-untrained checkpoint produced by --dry-run. May also point to a local directory containing only a config.json, in which case a fresh draft is initialized from that full speculator config. Takes precedence over all other model-definition options: it is mutually exclusive with --draft-config and the decoder-shaping flags (--num-layers, --draft-arch, --draft-hidden-act, --sliding-window, --full-attention-indices).

• --draft-config (str, default: "") HF id, directory, or JSON path of a decoder config (LlamaConfig for eagle3/peagle, Qwen3Config for DFlash-family models) used as the draft transformer_layer_config; the rest of the speculator is built from the other CLI args. The draft hidden_size must match the verifier (mismatch is not yet supported). If a full speculator config is passed, its nested transformer_layer_config is extracted. Mutually exclusive with --from-pretrained and with the decoder-shaping flags (--num-layers, --draft-arch, --draft-hidden-act, --sliding-window, --full-attention-indices).

• --dry-run (flag) Build the speculator, initialize weights, save a checkpoint to --save-path, then exit before training. Useful to validate the config/weights in vLLM before launching a full run; the saved checkpoint can be fed straight back via --from-pretrained.

• --num-layers (int, default: 5 for dflash/dspark/dflash2, 1 otherwise) Number of transformer layers in the draft model.

• --draft-arch (str, default: "llama") Architecture for the synthesized draft decoder layers. Options: llama, qwen3. Used by Eagle3 and P-EAGLE, which select the decoder layer class from this value; DFlash-family models always use a Qwen3-style decoder regardless. Both are supported in vLLM for inference, and the target and draft architectures do not have to match.

• --draft-hidden-act (str, default: "silu") Activation function for draft decoder layers. Setting as None will inherit activation function from the verifier model.

Data Arguments

• --data-path (str, default: "./data") Path to the processed training data directory.

• --on-missing (choice: generate|skip|warn|raise, default: generate) Behavior when cached hidden states are missing:

  • generate: Generate hidden states on-demand using vLLM endpoint
  • skip: Skip the sample silently, pads to fill batch.
  • warn: Skip the sample with a warning, pads to fill batch.
  • raise: Raise an error

• --on-generate (choice: cache|delete, default: "delete") Behavior after generating new hidden states (only applies if --on-missing=generate):

  • delete: Delete hidden states after loading (pure online training)
  • cache: Store hidden states for reuse in future epochs (hybrid training)

• --hidden-states-path (str, default: {data-path}/hidden_states) Path where cached hidden states files are stored (or will be stored if generating).

• --vllm-endpoint (str, default: "http://localhost:8000/v1") vLLM endpoint address for generating hidden states on-demand (online training). Ignored if --on-missing is not set to generate.

• --request-timeout (float, default: 180.0) Timeout in seconds for each individual vLLM request.

• --max-retries (int, default: 3) Maximum number of retry attempts per vLLM request on failure.

• --total-seq-len (int, default: 8192) Maximum total sequence length for training batches. Note: samples will be packed into batches with total combined sequence length {total-seq-len}.

Vocabulary Mapping Arguments

• --draft-vocab-size (int, default: None) Vocabulary size for the draft model. If not specified and no vocab mapping files are provided, uses full verifier vocabulary.

• --token-freq-path (str, default: {data-path}/token_freq.pt) Path to token frequency distribution file. This is used to determine which tokens to include in the reduced draft vocab.

• --d2t-path (str, default: None) Path to draft-to-target vocabulary mapping file (.npy). Must be provided with --t2d-path.

• --t2d-path (str, default: None) Path to target-to-draft vocabulary mapping file (.npy). Must be provided with --d2t-path.

• --mask-token-id (int, default: auto-detect) Token ID to use as mask token (for DFlash). Auto-detected if not provided.

• --target-layer-ids (int list, default: auto-select) Space-separated list of layer IDs for the auxiliary hidden states. Default: [2, num_layers//2, num_layers-3] If custom layers were specified when launching vLLM, pass the same ids here, excluding the final layer launch_vllm.py appends — that one reaches training separately as the verifier's last hidden states.

Distributed Training Arguments

• --fsdp-shard (flag) Shard model parameters across GPUs with FSDP. By default, parameters are fully replicated (DDP-like). Enable this when the model does not fit in a single GPU's memory.

Training Arguments

• --save-path (str, default: "./checkpoints") Directory to save model checkpoints.

• --epochs (int, default: 20) Number of training epochs.

• --lr (float, default: 1e-4) Learning rate.

• --train-data-ratio (float, default: 0.9) Ratio of data to use for training, the rest of the provided data will be used for validation.

• --no-resume-from-checkpoint (flag) Disable automatic checkpoint resumption. Without this flag, this script will automatically load the latest checkpoint in {save-path} if one exists.

• --logger (str, default: "") Metric logging backend(s). Options: trackio, wandb, tensorboard, mlflow Can specify multiple comma-separated: --logger tensorboard,wandb. Warning: backend must be pip installed before using.

• --log-dir (str, default: "./logs") Directory to save training logs. Only applies to some logging backends (e.g. tensorboard)

• --run-name (str, default: None) Name for the training run (used by logging backends).

• --seed (int, default: 42) Random seed for reproducibility.

• --hidden-states-dtype (str, default: "bfloat16") Data type for dataloader hidden states and autocast compute. Model master weights are always kept in fp32. Options: float32 (full precision, for debugging), bfloat16 (recommended for mixed precision training). Note: float16 is not supported as it requires gradient scaling to prevent underflow.

• --deterministic-cuda (flag) Enable deterministic CUDA operations. May impact performance.

• --loss-fn (str, default: "ce" for dflash, "kl_div" otherwise) Loss function specification. Pass a name for a single loss (kl_div, rkl, jsd, ce, tv, nla, lk_hybrid) or a JSON dict for a weighted combination, e.g. '{"ce": 0.1, "tv": 0.9}'. Required to be ce when --per-position-loss-weight dpace is used.

Optimizer Arguments

• --optimizer (str, default: "muon") Optimizer to use. Options: adamw, muon. The muon option applies the Muon optimizer to 2D weight matrices and AdamW to the remaining parameters (norms, biases, embeddings, lm_head).

• --weight-decay (float, default: 0.01) Weight decay for the AdamW optimizer (and the AdamW group in muon mode).

• --muon-lr (float, default: 10*lr) Learning rate for the Muon (2D weights) group. Only used with --optimizer muon. Defaults to 10× the --lr value.

• --muon-momentum (float, default: 0.95) Momentum for the Muon optimizer. Only used with --optimizer muon.

• --muon-weight-decay (float, default: 0.1) Weight decay for the Muon optimizer. Only used with --optimizer muon.

• --muon-ns-steps (int, default: 5) Number of Newton-Schulz steps for Muon. Only used with --optimizer muon.

• --muon-adjust-lr-fn (str, default: "match_rms_adamw") Muon LR adjustment strategy. Options: original, match_rms_adamw. Only used with --optimizer muon.

Eagle3-Specific Arguments

• --norm-before-residual / --no-norm-before-residual (flag, default: True) Toggle normalization before residual connections.

• --embed-requires-grad / --no-embed-requires-grad (flag, default: False) Whether to train embedding layer weights.

• --norm-before-fc / --no-norm-before-fc (flag, default: True for eagle3, False otherwise) Apply a single RMSNorm to the concatenated auxiliary hidden states before the FC projection (gpt-oss style). See --fc-norm for the per-layer alternative from the Eagle 3.1 paper.

• --fc-norm (flag, default: False) Apply per-layer RMSNorm to each auxiliary hidden state before concatenation and FC projection (Eagle 3.1 paper approach).

• --norm-output / --no-norm-output (flag, default: True for eagle3, False otherwise) Feed post-norm hidden states back across TTT steps to stabilize magnitude drift across speculation depths.

• --ttt-steps (int, default: 3) Number of test-time training steps

• --ttt-step-loss-decay (float, default: 1.0) Loss decay factor for test-time training steps.

P-EAGLE-Specific Arguments

• --num-depths (int, default: 8) Number of parallel prediction depths.

• --down-sample-ratio (float, default: 0.7) Geometric decay ratio for COD sampling.

• --down-sample-ratio-min (float, default: 0.2) Minimum retention ratio for COD sampling.

Attention Backend Arguments

• --draft-attn-impl (str, default: "simple_flex_attention") Attention implementation for draft layers. Options: simple_flex_attention, sdpa, eager. Use sdpa or eager on hardware where flex attention is unavailable (e.g. Ascend NPU). Applies to Eagle3, P-EAGLE, and DFlash-family models. Not supported for MTP.

DFlash-Specific Arguments

• --block-size (int, default: 16 for dflash, 8 otherwise) Block size for DFlash-family models.

• --sample-from-anchor / --no-sample-from-anchor (bool, default: algorithm-specific) Whether to sample from the anchor position. True: sample from anchor and all mask positions (default for dspark, produces block_size tokens). False: anchor is bonus token (default for dflash/dflash2, produces block_size-1 tokens).

• --max-anchors (int, default: 512) Maximum anchor positions for DFlash-family and P-EAGLE training.

• --dflash-decay-gamma (float, default: 4.0) Decay gamma for DFlash-family loss weighting.

• --per-position-loss-weight (str, default: "dpace" for dflash, "fixed-exp-decay" otherwise) Per-position loss weighting scheme. Options: fixed-exp-decay, dpace. Applies to DFlash-family models. dpace requires --loss-fn ce.

• --dpace-alpha (float, default: 0.5) Confidence smoothing constant for the D-PACE loss. Only used with --per-position-loss-weight dpace.

DFlash2-Specific Arguments

DFlash2 builds on DFlash, so all DFlash-specific arguments apply as well. It defaults to five draft layers, block size 8, and KL loss.

• --conv-kernel-size (int, default: 2) Local convolution kernel size.

• --conv-group-size (int, default: 16) Channel group size for local convolution.

• --selector-rank (int, default: 256) Low-rank dimension of the candidate selector.

• --selector-top-k (int, default: 16) Number of candidates retained per position.

• --selector-loss-alpha (float, default: 1.0) Weight of the candidate-selector K-way cross-entropy term.

DSpark-Specific Arguments

DSpark builds on DFlash, so all DFlash-specific arguments apply as well.

• --markov-rank (int, default: 256) Low-rank dim of the Markov logit-bias head. 0 disables it.

• --markov-head-type (str, default: "vanilla") Sequential head variant. Options: vanilla, gated, rnn.

• --enable-confidence-head / --no-enable-confidence-head (flag, default: True) Attach the per-position acceptance confidence head.

• --confidence-head-with-markov / --no-confidence-head-with-markov (flag, default: True) Feed the Markov previous-token embedding into the confidence head alongside the backbone hidden state.

• --confidence-head-alpha (float, default: 1.0) Weight of the confidence-head BCE term.

Sliding Window Attention Arguments

All speculator types (except mtp) use sliding window attention on all draft layers by default.

• --sliding-window (int, default: 2048) Sliding window size for sliding window attention layers.

• --full-attention-indices (int list, default: none) Space-separated draft layer indices that should use full attention instead of sliding window. Example: --full-attention-indices 0 2 makes layers 0 and 2 use full attention; the rest use sliding window.

• --sliding-window-non-causal (flag) Use non-causal (bidirectional) masking within draft blocks for sliding window attention layers. Full attention layers are always bidirectional. Note: vLLM currently doesn't support these models.

Dataloader Arguments

• --num-workers (int, default: 12) Number of dataloader worker processes.

• --prefetch-factor (int, default: 4) Number of batches to prefetch per worker.

• --noise-std (float, default: 0.05) Standard deviation for noise augmentation on hidden states.

Checkpoint Arguments

• --checkpoint-freq (int, default: 1) Save a checkpoint every N epochs. Must be ≥ 1.

• --save-best (flag) Save a symbolic link to the checkpoint with the lowest validation loss.

Learning Rate Scheduler Arguments

• --scheduler-type (str, default: "linear") Type of learning rate scheduler. Options: linear, cosine, none

• --scheduler-warmup-steps (int, default: None) Number of warmup steps for the scheduler.

• --scheduler-warmup-ratio (float, default: None) Warmup as a fraction of total scheduler steps, in [0, 1]. Ignored (with a warning) when --scheduler-warmup-steps is also set.

• --scheduler-total-steps (int, default: None) Total number of training steps for the scheduler.

• --scheduler-num-cosine-cycles (float, default: 0.5) Number of cosine cycles for cosine scheduler.

Examples

Online Training

# First, start vLLM server
python scripts/launch_vllm.py \
  meta-llama/Llama-3.1-8B-Instruct \
  -- --port 8000

# Then train with on-demand hidden states generation
speculators train \
  --verifier-name-or-path meta-llama/Llama-3.1-8B-Instruct \
  --data-path ./training_data \
  --vllm-endpoint http://localhost:8000/v1 \
  --on-missing generate \
  --on-generate delete \
  --save-path ./checkpoints \
  --draft-vocab-size 32000 \
  --epochs 10 \
  --lr 3e-5

Offline Training

# Train using pre-generated hidden states
speculators train \
  --verifier-name-or-path meta-llama/Llama-3.1-8B-Instruct \
  --data-path ./training_data \
  --hidden-states-path ./hidden_states \
  --on-missing raise \
  --save-path ./checkpoints \
  --draft-vocab-size 32000 \
  --epochs 10 \
  --lr 3e-5

Hybrid Training (Cache on First Epoch)

speculators train \
  --verifier-name-or-path meta-llama/Llama-3.1-8B-Instruct \
  --data-path ./training_data \
  --hidden-states-path ./hidden_states \
  --vllm-endpoint http://localhost:8000/v1 \
  --on-missing generate \
  --on-generate cache \
  --save-path ./checkpoints \
  --draft-vocab-size 32000 \
  --epochs 10 \
  --lr 3e-5

Multi-GPU Training with WandB Logging

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun \
  --standalone \
  --nproc_per_node 4 \
  -m speculators.train \
  --verifier-name-or-path meta-llama/Llama-3.1-70B-Instruct \
  --data-path ./training_data \
  --hidden-states-path ./hidden_states \
  --save-path ./checkpoints \
  --draft-vocab-size 32000 \
  --epochs 20 \
  --lr 1e-4 \
  --logger wandb \
  --run-name eagle3-llama-70b \
  --scheduler-type cosine \
  --scheduler-warmup-steps 100 \
  --checkpoint-freq 2 \
  --save-best \
  --fsdp-shard

Fine-tuning a Pretrained Model

speculators train \
  --verifier-name-or-path meta-llama/Llama-3.1-8B-Instruct \
  --from-pretrained ./pretrained_speculator \
  --data-path ./new_training_data \
  --hidden-states-path ./hidden_states \
  --save-path ./finetuned_checkpoints \
  --epochs 5 \
  --lr 5e-6

Initializing From a Decoder Config (with Dry-Run Validation)

# Build the speculator from a plain decoder config, initialize weights, save a
# checkpoint, and exit before training so it can be validated in vLLM first.
speculators train \
  --verifier-name-or-path Qwen/Qwen3-8B \
  --speculator-type dflash \
  --draft-config ./qwen3_draft_decoder_config.json \
  --draft-vocab-size 32000 \
  --save-path ./draft_init \
  --dry-run

# After validating ./draft_init in vLLM, train starting from it:
speculators train \
  --verifier-name-or-path Qwen/Qwen3-8B \
  --speculator-type dflash \
  --from-pretrained ./draft_init \
  --data-path ./training_data \
  --epochs 5 \
  --lr 5e-6