Refactor(regen): emit generation boundary as loss mask for on-policy data, eliminate redundant re-tokenize roundtrip (vllm-project#729)

WindChimeRan · web-flow · commit 6030a44f2f25 · 2026-07-10T15:04:51.000Z
# Purpose The current on-policy response regeneration has re-tokenization roundtrip redundancy. This PR helps eliminate the roundtrip: ``` BEFORE AFTER regen -> text regen -> return_token_ids write conversations write input_ids + loss_mask (boundary) re-tokenize + REGEX/HF mask _preprocess_batch passthrough (no masking) vLLM -> hidden states vLLM -> hidden states ``` The masking machinery still exists (off-policy needs it until PR2); it is just no longer on the on-policy path. ## Concrete example — a 2-turn conversation Before (one row of text, masked later by regex over the whole rendering): ```json {"id": "sample_7", "conversations": [{"from":"human","value":"2+2?"}, {"from":"gpt","value":"4"}, {"from":"human","value":"3+3?"}, {"from":"gpt","value":"6"}]} ``` After (one row per assistant turn, mask already applied): ```json {"id":"sample_7_turn0", "input_ids":[<prompt>,<"4">], "loss_mask":[0,0,0,1,1], "conversations":[...]} {"id":"sample_7_turn1", "input_ids":[<prompt+history>,<"6">], "loss_mask":[0,...,0,1,1], "conversations":[...]} ``` `loss_mask` 0 is the prompt the target conditioned on; 1 is what it generated. No `{% generation %}` markers, no regex — just the boundary the serving engine reports. ## Why On-policy data never needed regex masking: we generated the response, so we already know where the assistant tokens are, and capturing that boundary is exact. It is also correct for multi-turn reasoning models, where each turn supervises the `<think>` it actually generated while history keeps only parsed content — matching inference. The regex path cannot reproduce this, because the same turn tokenizes differently as "generated" vs as "history". ## What changes (observable) | | Before | After | |---|---|---| | Output schema | `conversations` (text) | `input_ids` + `loss_mask` (+ `conversations`, review-only) | | Rows per conversation | 1 | 1 per assistant turn | | Mask made | downstream regex / HF | generation boundary | | Endpoint | text out | must support `return_token_ids` | `conversations` is retained as an inert human-review twin of `input_ids`: `_preprocess_batch` routes on `input_ids`/`loss_mask` (checked first) and drops `conversations` before training, so it is never re-masked. ## Tests `build_boundary_sample` produces the boundary mask; pre-tokenized rows pass through `_preprocess_batch` without a processor, and a review `conversations` field is ignored. Off-policy suite unchanged (56 passed / 2 skipped). ## Scope — step 1 of 5 ("renderer is the single source of truth") 1. (this PR) on-policy → generation boundary; stop feeding the regex path. (Please note that this PR is a concrete standalone PR. No speculative scaffolding. ) 2. To unify the chat template generation tag. we should vendor generation-tagged template from trl huggingface/trl#4879. 3. on-policy → we have to diverge a bit from pure on-policy regen: real inference has a thinking compression mechanism - newer models, e.g., GLM5.2, Qwen3.6 only keep the latest turn's thinking, and drop 0:n-1 thinking. But we need all thinking because that's where speculative decoding can helps. PR1 is correct (completely on policy, no divergence) but inefficient because it fan out multiturn to multi-row jsonl. PR3 can linearize them back to single-row, with some tradeoff. This is also the same solution with trl, llama-factory for multi-turn SFT. We can discuss further on the tradeoff when I have a draft PR. 4. off-policy → vLLM `/render` this path can be decoupled from the regex & hf path 5. delete the regex/HF masking once both producers are redirected. small followup PRs: `--turn-dropout` not compatible with this PR. need to run some tests on real data to see if we should improve it or remove it. ## Checklist I have filled in: - [x] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)". - [x] The test plan/results, such as providing test command and pasting the results. - [ ] (Optional) The necessary documentation update. - [x] I (a human) have written or reviewed the code in this pr to the best of my ability. --------- Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
diff --git a/docs/cli/response_regeneration.md b/docs/cli/response_regeneration.md
@@ -93,7 +93,7 @@ python scripts/response_regeneration/script.py --dataset magpie
 
 - **`--outfile`** (str, default: auto-generated) Output JSONL path. If not specified, auto-generated as `{dataset}_{model}.jsonl`.
 
-- **`--resume`** (flag) Skip rows already present in the output file (matched by uuid or index).
+- **`--resume`** (flag) Skip conversations already present in the output file (matched by `primary_id`: the row's `id`/`uuid` if it has one, otherwise a content hash).
 
 ### Full Example
 
@@ -117,40 +117,41 @@ python scripts/response_regeneration/script.py \
 
 ## Output Format
 
-Each line of the output JSONL file pairs the original user prompt with the newly generated response, in a conversations format compatible with fine-tuning:
+Rows are pre-tokenized and ready for training: one row per assistant turn, 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
 {
-  "id": "sample_0",
+  "id": "conv-abc_turn0",
+  "primary_id": "conv-abc",
+  "input_ids": [151644, 872, ...],
+  "loss_mask": [0, 0, ..., 1, 1],
   "conversations": [
-    {"from": "human", "value": "What is the capital of France?"},
-    {"from": "gpt", "value": "The capital of France is Paris."}
+    {"role": "user", "content": "What is the capital of France?"},
+    {"role": "assistant", "content": "The capital of France is Paris."}
   ],
   "metadata": {
     "idx": 0,
     "finish_reason": "stop",
-    "latency_s": 1.234,
     "usage": {...},
-    "endpoint": "http://127.0.0.1:8000/v1/chat/completions",
-    "reasoning_content": "..."
+    "endpoint": "http://127.0.0.1:8000/v1/chat/completions"
   }
 }
 ```
 
-- The `id` field uses the dataset's UUID if available, otherwise falls back to `sample_{idx}`.
-- The `reasoning_content` field in metadata is only included when the model provides reasoning content (e.g., with reasoning models).
+- `loss_mask` is `0` over the prompt and `1` over the generated tokens. This *is* the generation boundary, so training applies no further masking.
+- A conversation with N assistant turns yields N rows, each carrying the history before it. Turn `k`'s row is `{primary_id}_turn{k}`.
+- `primary_id` is the conversation's stable id, used by `--resume`. The row `id` is turn-suffixed and never matches it.
+- `conversations` is a human-readable twin of `input_ids` for review only. Training drops it.
 
-On error, the response conversation turn is omitted and an `error` field is included in metadata:
+Rows are written only after every turn of a conversation succeeds. A conversation that fails partway writes nothing to the output file and one row to a sibling error file instead (`--outfile out.jsonl` gives `out.errors.jsonl`), so `--resume` retries it whole:
 
 ```json
 {
-  "id": "sample_0",
-  "conversations": [
-    {"from": "human", "value": "What is the capital of France?"}
-  ],
+  "id": "conv-abc",
   "metadata": {
     "idx": 0,
     "error": "ConnectionError(...)",
+    "turns_completed": 1,
     "endpoint": "http://127.0.0.1:8000/v1/chat/completions"
   }
 }
diff --git a/docs/user_guide/tutorials/response_regeneration.md b/docs/user_guide/tutorials/response_regeneration.md
@@ -50,25 +50,29 @@ For larger models, use data parallelism and/or tensor parallelism:
 
 ## Step 2: Verify the Output
 
-The output is a JSONL file where each line contains the original prompt and the generated response in a conversations format compatible with fine-tuning:
+The output is a JSONL file with one pre-tokenized row per assistant turn. `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
 {
-  "id": "sample_0",
+  "id": "conv-abc_turn0",
+  "primary_id": "conv-abc",
+  "input_ids": [151644, 872, ...],
+  "loss_mask": [0, 0, ..., 1, 1],
   "conversations": [
-    {"from": "human", "value": "What is the capital of France?"},
-    {"from": "gpt", "value": "The capital of France is Paris."}
+    {"role": "user", "content": "What is the capital of France?"},
+    {"role": "assistant", "content": "The capital of France is Paris."}
   ],
   "metadata": {
     "idx": 0,
     "finish_reason": "stop",
-    "latency_s": 1.234,
     "usage": {...},
     "endpoint": "http://127.0.0.1:8000/v1/chat/completions"
   }
 }
 ```
 
+A conversation with N assistant turns produces N rows, so expect more lines than input conversations. `conversations` is a review-only twin of `input_ids`; training drops it.
+
 Check that the output looks correct:
 
 ```bash
diff --git a/scripts/response_regeneration/script.py b/scripts/response_regeneration/script.py
@@ -6,7 +6,6 @@
 import os
 import re
 import sys
-import time
 from typing import Any
 
 import aiohttp
@@ -187,11 +186,16 @@ def _primary_identifier(row: dict[str, Any]) -> str:
 
 
 def load_seen(path: str) -> set[str]:
-    """Load previously written output ids from the output file.
+    """Load previously completed conversation ids from the output file.
 
-    Each record stores its stable id under the top-level ``id`` (equal to the
-    row's :func:`_primary_identifier`), so a resumed run skips a row when its
-    recomputed id is already present.
+    A conversation fans out to one row per assistant turn, whose ``id`` carries a
+    ``_turn<N>`` suffix; the conversation's own :func:`_primary_identifier` is
+    kept alongside it as ``primary_id``. Resume keys on that, since the suffixed
+    ids never match a recomputed one. Rows are written only after every turn
+    succeeds, so one row is enough to mark the conversation done.
+
+    ``id`` is the fallback for output files written before the fan-out, where the
+    top-level ``id`` *was* the primary identifier.
     """
     seen: set[str] = set()
     if not os.path.isfile(path):
@@ -203,8 +207,11 @@ def load_seen(path: str) -> set[str]:
                 obj = json.loads(line)
             except json.JSONDecodeError:
                 continue
-            if _is_present(obj.get("id")):
-                seen.add(str(obj["id"]))
+            key = obj.get("primary_id")
+            if not _is_present(key):
+                key = obj.get("id")
+            if _is_present(key):
+                seen.add(str(key))
     return seen
 
 
@@ -270,6 +277,19 @@ async def _post_chat(
         return await response.json()
 
 
+def build_boundary_sample(
+    prompt_token_ids: list[int],
+    completion_token_ids: list[int],
+) -> tuple[list[int], list[int]]:
+    """Build one training sample: prompt (loss_mask 0) + generated tokens (1).
+
+    The generation boundary is the mask -- no ``{% generation %}`` markers, no regex.
+    """
+    input_ids = [*prompt_token_ids, *completion_token_ids]
+    loss_mask = [0] * len(prompt_token_ids) + [1] * len(completion_token_ids)
+    return input_ids, loss_mask
+
+
 async def worker(
     session: aiohttp.ClientSession,
     queue: "asyncio.Queue[dict[str, Any]]",
@@ -280,12 +300,10 @@ async def worker(
     progress,
     stats: dict[str, int],
 ):
-    """Worker that pulls conversations from the queue and regenerates them.
+    """Regenerate each queued conversation into pre-tokenized training samples.
 
-    Each user turn is sent to the endpoint with the freshly generated prefix
-    (system + regenerated history); the original assistant turns are discarded
-    and replaced by the model's responses. Single-turn rows are the degenerate
-    case of one user turn.
+    One sample per assistant turn: the prompt the target conditioned on
+    (loss_mask 0) followed by the tokens it generated (1).
     """
     while True:
         item = await queue.get()
@@ -295,28 +313,23 @@ async def worker(
 
         idx = item["idx"]
         turns = item["turns"]
+        conv_id = item["primary_id"]
 
-        # The API prefix (role/content) and the output conversation (from/value)
-        # are built in lockstep as we walk the turns.
         prefix: list[dict[str, Any]] = []
-        out_convs: list[dict[str, Any]] = []
-        finish_reasons: list[str | None] = []
-        usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
-        start_time = time.time()
+        samples: list[dict[str, Any]] = []
         try:
             for turn in turns:
                 if turn["role"] == "system":
                     prefix.append({"role": "system", "content": turn["content"]})
-                    out_convs.append({"from": "system", "value": turn["content"]})
                     continue
 
                 prefix.append({"role": "user", "content": turn["content"]})
-                out_convs.append({"from": "human", "value": turn["content"]})
 
                 payload = {
                     "model": args.model,
                     "messages": prefix,
                     "max_tokens": args.max_tokens,
+                    "return_token_ids": True,  # prompt_token_ids + completion token_ids
                 }
                 data = await _post_chat(
                     session,
@@ -326,53 +339,60 @@ async def worker(
                 )
 
                 choice = data["choices"][0]
-                message = choice["message"]
-                generated_text = message.get("content")
-                reasoning_content = message.get("reasoning_content")
-                if reasoning_content is None:
-                    reasoning_content = message.get("reasoning")
-                finish_reasons.append(choice.get("finish_reason"))
-                turn_usage = data.get("usage") or {}
-                for key in usage:
-                    usage[key] += turn_usage.get(key) or 0
-
-                # Empty content (e.g. truncated mid-reasoning) would corrupt the
-                # next turn's prefix and emit a null target; fail the conversation.
+                generated_text = choice["message"].get("content")
+
+                # Empty content corrupts the next turn's prefix; fail the conversation.
                 if not generated_text:
-                    raise ValueError(f"empty assistant content (turn {len(out_convs)})")
-
-                prefix.append({"role": "assistant", "content": generated_text})
-                gpt_turn = {"from": "gpt", "value": generated_text}
-                # Stored per turn: data prep reads it here, not top-level metadata.
-                if reasoning_content is not None:
-                    gpt_turn["reasoning_content"] = reasoning_content
-                out_convs.append(gpt_turn)
-
-            metadata = {
-                "idx": idx,
-                "finish_reasons": finish_reasons,
-                "latency_s": round(time.time() - start_time, 3),
-                "usage": usage,
-                "endpoint": endpoint,
-            }
+                    raise ValueError(f"empty assistant content (turn {len(samples)})")
+
+                prompt_token_ids = data.get("prompt_token_ids")
+                completion_token_ids = choice.get("token_ids")
+                if not prompt_token_ids or not completion_token_ids:
+                    raise ValueError(
+                        "endpoint returned no token ids; it must support "
+                        "return_token_ids"
+                    )
 
-            output = {
-                "id": item["primary_id"],
-                "conversations": out_convs,
-                "metadata": metadata,
-            }
-            out_fh.write(json.dumps(output, ensure_ascii=False) + "\n")
+                input_ids, loss_mask = build_boundary_sample(
+                    prompt_token_ids, completion_token_ids
+                )
+                # History keeps parsed content; the generated <think> is supervised
+                # in this turn's completion tokens above.
+                assistant_msg = {"role": "assistant", "content": generated_text}
+                samples.append(
+                    {
+                        "id": f"{conv_id}_turn{len(samples)}",
+                        # Conversation-level key for --resume; the row `id` is
+                        # turn-suffixed and would never match a recomputed one.
+                        "primary_id": conv_id,
+                        "input_ids": input_ids,
+                        "loss_mask": loss_mask,
+                        # Review-only twin of input_ids; ignored by training.
+                        "conversations": [*prefix, assistant_msg],
+                        "metadata": {
+                            "idx": idx,
+                            "finish_reason": choice.get("finish_reason"),
+                            "usage": data.get("usage") or {},
+                            "endpoint": endpoint,
+                        },
+                    }
+                )
+                prefix.append(assistant_msg)
+
+            # Written only after every turn succeeds, so any row in the output
+            # file means the whole conversation is done (see load_seen).
+            for sample in samples:
+                out_fh.write(json.dumps(sample, ensure_ascii=False) + "\n")
             out_fh.flush()
             stats["ok"] += 1
         except Exception as e:  # noqa: BLE001
-            # Failures go to a separate error file, not the training output; an
-            # in-band marker would be invisible (the pipeline drops metadata).
+            # Failures go to a separate error file, not the training output.
             error_output = {
-                "id": item["primary_id"],
-                "conversations": out_convs,
+                "id": conv_id,
                 "metadata": {
                     "idx": idx,
                     "error": repr(e),
+                    "turns_completed": len(samples),
                     "endpoint": endpoint,
                 },
             }
diff --git a/src/speculators/data_generation/preprocessing.py b/src/speculators/data_generation/preprocessing.py
@@ -507,6 +507,37 @@ def _parse_conv_tools(conv_tools: object, idx: int) -> list | None:
         return None
 
 
+def _passthrough_pretokenized(
+    examples: dict, max_length: int, minimum_valid_tokens: int | None = None
+) -> dict[str, list]:
+    """Carry pre-tokenized ``(input_ids, loss_mask)`` rows through, truncated only.
+
+    On-policy regeneration already applied the boundary as the mask, so these rows
+    need no chat-template rendering or regex span detection.
+    """
+    results: dict[str, list] = {"input_ids": [], "loss_mask": [], "seq_len": []}
+    for ids, mask in zip(examples["input_ids"], examples["loss_mask"], strict=True):
+        # `strict=True` only pairs the columns; a per-row skew would survive it and
+        # the collator packs each key independently, silently shifting the mask
+        # against the ids for every sample packed after this one.
+        if len(ids) != len(mask):
+            raise ValueError(
+                f"Pre-tokenized row shape mismatch: "
+                f"input_ids={len(ids)}, loss_mask={len(mask)}"
+            )
+        trimmed_ids = ids[:max_length]
+        trimmed_mask = mask[:max_length]
+        if (
+            minimum_valid_tokens is not None
+            and sum(trimmed_mask) < minimum_valid_tokens
+        ):
+            continue
+        results["input_ids"].append(torch.tensor(trimmed_ids, dtype=torch.long))
+        results["loss_mask"].append(torch.tensor(trimmed_mask, dtype=torch.long))
+        results["seq_len"].append(len(trimmed_ids))
+    return results
+
+
 def _preprocess_batch(
     examples: dict,
     processor: ProcessorLike,
@@ -517,6 +548,11 @@ def _preprocess_batch(
 ) -> dict[str, list]:
     """Process a batch of conversations into tokenized format with loss masks."""
 
+    # On-policy regeneration rows are already masked (boundary); pass them through
+    # instead of re-tokenizing and re-masking.
+    if "input_ids" in examples and "loss_mask" in examples:
+        return _passthrough_pretokenized(examples, max_length, minimum_valid_tokens)
+
     results: dict[str, list] = {"input_ids": [], "loss_mask": [], "seq_len": []}
     conversations: list[dict] = examples.get("conversations", [])
 
@@ -617,8 +653,21 @@ def build_eagle3_dataset(
                      conversation
         minimum_valid_tokens: Number of tokens to consider for a valid sample
     """
+    original_cols = dataset.column_names
+    # These rows carry the generation boundary as their mask, so _preprocess_batch
+    # passes them through: no chat template, no span detection, no turn dropout.
+    pretokenized = {"input_ids", "loss_mask"} <= set(original_cols)
+
+    if pretokenized:
+        log.info("Pre-tokenized rows: using their loss mask, skipping chat template")
+        if turn_dropout:
+            log.warning("turn_dropout does not apply to pre-tokenized rows; ignoring")
+        if assistant_pattern is not None:
+            log.warning(
+                "assistant_pattern does not apply to pre-tokenized rows; ignoring"
+            )
     # Detect and use provided assistant message pattern
-    if assistant_pattern is not None:
+    elif assistant_pattern is not None:
         log.info(f"Using custom assistant pattern: {str(assistant_pattern)[:80]}...")
     elif _supports_assistant_mask(processor):
         assistant_pattern = None  # Signal to use HF mask in _preprocess_batch
@@ -627,8 +676,6 @@ def build_eagle3_dataset(
         assistant_pattern = _detect_assistant_pattern(processor)
         log.info(f"Detected assistant pattern: {str(assistant_pattern)[:80]}...")
 
-    original_cols = dataset.column_names
-
     # Avoid CPU contention for MM processing:
     # https://github.com/vllm-project/vllm/pull/31879
     with (
diff --git a/tests/unit/scripts/test_response_regeneration.py b/tests/unit/scripts/test_response_regeneration.py
