[Data] Derive off-policy loss masks from vLLM render boundaries

docs/cli/prepare_data.md

@@ -1,23 +1,35 @@
 # prepare_data.py
 
-Prepares data for speculator training by:
+Converts on-policy target-model data into the format consumed by speculator training. It accepts either:
 
-1. Applying chat template and tokenizing each sample
-2. Producing a loss/assistant mask for each sample
-3. Recording token frequency statistics
+1. Natural-language conversations whose assistant responses were produced by the target model.
+2. Speculator-format rows that already contain `input_ids` and `loss_mask`.
 
-The output is a processed dataset ready for online training or offline hidden states generation.
+For natural-language conversations, `prepare_data.py` asks the target model's vLLM `/render` endpoint to apply the serving chat template, tokenize each assistant turn, and derive its loss mask. Rendering only converts the data's representation: it does not generate responses or turn an arbitrary dataset into on-policy data.
+
+The output is ready for online training or offline hidden-state generation.
 
 ## Basic Usage
 
+Given a natural-language JSONL file such as:
+
+```json
+{"conversations":[{"role":"user","content":"Hello"},{"role":"assistant","content":"Hello! How can I help?"}]}
+```
+
+where the assistant response came from the target model:
+
 ```bash
 python scripts/prepare_data.py \
   --model meta-llama/Llama-3.1-8B-Instruct \
-  --data sharegpt \
+  --data ./on_policy_conversations.jsonl \
+  --render-endpoint http://localhost:8000 \
   --output ./training_data \
   --max-samples 5000
 ```
 
+`--render-endpoint` is not needed when every input row already contains `input_ids` and `loss_mask`.
+
 ## Arguments
 
 ### Model Arguments
@@ -30,19 +42,19 @@ python scripts/prepare_data.py \
 
 ### Data Arguments
 
-- **`--data`** (str, required, repeatable) Path to training data. Can be a HuggingFace dataset name or local path. Use multiple times to specify multiple datasets.
+- **`--data`** (str, required, repeatable) On-policy target-model data. Use a local JSON/JSONL file or directory, or an `hf:` dataset spec. Use multiple times to combine datasets.
 
-  Example: `--data sharegpt --data ./custom_data.jsonl`
+  Example: `--data ./target_responses.jsonl --data hf:my-org/more-target-responses`
 
-  The input conversation should be provided in the `conversations` column. Tool-calling datasets that include separate columns for tools are also supported, as demonstrated in [llamafactory/reason-tool-use-demo-1500](https://huggingface.co/datasets/llamafactory/reason-tool-use-demo-1500) and [interstellarninja/hermes_reasoning_tool_use](https://huggingface.co/datasets/interstellarninja/hermes_reasoning_tool_use).
+  Natural-language input uses a `conversations` column and requires `--render-endpoint`. Assistant responses must already have been produced by the target model. Tool-calling datasets may also include a separate `tools` column. Speculator-format input uses `input_ids` and `loss_mask`.
 
 - **`--seq-length`** (int, default: `8192`) Maximum sequence length for each sample. Longer samples will be truncated.
 
 - **`--max-samples`** (int, default: `None`) Maximum number of samples to process. If `None`, processes all samples.
 
 - **`--token-freq-path`** (str, default: `{output}/token_freq.pt`) Path to save token frequency distribution. Defaults to `token_freq.pt` in the output directory.
 
-- **`--assistant-pattern`** (str, default: `None`) Custom regex pattern for matching assistant responses. If not provided, auto-detected from chat template.
+- **`--render-endpoint`** (str, default: `None`) Base URL of the target model's running vLLM server (e.g. `http://localhost:8000`). The instance launched for hidden-state extraction ([launch_vllm.py](launch_vllm.md)) serves this too, so no second server is needed. Pass the base URL only: `/v1/chat/completions/render` is appended to it, so the `/v1`-suffixed form that [data_generation_offline.py](data_generation_offline.md) `--endpoint` takes will 404. Required for natural-language conversations; omit it when every input already contains `input_ids` and `loss_mask`.
 
 - **`--minimum-valid-tokens`** (int, default: `None`) Drop samples whose loss mask contains fewer than this many trainable tokens.
 
@@ -63,8 +75,9 @@ python scripts/prepare_data.py \
 ```bash
 python scripts/prepare_data.py \
   --model meta-llama/Llama-3.1-8B-Instruct \
-  --data sharegpt \
-  --data ./custom_conversations.jsonl \
+  --data ./target_responses_part1.jsonl \
+  --data ./target_responses_part2.jsonl \
+  --render-endpoint http://localhost:8000 \
   --output ./prepared_data \
   --seq-length 4096 \
   --max-samples 10000 \

docs/cli/response_regeneration.md

@@ -1,6 +1,6 @@
 # response_regeneration
 
-Regenerates assistant responses in existing datasets using a vLLM-served model. Given a dataset containing conversations (e.g., Magpie, UltraChat, GSM8K), this pipeline extracts conversation turns, regenerates each assistant response turn-by-turn against the model's own prior outputs, and produces pre-tokenized training samples. For multi-turn conversations, each turn conditions on the regenerated history, producing on-policy training data.
+Regenerates assistant responses in existing datasets using a vLLM-served model. Given a dataset containing conversations (e.g., Magpie, UltraChat, GSM8K), this pipeline extracts conversation turns, regenerates each assistant response turn-by-turn against the model's own prior outputs, and produces speculator-format training samples. For multi-turn conversations, each turn conditions on the regenerated history, producing on-policy training data.
 
 The pipeline consists of two scripts:
 
@@ -57,7 +57,7 @@ All other arguments are passed through to `script.py`.
 
 ## script.py
 
-Extracts conversation turns from a dataset, regenerates each assistant response turn-by-turn via a vLLM chat completion endpoint, and writes out pre-tokenized training samples with generation boundaries marked in the loss mask.
+Extracts conversation turns from a dataset, regenerates each assistant response turn-by-turn via a vLLM chat completion endpoint, and writes out speculator-format training samples with generation boundaries marked in the loss mask.
 
 ### Features
 
@@ -136,11 +136,11 @@ The text presets from the shared dataset registry (`DATASET_CONFIGS` in `specula
 | `open-perfectblend` | `mlabonne/open-perfectblend`                      | `train`       |
 | `hermes-fc`         | `NousResearch/hermes-function-calling-v1`         | `train`       |
 
-The registry's multimodal preset, `sharegpt4v_coco`, is **off-policy only** and `--dataset` rejects it. Its turns carry image content parts, which the Chat Completions API rejects, and the pre-tokenized output row has nowhere to keep pixel data. Use it with `prepare-data`.
+The registry's multimodal preset, `sharegpt4v_coco`, is rejected because this regeneration pipeline cannot send its image content or retain it in a speculator-format row. Generate target responses with a multimodal-capable workflow, save the resulting natural-language conversations, and convert them with `prepare_data.py`.
 
 ## Output Format
 
-Rows are pre-tokenized and ready for training: one row per target generation, holding the prompt the target conditioned on followed by the tokens it generated. The endpoint must support `return_token_ids`, which the script uses to read the generation boundary directly instead of re-tokenizing the text and recovering the boundary with a regex.
+Rows are in speculator format and ready for training: one row per target generation, holding the prompt the target conditioned on followed by the tokens it generated. The endpoint must support `return_token_ids`, which the script uses to read the generation boundary directly instead of re-tokenizing the text and recovering the boundary with a regex.
 
 ```json
 {
@@ -184,7 +184,7 @@ Rows are written only once a conversation finishes. A conversation that fails pa
 
 If a source row carries a `tools` schema, it is forwarded to the endpoint on every request and the target regenerates its own tool calls, which are supervised like any other generation.
 
-Tools are **not executed**. The target's *k*-th regenerated call is paired with the *k*-th cached tool result already present in the source row, spliced back as a `tool` message so the conversation can continue. This keeps the call tokens on-policy while the results stay off-policy.
+Tools are **not executed**. The target's *k*-th regenerated call is paired with the *k*-th cached tool result already present in the source row, spliced back as a `tool` message so the conversation can continue. Tool results are environment observations rather than policy outputs; all assistant and tool-call tokens are generated by the target model.
 
 A conversation stops early — keeping the rows completed so far — when the target emits a call that cannot be paired 1:1 with a cached result: it has exhausted the cached results, emitted parallel calls in a single generation, or called a different tool than the next cached result answers. Such conversations are counted under `truncated` in the progress bar.
 

docs/user_guide/tutorials/response_regeneration.md

@@ -1,6 +1,6 @@
 # Response Regeneration
 
-This tutorial walks you through regenerating assistant responses in an existing dataset using a target model served by vLLM. The resulting dataset pairs the original user prompts with freshly generated responses (on-policy data), and is the recommended starting point for speculator training: the drafter learns to predict what the target model actually generates, not what the dataset's original authors wrote. For multi-turn conversations, each assistant turn is regenerated sequentially against the model's own prior responses, keeping the entire history on-policy. Training directly on the dataset's original responses (off-policy) is a cheaper fallback, since it skips a full target-model pass over the data, but costs acceptance length at inference time.
+This tutorial walks you through regenerating assistant responses in an existing dataset using a target model served by vLLM. The resulting dataset pairs the original user prompts with freshly generated responses (on-policy data) for speculator training: the drafter learns to predict what the target actually generates, not what the dataset's original authors wrote. For multi-turn conversations, each assistant turn is regenerated sequentially against the model's own prior responses, keeping the entire assistant history on-policy.
 
 ## Overview
 
@@ -27,7 +27,7 @@ This will:
 
 1. Start a vLLM server with the specified model
 2. Extract conversation turns from the dataset and regenerate assistant responses turn-by-turn
-3. Save pre-tokenized results to a JSONL file (e.g., `magpie_Llama-3.3-70B-Instruct.jsonl`)
+3. Save speculator-format results to a JSONL file (e.g., `magpie_Llama-3.3-70B-Instruct.jsonl`)
 4. Stop the server
 
 ### Multi-GPU Configurations
@@ -72,13 +72,13 @@ For tool-calling datasets (e.g. `hermes-fc`), pass the model's `--tool-call-pars
   --dataset hermes-fc
 ```
 
-This is **semi-on-policy** tool-call regeneration: the target regenerates the tool-call tokens on-policy, but tools are not executed. The *i*-th cached tool result from the source data is spliced positionally after the target's *i*-th regenerated call.
+This regenerates all assistant and tool-call tokens on-policy, but does not execute tools. The *i*-th cached tool result from the source data is treated as an environment observation and spliced positionally after the target's *i*-th regenerated call.
 
 **Limitation:** parallel tool calls are under development; the turn is currently truncated to the first call.
 
 ## Step 2: Verify the Output
 
-The output is a JSONL file with one pre-tokenized row per target generation. `loss_mask` is `0` over the prompt the target conditioned on and `1` over the tokens it generated, so training needs no further masking:
+The output is a JSONL file with one speculator-format row per target generation. `loss_mask` is `0` over the prompt the target conditioned on and `1` over the tokens it generated, so training needs no further masking:
 
 ```json
 {

docs/user_guide/tutorials/train.md

@@ -98,41 +98,39 @@ Note: if you are using an experiment tracker (e.g. trackio, wandb, tensorboard,
 
 ## Step 1: Prepare Your Data
 
-**Recommended:** regenerate the dataset's responses with your target model first (see [Response Regeneration](response_regeneration.md)) and pass the resulting JSONL to `--data`. This on-policy data aligns the drafter with what the target actually generates. Using the dataset's original responses, as shown below, is a cheaper off-policy fallback that skips a full target-model pass over the data, at the cost of lower acceptance length. **For MTP this is not optional** -- it requires data generated by the target model itself.
+Speculator training data must contain responses produced by the target model. You can create it with [Response Regeneration](response_regeneration.md) or supply on-policy data from your own generation pipeline.
 
-First, preprocess your training dataset:
+Response Regeneration writes speculator-format rows containing `input_ids` and `loss_mask`, which `prepare_data.py` can use directly:
 
 ```bash
 # in speculators venv
 python scripts/prepare_data.py \
   --model Qwen/Qwen3-8B \
-  --data sharegpt \
+  --data ./target_responses.jsonl \
   --output ./output \
   --max-samples 5000 \
   --seq-length 8192
 ```
 
-For MTP, point `--data` at a regenerated dataset. You can produce one with [Response Regeneration](response_regeneration.md), or download a pre-regenerated one:
+If your generation pipeline saves natural-language conversations instead, start the target model's vLLM server as described in Step 2, then use its render endpoint to convert those responses into speculator format:
 
 ```bash
-# in speculators venv
-hf download \
-  inference-optimization/Qwen3.5-9B-responses gsm8k.jsonl \
-  --repo-type dataset \
-  --local-dir ./output/dataset
-
 python scripts/prepare_data.py \
-  --model Qwen/Qwen3.5-9B \
-  --data ./output/dataset/gsm8k.jsonl \
+  --model Qwen/Qwen3-8B \
+  --data ./on_policy_conversations.jsonl \
+  --render-endpoint http://localhost:8000 \
   --output ./output \
   --max-samples 5000 \
   --seq-length 8192
 ```
 
+The render endpoint applies the serving chat template, tokenizes each turn, and derives its loss mask. It does not generate responses or make a dataset on-policy, so the assistant responses must already come from the same target model and generation configuration used for training.
+
 **Parameters explained:**
 
 - `--model` - The target model you want to accelerate
-- `--data` - Dataset to use (built-in support for `sharegpt`, `ultrachat`. Otherwise provide a custom path to a jsonl file). Can be supplied multiple times to combine multiple datasets.
+- `--data` - On-policy target-model data, either natural-language `conversations` or speculator-format `input_ids` and `loss_mask`. Can be supplied multiple times to combine datasets.
+- `--render-endpoint` - Target model's vLLM base URL; required only for natural-language conversations.
 - `--output` - Where to save preprocessed data
 - `--max-samples` - Limit samples (optional, good for testing/getting started)
 - `--seq-length` - Maximum sequence length

pyproject.toml

@@ -45,6 +45,7 @@ dependencies = [
     "click",
     "datasets>=4.0.0,<=5.0.1",
     "hs-connectors",
+    "httpx",
     "huggingface-hub",
     "loguru>=0.7.2,<=0.7.3",
     "numpy>=2.0.0,<=2.4.6",

scripts/prepare_data.py

@@ -2,11 +2,17 @@
 """
 Prepare data for speculator training
 
-This script processes an input dataset and:
-1. Applies chat template + tokenizes each sample
-2. Produces a loss/assistant mask for each sample
+Accepted inputs contain responses produced by the target model, either as
+natural-language conversations or as speculator-format ``input_ids`` and
+``loss_mask`` rows. For natural-language input this script:
+
+1. Uses the target model's vLLM endpoint to render each conversation
+2. Derives a loss mask from each assistant-turn boundary
 3. Records token frequency statistics
 
+Rendering converts an existing on-policy conversation into speculator format.
+It does not generate responses or make an arbitrary conversation on-policy.
+
 The output of this script is:
 1. Processed dataset ready for online training or offline datagen in output_dir
 2. Token frequency statistics file at token_freq_path
@@ -17,7 +23,8 @@
 Usage:
     python prepare_data.py \
         --model meta-llama/Llama-3.1-8B-Instruct \
-        --data sharegpt \
+        --data ./on_policy_conversations.jsonl \
+        --render-endpoint http://localhost:8000 \
         --output ./training_data \
         --max-samples 5000
 """
@@ -101,7 +108,11 @@ def parse_args():
         type=str,
         action="append",
         required=True,
-        help="Path to training data (same as used in preprocessing)",
+        help=(
+            "On-policy target-model data as natural-language conversations or "
+            "speculator-format input_ids/loss_mask rows. Assistant responses "
+            "must come from the target model; this command does not generate them."
+        ),
     )
     parser.add_argument(
         "--seq-length",
@@ -125,12 +136,19 @@ def parse_args():
         ),
     )
     parser.add_argument(
-        "--assistant-pattern",
+        "--render-endpoint",
         type=str,
         default=None,
         help=(
-            "Custom regex pattern for matching assistant responses. "
-            "If not provided, auto-detected from chat template."
+            "Base URL of a running vLLM server (e.g. http://localhost:8000). "
+            "The instance launched for hidden-state extraction serves this "
+            "too, so no second server is needed. Pass the base URL only: "
+            "/v1/chat/completions/render is appended to it, so the "
+            "/v1-suffixed form that data_generation_offline.py --endpoint "
+            "takes will 404. Conversations are tokenized by that endpoint and "
+            "the loss mask is derived from the render boundary. Rendering does "
+            "not generate responses or make arbitrary data on-policy. Required "
+            "unless every --data input already contains input_ids and loss_mask."
         ),
     )
 
@@ -224,7 +242,7 @@ def main():
         seed=args.seed,
         max_samples=args.max_samples,
         token_freq_path=token_freq_path,
-        assistant_pattern=args.assistant_pattern,
+        render_endpoint=args.render_endpoint,
         minimum_valid_tokens=args.minimum_valid_tokens,
         allow_empty_output=args.allow_empty_output,
         trust_remote_code=args.trust_remote_code,

scripts/response_regeneration/script.py

@@ -26,8 +26,9 @@
 
 logger = logging.getLogger(__name__)
 
-# On-policy regeneration has no multimodal support yet; off-policy `prepare-data`
-# does, so these presets are gated here rather than dropped from the registry.
+# On-policy regeneration has no multimodal support yet. Users can generate
+# multimodal target responses externally and then convert those conversations
+# with `prepare_data.py`.
 MULTIMODAL_DATASETS = {"sharegpt4v_coco"}
 REGEN_DATASETS = [name for name in DATASET_CONFIGS if name not in MULTIMODAL_DATASETS]
 
@@ -36,8 +37,9 @@ def _dataset_choice(name: str) -> str:
     """Reject multimodal presets with a reason, not a bare invalid choice."""
     if name in MULTIMODAL_DATASETS:
         raise argparse.ArgumentTypeError(
-            f"{name!r} is multimodal; on-policy regeneration does not support "
-            "images yet. Use it off-policy with `prepare-data`."
+            f"{name!r} is multimodal; response regeneration does not support "
+            "images yet. Generate target-model responses with a multimodal-capable "
+            "workflow, then convert them with `prepare_data.py`."
         )
     return name
 
@@ -452,8 +454,11 @@ def build_boundary_sample(
 
 
 def _tool_result_message(tool_call: dict, content: str) -> dict[str, Any]:
-    """Build the ``tool`` message that feeds a cached (off-policy) result back to
-    the target, paired to the id of the call the target just generated."""
+    """Feed a cached tool result back so target-model generation can continue.
+
+    Tool results are environment observations rather than policy outputs; only
+    assistant tokens need to be regenerated by the target model.
+    """
     message: dict[str, Any] = {"role": "tool", "content": content}
     call_id = tool_call.get("id")
     if call_id: