# train

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

## Basic Usage

**Single-GPU:**

```bash
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):**

```bash
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):**

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

```bash
# 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

```bash
# 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)

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

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

```bash
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)

```bash
# 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
```