Skip to content

docs(speculative): add subsystem README, fold in regeneration guide #2448

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
HuiyingLi merged 1 commit into from
Jun 9, 2026
Merged

docs(speculative): add subsystem README, fold in regeneration guide #2448

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
333 changes: 333 additions & 0 deletions examples/speculative/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,333 @@
# Speculative decoding: draft-model training in NeMo AutoModel

This directory trains the **draft models** used for speculative decoding. A draft
proposes several future tokens cheaply, the frozen **target** model verifies them
in one forward pass, and every accepted token is a step the target never had to
run autoregressively. The faster and more accurate the draft, the higher the
acceptance length and the larger the inference speedup.

AutoModel trains the draft; you then serve `draft + target` together in an
inference engine (SGLang or vLLM). The training code lives in
`nemo_automodel/components/speculative/`, the recipes in
`nemo_automodel/recipes/llm/`, and ready-to-run configs in this folder.

## Supported methods

| Method | What the draft does | Recipe (`recipe:` key) | Configs |
|---|---|---|---|
| **EAGLE-1** | Single decoder block that predicts the next target hidden state; supervised by SmoothL1 hidden-state loss plus a token loss through the frozen target `lm_head`. | `TrainEagle1Recipe` | `eagle1/` |
| **EAGLE-2** | Same training objective as EAGLE-1 (it differs only in the inference-time tree policy), so the recipe is a thin subclass. | `TrainEagle2Recipe` | `eagle2/` |
| **EAGLE-3** | Draft with its own `lm_head` over a (optionally compressed) draft vocab, fed three auxiliary target hidden states, trained with a **test-time-training (TTT)** unroll. | `TrainEagle3Recipe` | `eagle3/` |
| **EAGLE-3.1** | EAGLE-3 plus two drafter toggles (`fc_norm`, `norm_output`), matching the vLLM EAGLE-3.1 architecture. | `TrainEagle3Recipe` | `eagle3_1/` |
| **P-EAGLE** | Parallel-drafting EAGLE-3: predicts all `num_depths` tokens in one forward over a COD-subsampled sequence instead of the TTT unroll. Serves on vLLM only. | `TrainEagle3Recipe` (`parallel_drafting: true`) | `p-eagle/` |
| **DFlash** | Block-parallel drafting: drafts a whole block of `block_size` tokens in one non-causal "denoising" forward over `[anchor, MASK, MASK, ...]`. | `TrainDFlashRecipe` (LLM) / `DiffusionLMSFTRecipe` (DLLM SFT) | `dflash/`, `../dllm_sft/*dflash*` |

EAGLE-1/2/3 keep their separate code paths: `*_v12.py` files are the **EAGLE-1/2**
("v1/v2") implementation; the unsuffixed files are the EAGLE-3/3.1 path. Inside
EAGLE-3, `fc_norm`/`norm_output` upgrade to 3.1 and `parallel_drafting` upgrades
to P-EAGLE, all from the same draft class.

## Supported target models

A target's `config.architectures` string selects the draft architecture through a
registry (`eagle/registry.py`, `dflash/registry.py`). Capability is per registry,
not per (method, target) pair: the EAGLE-3.1 (`fc_norm` / `norm_output`) and
P-EAGLE (`parallel_drafting`) toggles ride on the same EAGLE-3 dense draft, so
they apply to any target the EAGLE-3 registry maps. The shipped example configs
only cover a subset.

- **EAGLE-1/2/3 (and the EAGLE-3.1 / P-EAGLE toggles on top of EAGLE-3)**:
`LlamaForCausalLM`, `Phi3ForCausalLM`, `Qwen3ForCausalLM`, `Qwen3MoeForCausalLM`.
- **gpt-oss** (`GptOssForCausalLM`): EAGLE-3 only, via a dedicated draft class.
- **DFlash**: `Qwen3ForCausalLM`, `Qwen3MoeForCausalLM`.

Qwen3-MoE is handled exactly like a dense target: the draft only consumes
post-block hidden states, never per-expert routing. gpt-oss uses a dedicated
draft class that reuses the target's YaRN rotary embedding but keeps the on-disk
`architectures` string as the Llama EAGLE-3 draft so inference engines load it
unchanged.

## Quick start

The CLI is `automodel` (alias `am`); it drives `torchrun` internally.

```bash
automodel examples/speculative/eagle3/llama_eagle3_mvp.yaml --nproc-per-node 8
```

Override any config key inline:

```bash
automodel examples/speculative/eagle3/llama_eagle3_perfectblend.yaml --nproc-per-node 8 --recipe_args.micro_batch_size=2
```

The DFlash **DLLM SFT** configs under `../dllm_sft/` use the standard AutoModel
SFT entry script instead:

```bash
torchrun --nproc-per-node 8 examples/dllm_sft/finetune.py -c examples/dllm_sft/qwen3_4b_dflash.yaml
```

Each config family ships an `*_mvp.yaml` (tiny single-GPU smoke with placeholder
data paths) and a `*_perfectblend.yaml` (real run on
`frankleeeee/PerfectBlend-Regenerated-Llama-3.1-8B-Instruct`).

## Operators, kernels, and attention backends

Beyond the methods themselves, the subsystem supports several compute backends.
Pick them through config; all degrade gracefully when a dependency is missing.

### Draft attention backend

| Method | Backends | How to select |
|---|---|---|
| EAGLE-3 / 3.1 | `eager`, `flash_attention_2` | `recipe_args.draft_attn_implementation` (default `eager`) |
| EAGLE-1/2 | `eager` only | n/a |
| P-EAGLE | `flex_attention` (compiled when CUDA + head_dim ≥ 16, else eager flex) | automatic |
| DFlash | `flex_attention`, `sdpa` | `recipe_args.attention_backend` (default `flex_attention`) |

For EAGLE-3, FlashAttention-2 is real FA2 over the TTT attention pattern: FA2
computes the `T×T` causal block and returns `softmax_lse`, and the diagonal
extension columns for cached TTT steps are merged in log space via `logaddexp`.
The draft declares `_supports_flash_attn = True`, FA2 availability is probed
defensively (`_HAS_FA`), and requesting `flash_attention_2` without `flash-attn`
installed raises rather than silently falling back. A ready example is
`eagle3/llama_eagle3_mvp_flash_attn.yaml`. FA2 requires a right-padded attention
mask (enforced at runtime).

### Fused Triton soft cross-entropy

EAGLE-3/P-EAGLE supervise the draft with a masked soft cross-entropy that uses a
**fused Triton kernel** when Triton is available and the logits are on CUDA
(`components/loss/soft_ce.py`, `components/loss/triton/soft_cross_entropy.py`),
falling back to a pure-PyTorch `log_softmax` path otherwise. The masked reduction
normalizes by valid-position count.

### Draft-vocab compression (d2t / t2d)

EAGLE-3 can shrink the draft `lm_head` to `draft_vocab_size < target_vocab_size`.
Training carries two tensors: `selected_token_ids` (draft index to target id, the
"d2t" direction) and `selected_token_mask` (a boolean membership mask over the
full target vocab, the "t2d" direction). The mapping is built and cached by
`components/datasets/llm/eagle3.py`; set `recipe_args.draft_vocab_size` to enable,
or point `recipe_args.selected_token_ids_path` at a precomputed map. `t2d` is
unset when the draft vocab is uncompressed.

## Target backends

The frozen target produces the supervision signal (aux hidden states plus the
target distribution). EAGLE-3 supports three ways to run it.

| Backend | `recipe_args.target_model_backend` | When to use |
|---|---|---|
| **Co-located (default)** | `colocated` | Target and draft share the same GPUs. Simplest; default for every config. |
| **Remote** | `remote` | Target served on separate GPUs/host; training streams supervision over HTTP (control) + NCCL (data, with a binary wire fallback). Numerically identical to co-located. |
| **Offline cache** | set `cached_target_path` | Precompute target outputs once to disk, then train without the target loaded. Disk-heavy and largely superseded by the remote backend. |

Remote serving (`eagle3/llama_eagle3_remote.yaml`): start a server first, then
point training at it.

```bash
python -m nemo_automodel.components.speculative.serve_target --target meta-llama/Llama-3.1-8B-Instruct --host 0.0.0.0 --port 8001
```

```yaml
recipe_args:
target_model_backend: remote
remote_urls: ["http://localhost:8001"]
target_prefetch_depth: 1
```

Offline cache is produced by `precompute_eagle3.py`
(`python -m nemo_automodel.components.speculative.precompute_eagle3 --target-model ... --input-data ... --output-dir ...`),
then consumed via `cached_target_path`. EAGLE-3 only.

EAGLE drafters learn best when the assistant turns in the training data are
produced by the **same model** that will serve as the inference target. Most
public chat datasets were generated by other models, so their assistant tokens
are off-distribution for the drafter. `components/speculative/regenerate.py`
replaces those answers with fresh ones from the target model.

Use it when you want a drafter for a specific target but your only conversational
data came from a different model, or when you have a curated prompts dataset
(ShareGPT, UltraChat, an internal corpus) and want its answer distribution
aligned with the target. If a regenerated set already exists on the Hub (for
example `frankleeeee/PerfectBlend-Regenerated-Llama-3.1-8B-Instruct`), skip this
and point `recipe_args.train_data_path` straight at it.

### Two-step flow

The script talks to an OpenAI-compatible chat-completion endpoint, so the target
must already be served. The examples use SGLang; vLLM or any other
OpenAI-compatible server works too.

Step 1, start the target server (use `--tp 2` or higher to shard a multi-GPU
target):

```bash
python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --port 30000
```

Step 2, regenerate against the running server:

```bash
python -m nemo_automodel.components.speculative.regenerate --input-data Aeala/ShareGPT_Vicuna_unfiltered --output-dir ./regenerated/sharegpt_llama31_8b --target-server http://localhost:30000/v1 --model meta-llama/Llama-3.1-8B-Instruct --concurrency 64 --shard-size 1000
```

For each sample the script loads the `messages` column (HF Hub id, local
parquet, or JSON/JSONL, same loader as `ChatDataset`), drops every trailing
`assistant` turn while keeping the leading `system / user / ...` context
(intermediate assistant turns in multi-turn conversations are kept), calls
`/v1/chat/completions` on the target with that prompt, appends the response as the
new assistant turn, and writes the rebuilt conversations to `shard-NNNNNN.parquet`
files of `--shard-size` rows each.

The run is resumable: rerun with the same `--output-dir` and `--resume` to skip
shards already on disk. A `manifest.json` guards resume, so changing the input
dataset, split, target model, or shard sizing fails fast instead of silently
mixing incompatible shards.

The output is a parquet dataset with a `messages` column, exactly what
`ChatDataset` (used by `build_eagle3_dataloader`) consumes. Point the recipe at
it:

```yaml
recipe_args:
target_model_name_or_path: meta-llama/Llama-3.1-8B-Instruct
train_data_path: ./regenerated/sharegpt_llama31_8b
val_data_path: null
```

### Regeneration tuning knobs

| Flag | Default | Notes |
|---|---|---|
| `--concurrency` | 32 | In-flight requests; raise to saturate the target server. |
| `--shard-size` | 1000 | Smaller shards mean more frequent checkpointing and more files. |
| `--max-new-tokens` | 1024 | Cap per-answer length. |
| `--temperature` | 0.0 | Greedy by default; drafters are typically trained against argmax answers. |
| `--top-p` | 1.0 | Only relevant with `temperature > 0`. |
| `--timeout-s` | 600 | Per-request timeout; bump for very long generations. |
| `--max-retries` | 3 | Retries on 5xx, 429, and transport errors with exponential backoff. |
| `--split` | `train` | Supports HF slice syntax, e.g. `train[:10000]`. |
| `--shuffle-seed` | unset | Optional shuffle before slicing. |

### Regeneration pitfalls

- **Wrong model name.** `--model` is the name sent in the OpenAI payload; it must
match what the server serves. SGLang uses `--model-path` as the served name by
default, so mirror `--served-model-name` here if you set it.
- **Server not warm.** Send one curl request to the server first; otherwise the
script retries then fails on the first batch.
- **Tokenizer mismatch.** The regenerated dataset is consumed by `ChatDataset`,
which applies the target model's chat template at training time. Make sure the
recipe's tokenizer comes from the same model id you used for `--model`, or the
loss-mask alignment silently drifts.

Datasets are consumed by `ChatDataset`: a `messages` list of `{role, content}`,
or a `conversations` column (ShareGPT or OpenAI style) that is auto-converted.

## Serve and benchmark a trained draft

After training, serve `target + draft` through SGLang:

```bash
python -m nemo_automodel.components.speculative.serve_sglang --target meta-llama/Llama-3.1-8B-Instruct --draft /path/to/run/epoch_0_step_1000/model --algorithm EAGLE3 --num-steps 3 --num-draft-tokens 4
```

`serve_sglang.py` resolves the consolidated `model/` directory, rewrites the
draft `architectures` to SGLang's canonical name, and regenerates SGLang's
speculative token map from `eagle_meta.pt` when needed. SGLang is not bundled; the
tool exits with an install hint if it is missing.

Measure acceptance length and speedup against the running server:

```bash
python -m nemo_automodel.components.speculative.bench_sglang --server http://localhost:30000 --model meta-llama/Llama-3.1-8B-Instruct --input-data <prompts-dataset> --baseline-server http://localhost:30001
```

It reports `accept_length` (mean tokens per verify step), `acceptance_rate`,
output throughput, and a `speedup` ratio versus an optional non-speculative
baseline server. Point it at a freshly started server, since SGLang reports a
server-cumulative average.

## Inference-engine compatibility

| Draft | SGLang | vLLM |
|---|---|---|
| EAGLE-1/2/3, EAGLE-3.1 | yes | yes |
| P-EAGLE | no (tracked upstream) | yes (parallel-drafting runtime) |

`serve_sglang.py` rejects P-EAGLE drafts with an actionable error; serve those on
vLLM.

## Config reference (EAGLE-style schema)

EAGLE-1/2/3/3.1, P-EAGLE, and the LLM DFlash recipe share one schema. The DFlash
**DLLM SFT** configs under `../dllm_sft/` use the standard AutoModel SFT schema
(`step_scheduler` / `model._target_` / `dataset._target_` / `dllm` / `dflash`
blocks) instead.

### Top-level sections

| Section | Purpose |
|---|---|
| `recipe` | Recipe class name (required). |
| `recipe_args` | Main training block (below). |
| `dist_env` | `backend` (nccl), `timeout_minutes`. |
| `distributed` | Optional; only for MoE / large targets. `strategy: fsdp2`, `tp_size`, `pp_size`, `cp_size`, `ep_size`, `activation_checkpointing`, `sequence_parallel`. Absent means DDP. |
| `optimizer` | `lr`, `betas`, `weight_decay`, optional `warmup_ratio` (0.05), `min_lr_ratio` (0.1). |
| `checkpoint` | `enabled`, `checkpoint_dir`, `model_save_format: safetensors`, `save_consolidated`, optional `restore_from` (`LATEST` / subdir / path). |
| `wandb` | Optional; `project`, `entity`, `name`. |

### `recipe_args` common to all methods

`target_model_name_or_path`, `train_data_path`, `val_data_path`, `train_split`,
`val_split`, `output_dir`, `seq_length`, `micro_batch_size`,
`grad_accumulation_steps`, `num_workers`, `num_epochs`, `freeze_embeddings`,
`trust_remote_code`, `shuffle_seed`, `log_every_steps`, `max_grad_norm`. Optional
checkpoint cadence: `ckpt_every_steps`, `save_checkpoint_every_epoch`.

### Method-specific `recipe_args`

| Key | Methods | Notes |
|---|---|---|
| `draft_num_hidden_layers` | EAGLE-1/2, DFlash | Stacked draft decoder layers. |
| `hidden_loss_weight`, `token_loss_weight` | EAGLE-1/2 | Defaults 1.0 / 0.1. |
| `ttt_steps` | EAGLE-3 / 3.1 | TTT unroll depth; integer ≥ 1 (required). |
| `draft_vocab_size` | EAGLE-3 family | Compress the draft `lm_head`; omit for full vocab. |
| `selected_token_ids_path` | EAGLE-3 family | Reuse a cached draft-vocab map. |
| `aux_layer_ids` | EAGLE-3 | Override the default low/mid/high recipe `[1, n//2-1, n-4]`. |
| `draft_attn_implementation` | EAGLE-3 | `eager` (default) or `flash_attention_2`. |
| `fc_norm`, `norm_output` | EAGLE-3.1 | Both default false; either alone is a valid intermediate config. |
| `target_model_backend`, `remote_urls`, `target_prefetch_depth`, `remote_timeout`, `remote_max_retries` | EAGLE-3 remote | See Target backends. |
| `cached_target_path` | EAGLE-3 offline | Path to a `precompute_eagle3` cache. |
| `parallel_drafting`, `num_depths`, `num_draft_layers`, `down_sample_ratio`, `down_sample_ratio_min`, `mask_token_id`, `sequence_partitions` | P-EAGLE | `mask_token_id` is required (no default). `sequence_partitions > 1` splits each sequence by dependency lineage to bound long-context memory. |
| `block_size`, `num_anchors`, `loss_decay_gamma`, `mask_token_id`, `target_layer_ids`, `attention_backend` | DFlash (LLM recipe) | Block drafting knobs. |

## Directory layout

Configs in this folder:

```
examples/speculative/
eagle1/ eagle2/ eagle3/ eagle3_1/ p-eagle/ dflash/
README.md # this file (includes the dataset regeneration guide)
examples/dllm_sft/ # DFlash DLLM SFT configs (standard SFT schema)
```

Implementation:

```
nemo_automodel/components/speculative/
eagle/ core(.py/_v12), draft_llama(.py/_v12), draft_gpt_oss, backend,
registry, target(.py/_v12), peagle_*, remote/
dflash/ core, draft_qwen3, registry, target
regenerate.py # dataset regeneration with the target model
precompute_eagle3.py # offline target-output cache
serve_target.py # remote target server (HTTP + NCCL)
serve_sglang.py # serve a trained draft via SGLang
bench_sglang.py # acceptance-length / speedup benchmark
nemo_automodel/recipes/llm/
train_eagle1.py train_eagle2.py train_eagle3.py train_dflash.py peagle_recipe.py
```
Loading
Loading