From d9708cdf7a2f5c5724b326417845ec92b275dd82 Mon Sep 17 00:00:00 2001 From: Peter Hicks Date: Thu, 15 Oct 2026 16:28:58 +0000 Subject: [PATCH 1/5] [8/8] Add an external-draft training runbook Walks the full workflow end to end: dump a teacher policy once, then train several drafts against it, including one whose tokenizer differs from the target's and so needs the vocabulary swap. Documents the objectives and what each one actually scores, the export path (no export step -- the draft is already a plain HF checkpoint), and the failure modes that produce a falling loss curve with no acceptance. Co-Authored-By: Claude Opus 5 Signed-off-by: Peter Hicks --- examples/speculative_decoding/README.md | 4 + .../doc/external_draft.md | 253 ++++++++++++++++++ 2 files changed, 257 insertions(+) create mode 100644 examples/speculative_decoding/doc/external_draft.md diff --git a/examples/speculative_decoding/README.md b/examples/speculative_decoding/README.md index 786944ca9ef..f5d55f8117f 100644 --- a/examples/speculative_decoding/README.md +++ b/examples/speculative_decoding/README.md @@ -187,6 +187,10 @@ out-of-support term charges it for the mass the dump itself discarded. That is t same asymmetry dense `tvd` has against a full teacher, and it is intended, but it means `--top-p` is only symmetric between the two sides under `tvd_deploy`. +See [doc/external_draft.md](doc/external_draft.md) for the full reference: the recipe +parameters, the two teacher-data formats, the objectives, and the cross-tokenizer +vocabulary swap. + `external.external_loss` selects the objective: | value | what it scores | diff --git a/examples/speculative_decoding/doc/external_draft.md b/examples/speculative_decoding/doc/external_draft.md new file mode 100644 index 00000000000..6efb9f2bc02 --- /dev/null +++ b/examples/speculative_decoding/doc/external_draft.md @@ -0,0 +1,253 @@ +# External Draft — Training a Standalone Draft Model + +An external draft is an ordinary pretrained causal LM trained to propose tokens for a +target model. Unlike the head-based modes (EAGLE, Medusa, DFlash) nothing is grafted onto +the target: it contributes its vocabulary and, on the dense path, its `lm_head`. The result +exports as a plain HuggingFace checkpoint. + +Recipe: [`modelopt_recipes/general/speculative_decoding/external_draft.yaml`](../../../modelopt_recipes/general/speculative_decoding/external_draft.yaml) | +Swap implementation: [`vocab_swap.py`](../../../modelopt/torch/speculative/external/vocab_swap.py) + +## Architecture + +```text +Target (frozen, never modified) + │ + ├─ sparse path: top-k deployment policy ──► dumped to disk ──┐ + │ │ + └─ dense path: hidden states ──► target lm_head ──► logits ──┤ + │ + TVD / soft_ce + │ +Draft (a complete causal LM, trained) ──► its own lm_head ──► draft logits +``` + +**Key components:** + +- **No grafted head.** The draft owns its embeddings and `lm_head` and runs standalone at + serving time. Engines that already support a separate draft model load it with no + mode-specific support. +- **Acceptance-shaped objectives.** Under rejection sampling the probability a proposed + token survives verification is `sum(min(p, q))`, so total-variation distance is + `1 - acceptance`. Training on TVD optimises the serving metric directly. +- **Two data formats.** The TVD objectives need only the target's truncated next-token + distribution, which is far smaller than a hidden state per token and keeps the target off + the GPU at training time. `soft_ce` needs full logits and so needs dumped hidden states. +- **Vocabulary coupling.** The draft must index the target's vocabulary. A draft from + another family can be re-indexed onto it, at the cost of replacing a large fraction of a + small model. + +## Training + +### Quick Start + +```bash +torchrun --nproc_per_node 8 main.py \ + --config ../../modelopt_recipes/general/speculative_decoding/external_draft.yaml \ + model.model_name_or_path=$TARGET \ + draft_model_name_or_path=$DRAFT \ + data.sparse_data_path=$POLICY_DIR \ + training.output_dir=$OUTPUT_DIR +``` + +### Recipe + +See [`modelopt_recipes/general/speculative_decoding/external_draft.yaml`](../../../modelopt_recipes/general/speculative_decoding/external_draft.yaml) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `model.model_name_or_path` | — | The **target**, used for its vocabulary and `lm_head` | +| `draft_model_name_or_path` | — | The draft, the model actually trained | +| `external.external_loss` | `tvd` | `soft_ce`, `tvd`, `tvd_deploy` or `tvd_ce` (see below) | +| `external.external_top_k` | 20 | Serving top-k; must match the dump | +| `external.external_top_p` | 0.95 | Serving top-p; must match the dump | +| `external.external_tvd_alpha` | 0.9 | Weight on the TVD term under `tvd_ce` | +| `external.external_ce_alpha` | 0.1 | Weight on the CE term under `tvd_ce` | +| `external.external_vocab_swap` | false | Re-index the draft onto the target's vocabulary | +| `external.external_report_acc` | true | Log top-1 agreement during training | +| `data.sparse_data_path` | — | Dumped top-k policy; selects `data.mode=sparse` | +| `data.sparse_cache_dir` | — | Where the memmapped cache is built | +| `data.offline_data_path` | — | Dumped hidden states; the alternative to the above | +| `model.use_fake_base_for_offline` | true | Load only the target's `lm_head`, not its layers | + +> **Note on `external_top_k` / `external_top_p`:** the stored policy *is* the teacher +> distribution, so a dump narrower than the objective expects charges the draft for mass on +> tokens the teacher actually supported. A top-k mismatch draws a warning under +> `tvd_deploy`; the dump does not record its top-p, so a top-p mismatch goes undetected. + +### Objectives + +| `external_loss` | Scores | Needs | +|-----------------|--------|-------| +| `soft_ce` | Cross-entropy against the target's soft distribution (the EAGLE offline objective) | hidden states | +| `tvd` | Total-variation distance between the two distributions | either format | +| `tvd_deploy` | The same, with both sides put through the serving filter first | either format | +| `tvd_ce` | `tvd` plus a cross-entropy term on the target's argmax token | either format | + +On the sparse path the target side is truncated by the dump whichever objective runs, since +that is all the dump holds; `tvd` and `tvd_deploy` then differ only in whether the *draft* +is put through the filter too. Only `tvd_deploy` applies the same nucleus to both sides. +Pairing `soft_ce` with `data.sparse_data_path` is a hard error, not a quiet degradation. + +`tvd_ce` exists because TVD is bounded and so ranks weakly among tokens inside the target's +support; cross-entropy's log penalty has no such ceiling. + +### Teacher policy data + +```bash +python collect_hidden_states/compute_sparse_policy_hf.py \ + --model $TARGET --input-data $CONVERSATIONS --output-dir $POLICY_DIR \ + --top-k 20 --top-p 0.95 --max-seq-len 4096 \ + --answer-only-loss --chat-template ./template.jinja +``` + +Set `--max-seq-len` to the `training.training_seq_len` you will train at. The producer skips +conversations longer than it, and the collator pads every sample to `training_seq_len`, so +over-sizing spends the difference on padding and under-sizing silently drops records that +carry no scored position. + +On a large corpus the first run spends a long time building the memmapped cache beside the +shards. Under `torchrun` only rank 0 builds it and the others wait, so build it once in a +single process first and point every run at the same `data.sparse_cache_dir`: + +```bash +python - < **Note on `answer_only_loss` and chat templates:** when `--answer-only-loss` is set, the +> chat template must wrap assistant content in `{% generation %}` / `{% endgeneration %}` +> tags. Most published templates do not. Copy the model's own template and add them: +> +> 1. Print `AutoTokenizer.from_pretrained(TARGET).chat_template` to a file +> 2. Wrap the assistant content each branch emits, including a reasoning branch if present +> 3. Pass it with `--chat-template` to **every** dump command for that target +> +> Verify by decoding, not by counting: `tokenizer.decode(record["gen_ids"])` should read as +> the assistant's reply alone. A misplaced wrap still cuts at the wrong boundary, so a +> length check passes while the span is wrong. + +The dense alternative dumps hidden states instead. Pass `--no-aux-hidden-states`: this mode +never reads the aux planes and they dominate the dump size. + +```bash +python collect_hidden_states/compute_hidden_states_hf.py \ + --model $TARGET --input-data $CONVERSATIONS --output-dir $HIDDEN_DIR \ + --no-aux-hidden-states --answer-only-loss --chat-template ./template.jinja +``` + +### Cross-tokenizer vocabulary swap (`external_vocab_swap`) + +Speculative decoding needs the draft to propose tokens the target can verify, so the two +must index the same vocabulary. When the target's family has no small member to draft with, +`external.external_vocab_swap=true` rebuilds the draft's embeddings and `lm_head` against +the target's vocabulary before conversion. Rows for tokens both tokenizers spell identically +carry over; the rest are drawn from the draft's own embedding statistics. + +The run reports what carried over, typically a minority of the table and nearly all of the +most frequent ids, since the rare tail is where tokenizers disagree. + +> **Note on judging a swap:** the match percentage describes the swap, not the draft. +> Overlap figures far apart have produced similar acceptance. Measure acceptance on held-out +> data before committing to a swapped draft, and prefer a vocabulary-matched draft where one +> exists: the replacement matrix is `base_vocab_size × hidden_size` parameters, a large +> fraction of a small draft, and a swapped draft stays behind a matched one of similar size. + +A swapped draft starts with a large untrained fraction and wants a higher learning rate than +a matched draft of comparable size. + +### Checkpoint resume + +A resumed draft keeps the embeddings it was saved with. The swap reads +`external_vocab_swap` out of the checkpoint's own ModelOpt state, so re-running with the +flag on does not rewrite rows that are already base-indexed. + +### Export + +Training writes an ordinary HuggingFace checkpoint, shipped with the **target's** tokenizer +when the vocabulary was swapped, since the draft's rows are indexed by the target's ids from +that point on. Nothing in the saved artifact depends on ModelOpt at load time. + +## Reference Configuration (Llama-3.2-1B drafting for Qwen3.8-27B) + +A cross-tokenizer draft: Llama-3.2-1B-Instruct re-indexed onto Qwen3.8-27B's +248,320-token vocabulary. Trained on a sparse teacher policy dumped once from the target, +8 GPUs, one epoch. + +### Teacher policy + +```bash +python collect_hidden_states/compute_sparse_policy_hf.py \ + --model $TARGET --input-data $CONVERSATIONS --output-dir $POLICY_DIR \ + --top-k 20 --top-p 0.95 --max-seq-len 12288 \ + --answer-only-loss --chat-template ./template.jinja +``` + +### Training + +```bash +torchrun --nproc_per_node 8 main.py \ + --config ../../modelopt_recipes/general/speculative_decoding/external_draft.yaml \ + model.model_name_or_path=$TARGET \ + draft_model_name_or_path=meta-llama/Llama-3.2-1B-Instruct \ + external.external_vocab_swap=true \ + external.external_loss=tvd \ + external.external_top_k=20 external.external_top_p=0.95 \ + data.sparse_data_path=$POLICY_DIR \ + data.sparse_cache_dir=$CACHE_DIR \ + training.training_seq_len=12288 \ + training.per_device_train_batch_size=1 \ + training.gradient_accumulation_steps=2 \ + training.learning_rate=2.0e-5 \ + training.warmup_steps=250 \ + training.max_grad_norm=1.0 \ + training.num_train_epochs=1 \ + training.seed=42 \ + training.output_dir=$OUTPUT_DIR +``` + +### Configuration + +| Parameter | Value | +|-----------|-------| +| Objective | `tvd` | +| Vocabulary swap | yes (Llama tokenizer onto Qwen's) | +| Sequence length | 12,288 | +| Effective batch | 16 (1 per device × 2 accumulation × 8 GPUs) | +| Learning rate | 2.0e-5 | +| Warmup steps | 250 | +| Epochs | 1 | +| Serving policy | top-k 20, top-p 0.95, temperature 1.0 | + +> **Note on the learning rate:** a swapped draft begins with a large untrained fraction and +> wants a higher rate than a matched draft of comparable size. The vocabulary-matched drafts +> trained against the same dump used 1.0e-5 (4B) and 2.0e-6 (0.8B). These were each selected +> by sweeping against held-out acceptance; the ordering does not transfer to another target +> or dump, so re-sweep for your own pair. + +## Open Items + +### Validated end-to-end + +- Sparse teacher-policy training with `tvd`, from dump through training to an exported + checkpoint, with acceptance measured on held-out conversations. +- The cross-tokenizer vocabulary swap, on a draft trained and measured this way. + +### Covered by tests, not yet exercised in a full run + +- `soft_ce`, `tvd_deploy` and `tvd_ce`, and the dense hidden-states data path. +- Resume, with and without a swap. +- Untied output heads, and heads carrying a bias. + +### Not yet validated + +- VLM targets. `main.py` routes through `load_vlm_or_llm`, so the path exists, but no + external draft has been trained against one. From acdeac71514b309b5997e375bb77eddae5217c16 Mon Sep 17 00:00:00 2001 From: Peter Hicks Date: Thu, 22 Oct 2026 11:34:44 +0000 Subject: [PATCH 2/5] Point external-draft teacher dumps at target-generated conversations The teacher policy section handed the reader $CONVERSATIONS without saying where it comes from, which reads as 'use whatever responses you have'. The objective scores the draft at positions it meets when serving, and those lie along the target's own samples, so the responses should be the target's. README already covers producing them under Data Synthesis; link it rather than restate it. Co-Authored-By: Claude Opus 5 Signed-off-by: Peter Hicks --- examples/speculative_decoding/doc/external_draft.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/examples/speculative_decoding/doc/external_draft.md b/examples/speculative_decoding/doc/external_draft.md index 6efb9f2bc02..4840b821f2d 100644 --- a/examples/speculative_decoding/doc/external_draft.md +++ b/examples/speculative_decoding/doc/external_draft.md @@ -94,6 +94,12 @@ support; cross-entropy's log penalty has no such ceiling. ### Teacher policy data +`$CONVERSATIONS` should carry responses the **target generated itself**, not the responses +that shipped with the prompts. The objective trains the draft to match the target at the +positions it will meet when serving, and those lie along the target's own samples; scoring +it against another model's text trains it on positions the target would rarely produce. +[Data Synthesis](../README.md#data-synthesis) covers generating them. + ```bash python collect_hidden_states/compute_sparse_policy_hf.py \ --model $TARGET --input-data $CONVERSATIONS --output-dir $POLICY_DIR \ From 80863b54fccf6e4be5117ee6804c2fd50a2be20f Mon Sep 17 00:00:00 2001 From: Peter Hicks Date: Thu, 22 Oct 2026 12:30:28 +0000 Subject: [PATCH 3/5] Address review on the external-draft runbook - Quick start pinned training_seq_len=4096 and sparse_cache_dir, which the recipe defaults (2048, cache beside the shards) contradicted. - Dropped the recipe link repeated from the header, and the learning-rate sentence repeated from the swap section. - max-seq-len: the producer skips whole conversations and prints the count, so the previous 'silently drops records' was wrong twice over. - Dense dump now sets --max-seq-len explicitly; it defaults to 3072 and so would have skipped what the sparse dump at 4096 kept. - Second teacher-policy block replaced by a cross-reference, so the two cannot drift. Co-Authored-By: Claude Opus 5 Signed-off-by: Peter Hicks --- .../doc/external_draft.md | 29 ++++++++----------- 1 file changed, 12 insertions(+), 17 deletions(-) diff --git a/examples/speculative_decoding/doc/external_draft.md b/examples/speculative_decoding/doc/external_draft.md index 4840b821f2d..9f2154a8e80 100644 --- a/examples/speculative_decoding/doc/external_draft.md +++ b/examples/speculative_decoding/doc/external_draft.md @@ -47,13 +47,13 @@ torchrun --nproc_per_node 8 main.py \ model.model_name_or_path=$TARGET \ draft_model_name_or_path=$DRAFT \ data.sparse_data_path=$POLICY_DIR \ + data.sparse_cache_dir=$CACHE_DIR \ + training.training_seq_len=4096 \ training.output_dir=$OUTPUT_DIR ``` ### Recipe -See [`modelopt_recipes/general/speculative_decoding/external_draft.yaml`](../../../modelopt_recipes/general/speculative_decoding/external_draft.yaml) - | Parameter | Default | Description | |-----------|---------|-------------| | `model.model_name_or_path` | — | The **target**, used for its vocabulary and `lm_head` | @@ -108,9 +108,9 @@ python collect_hidden_states/compute_sparse_policy_hf.py \ ``` Set `--max-seq-len` to the `training.training_seq_len` you will train at. The producer skips -conversations longer than it, and the collator pads every sample to `training_seq_len`, so -over-sizing spends the difference on padding and under-sizing silently drops records that -carry no scored position. +whole conversations longer than it, scored turns included, and prints how many it skipped. The +collator pads every sample to `training_seq_len`, so over-sizing spends the difference on +padding and under-sizing drops the long conversations from the dump. On a large corpus the first run spends a long time building the memmapped cache beside the shards. Under `torchrun` only rank 0 builds it and the others wait, so build it once in a @@ -147,7 +147,7 @@ never reads the aux planes and they dominate the dump size. ```bash python collect_hidden_states/compute_hidden_states_hf.py \ --model $TARGET --input-data $CONVERSATIONS --output-dir $HIDDEN_DIR \ - --no-aux-hidden-states --answer-only-loss --chat-template ./template.jinja + --max-seq-len 4096 --no-aux-hidden-states --answer-only-loss --chat-template ./template.jinja ``` ### Cross-tokenizer vocabulary swap (`external_vocab_swap`) @@ -190,12 +190,8 @@ A cross-tokenizer draft: Llama-3.2-1B-Instruct re-indexed onto Qwen3.8-27B's ### Teacher policy -```bash -python collect_hidden_states/compute_sparse_policy_hf.py \ - --model $TARGET --input-data $CONVERSATIONS --output-dir $POLICY_DIR \ - --top-k 20 --top-p 0.95 --max-seq-len 12288 \ - --answer-only-loss --chat-template ./template.jinja -``` +The [teacher policy data](#teacher-policy-data) dump, with `--max-seq-len 12288` to match +`training.training_seq_len` below. ### Training @@ -233,11 +229,10 @@ torchrun --nproc_per_node 8 main.py \ | Epochs | 1 | | Serving policy | top-k 20, top-p 0.95, temperature 1.0 | -> **Note on the learning rate:** a swapped draft begins with a large untrained fraction and -> wants a higher rate than a matched draft of comparable size. The vocabulary-matched drafts -> trained against the same dump used 1.0e-5 (4B) and 2.0e-6 (0.8B). These were each selected -> by sweeping against held-out acceptance; the ordering does not transfer to another target -> or dump, so re-sweep for your own pair. +> **Note on the learning rate:** the vocabulary-matched drafts trained against the same dump +> used 1.0e-5 (4B) and 2.0e-6 (0.8B). Each rate was selected by sweeping against held-out +> acceptance; the ordering does not transfer to another target or dump, so re-sweep for your +> own pair. ## Open Items From 2879755868edea8dbcd05dafdb87d85e4fe08e67 Mon Sep 17 00:00:00 2001 From: Peter Hicks Date: Fri, 23 Oct 2026 15:42:37 +0000 Subject: [PATCH 4/5] Drive the Training section's commands from one SEQ_LEN The sequence length appeared as a literal 4096 in four places: the quick start, the sparse dump, the cache pre-build and the dense dump. All four have to agree, so a reader changing one silently got a mismatched set. Declare it once and have the commands read it. The Reference Configuration keeps its literal 12288, since that is a record of one particular run rather than a template. Co-Authored-By: Claude Opus 5 Signed-off-by: Peter Hicks --- examples/speculative_decoding/doc/external_draft.md | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/examples/speculative_decoding/doc/external_draft.md b/examples/speculative_decoding/doc/external_draft.md index 9f2154a8e80..e5c9d801858 100644 --- a/examples/speculative_decoding/doc/external_draft.md +++ b/examples/speculative_decoding/doc/external_draft.md @@ -42,13 +42,16 @@ Draft (a complete causal LM, trained) ──► its own lm_head ──► draft ### Quick Start ```bash +# Every command in this section reads $SEQ_LEN; change it here and they stay consistent. +SEQ_LEN=4096 + torchrun --nproc_per_node 8 main.py \ --config ../../modelopt_recipes/general/speculative_decoding/external_draft.yaml \ model.model_name_or_path=$TARGET \ draft_model_name_or_path=$DRAFT \ data.sparse_data_path=$POLICY_DIR \ data.sparse_cache_dir=$CACHE_DIR \ - training.training_seq_len=4096 \ + training.training_seq_len=$SEQ_LEN \ training.output_dir=$OUTPUT_DIR ``` @@ -103,7 +106,7 @@ it against another model's text trains it on positions the target would rarely p ```bash python collect_hidden_states/compute_sparse_policy_hf.py \ --model $TARGET --input-data $CONVERSATIONS --output-dir $POLICY_DIR \ - --top-k 20 --top-p 0.95 --max-seq-len 4096 \ + --top-k 20 --top-p 0.95 --max-seq-len $SEQ_LEN \ --answer-only-loss --chat-template ./template.jinja ``` @@ -124,7 +127,7 @@ from modelopt.torch.speculative.external.sparse_data import ( ) SparsePolicyDataset( - discover_shards("$POLICY_DIR"), max_length=4096, cache_dir="$CACHE_DIR" + discover_shards("$POLICY_DIR"), max_length=$SEQ_LEN, cache_dir="$CACHE_DIR" ) PY ``` @@ -147,7 +150,7 @@ never reads the aux planes and they dominate the dump size. ```bash python collect_hidden_states/compute_hidden_states_hf.py \ --model $TARGET --input-data $CONVERSATIONS --output-dir $HIDDEN_DIR \ - --max-seq-len 4096 --no-aux-hidden-states --answer-only-loss --chat-template ./template.jinja + --max-seq-len $SEQ_LEN --no-aux-hidden-states --answer-only-loss --chat-template ./template.jinja ``` ### Cross-tokenizer vocabulary swap (`external_vocab_swap`) From c0ce90f98eb4921e4794cb3a57cfbb8a0418f045 Mon Sep 17 00:00:00 2001 From: Peter Hicks Date: Fri, 23 Oct 2026 16:22:36 +0000 Subject: [PATCH 5/5] Collapse the README external-draft section into a pointer The README section and doc/external_draft.md had grown two copies of the dump command, the training command, the top-k/top-p mismatch note, the objectives table and the acceptance explanation, and the wording had already drifted between them. Follow the DFlash pattern instead: a short intro, the two commands, and a link. Three points lived only in the deleted copies and move into the doc so nothing is lost: - streaming is accepted but unexercised, and online training is rejected, into Open Items; - report which truncation you trained under, into Objectives; - a ragged third-party dump needs top_k passed to SparsePolicyDataset, from the recipe comment into the cache pre-build paragraph. The recipe comment held a third copy of the cache pre-build snippet, with a max_length that disagreed with the doc; it is now a pointer. The copy in the TimeoutError message stays self-contained, since that is where a user meets it. Co-Authored-By: Claude Opus 5 Signed-off-by: Peter Hicks --- examples/speculative_decoding/README.md | 55 ++----------------- .../doc/external_draft.md | 11 ++++ .../speculative_decoding/external_draft.yaml | 11 +--- 3 files changed, 17 insertions(+), 60 deletions(-) diff --git a/examples/speculative_decoding/README.md b/examples/speculative_decoding/README.md index f5d55f8117f..f593fc2a009 100644 --- a/examples/speculative_decoding/README.md +++ b/examples/speculative_decoding/README.md @@ -138,72 +138,25 @@ For large base models, you can stream hidden states from a live `vllm serve` ins The modes above graft a draft head onto the target. The `external` mode instead trains an ordinary pretrained causal LM as the draft: it keeps its own embeddings and lm_head, the target is never modified, and the result exports as a plain HF -checkpoint that TRT-LLM, vLLM and SGLang load with no mode-specific support. The -draft must index the same vocabulary as the target. +checkpoint that TRT-LLM, vLLM and SGLang load with no mode-specific support. -Training consumes dumped base hidden states, as in the offline flow above. Pass -`--no-aux-hidden-states` when dumping: this mode never reads them and they -dominate the dump size. Streaming feeds the same hidden states from a live -`vllm serve` and is accepted here, though it has not been exercised for this -mode. Online training is rejected: it produces no cached base outputs for the -draft to be scored against. - -The TVD objectives only ever need the base's *truncated* next-token distribution, -not a full hidden state, so they can instead read a much smaller dump of the -base's top-k deployment policy. Produce one with: +Dump the target's top-k deployment policy once, then train against it: ```bash python collect_hidden_states/compute_sparse_policy_hf.py \ --model $BASE_MODEL --input-data $CONVERSATIONS --output-dir $POLICY_DIR \ --top-k 20 --top-p 0.95 --answer-only-loss -``` -and train against it with `data.sparse_data_path=$POLICY_DIR` in place of -`data.offline_data_path` — the two are alternatives, not a pair: - -```bash torchrun --nproc_per_node 8 main.py \ --config ../../modelopt_recipes/general/speculative_decoding/external_draft.yaml \ model.model_name_or_path=$BASE_MODEL \ draft_model_name_or_path=$DRAFT_MODEL \ data.sparse_data_path=$POLICY_DIR \ - external.external_loss=tvd \ - external.external_top_k=20 external.external_top_p=0.95 \ training.output_dir=$OUTPUT_DIR ``` -Keep `--top-k` / `--top-p` equal to `external.external_top_k` / -`external.external_top_p`: the stored policy *is* the teacher distribution, so a dump -narrower than the objective expects charges the draft for mass on tokens the teacher -actually supported. A top-k mismatch draws a warning under `tvd_deploy`; the dump does not -record its top-p, so a top-p mismatch goes undetected, and only a stored row that does not -sum to 1 draws a warning. This dump -also requires a TVD objective; `soft_ce` needs the full base distribution and is -rejected against `sparse_data_path`. - -Only `tvd_deploy` puts the draft through the same nucleus the dump was written with. -Under `tvd` and `tvd_ce` the draft side is an untruncated softmax, so the -out-of-support term charges it for the mass the dump itself discarded. That is the -same asymmetry dense `tvd` has against a full teacher, and it is intended, but it -means `--top-p` is only symmetric between the two sides under `tvd_deploy`. - -See [doc/external_draft.md](doc/external_draft.md) for the full reference: the recipe -parameters, the two teacher-data formats, the objectives, and the cross-tokenizer -vocabulary swap. - -`external.external_loss` selects the objective: - -| value | what it scores | -|---|---| -| `soft_ce` | cross-entropy against the base's soft distribution (the EAGLE offline objective) | -| `tvd` | total-variation distance between the two full distributions | -| `tvd_deploy` | the same, with both sides put through the serving filter (`external_top_k`, then `external_top_p`) first, so the loss is the acceptance the deployed pair sees | -| `tvd_ce` | `tvd` plus a cross-entropy term on the base's argmax token, weighted by `external_tvd_alpha` / `external_ce_alpha` (default 0.9 / 0.1). TVD is bounded and so ranks weakly among tokens inside the base's support; CE's log penalty does not | - -Under rejection sampling the acceptance probability is exactly `sum(min(p, q))`, -so `tvd_deploy` is `1 - acceptance` — the deployment metric itself rather than a -proxy for it. Report which truncation you trained under, since `top_k` and -`top_p` settings are not interchangeable. +See [doc/external_draft.md](doc/external_draft.md) for the recipe parameters, the +dense hidden-state alternative, the objectives, and the cross-tokenizer vocabulary swap. ## Model Validation diff --git a/examples/speculative_decoding/doc/external_draft.md b/examples/speculative_decoding/doc/external_draft.md index e5c9d801858..477f706bd07 100644 --- a/examples/speculative_decoding/doc/external_draft.md +++ b/examples/speculative_decoding/doc/external_draft.md @@ -92,6 +92,9 @@ that is all the dump holds; `tvd` and `tvd_deploy` then differ only in whether t is put through the filter too. Only `tvd_deploy` applies the same nucleus to both sides. Pairing `soft_ce` with `data.sparse_data_path` is a hard error, not a quiet degradation. +Report which truncation you trained under. `top_k` and `top_p` settings are not +interchangeable, so a number quoted without them is not comparable to anything. + `tvd_ce` exists because TVD is bounded and so ranks weakly among tokens inside the target's support; cross-entropy's log penalty has no such ceiling. @@ -132,6 +135,11 @@ SparsePolicyDataset( PY ``` +Use `discover_shards` and the same `cache_dir` the training run will use: a cache keyed on a +different shard list is ignored, and `torchrun` rebuilds it from scratch. A ragged +third-party dump also needs `top_k=` here, since the cache width is otherwise +inferred from the first record and the recipe has no field to set it. + > **Note on `answer_only_loss` and chat templates:** when `--answer-only-loss` is set, the > chat template must wrap assistant content in `{% generation %}` / `{% endgeneration %}` > tags. Most published templates do not. Copy the model's own template and add them: @@ -253,5 +261,8 @@ torchrun --nproc_per_node 8 main.py \ ### Not yet validated +- Streaming the teacher data from a live `vllm serve`. It is accepted for this mode but has + not been exercised. Online training is rejected outright, since it produces no cached + target outputs for the draft to be scored against. - VLM targets. `main.py` routes through `load_vlm_or_llm`, so the path exists, but no external draft has been trained against one. diff --git a/modelopt_recipes/general/speculative_decoding/external_draft.yaml b/modelopt_recipes/general/speculative_decoding/external_draft.yaml index 235afe4f18b..8c276a2339e 100644 --- a/modelopt_recipes/general/speculative_decoding/external_draft.yaml +++ b/modelopt_recipes/general/speculative_decoding/external_draft.yaml @@ -29,15 +29,8 @@ data: # base lm_head at train time and requires a TVD objective (tvd, tvd_deploy or tvd_ce). # # The first run over a dump builds a memmapped cache beside the shards, which takes - # hours on a large corpus. Under torchrun only rank 0 builds it and the others wait, - # so build it once in a single process first: - # python -c "from modelopt.torch.speculative.external.sparse_data import \ - # SparsePolicyDataset as D, discover_shards as F; \ - # D(F('DUMP'), max_length=12288, cache_dir=None)" - # Use discover_shards and the same cache_dir training will use: a cache keyed on a - # different shard list is ignored, and torchrun rebuilds it from scratch. A ragged - # third-party dump also needs top_k= here, since the cache width is - # otherwise inferred from the first record and there is no recipe field for it. + # hours on a large corpus. Build it once before torchrun: see + # examples/speculative_decoding/doc/external_draft.md#teacher-policy-data sparse_data_path: # Where that cache lives. Defaults to a .sparse_cache beside the shards, which fails # on a read-only dump and makes concurrent runs over the same shards contend.