From d673cb9d8f81cc6bddd24ad742daa03861543a00 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 4 Aug 2026 15:04:44 +0000 Subject: [PATCH 1/4] build(deps): update datasets requirement from <=5.0.0,>=4.0.0 to >=4.0.0,<=5.0.1 (#917) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Updates the requirements on [datasets](https://github.com/huggingface/datasets) to permit the latest version.
Release notes

Sourced from datasets's releases.

5.0.1

Bug fixes

Docs

New Contributors

... (truncated)

Commits

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 78e8ad5a3..365c11a76 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,7 @@ keywords = [ ] dependencies = [ "click", - "datasets>=4.0.0,<=5.0.0", + "datasets>=4.0.0,<=5.0.1", "hs-connectors", "huggingface-hub", "loguru>=0.7.2,<=0.7.3", From 7fffc8de306e22ba40b86ac674371ea9c44dca3d Mon Sep 17 00:00:00 2001 From: Ranran Date: Wed, 5 Aug 2026 12:42:50 -0500 Subject: [PATCH 2/4] [Data] Derive off-policy loss masks from vLLM render boundaries (#794) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # Purpose Derive the **off-policy** training loss mask from a **render boundary**, with the vLLM instance as the single tokenization source. For each assistant turn we render the prompt and the prompt-plus-turn via `/v1/chat/completions/render`, and supervise the tokens the second render adds past the first. **Then we deprecate the fallback hf tokenization & fallback regex paths** This is the same prompt/completion boundary the engine already reports for on-policy regeneration (#729), reconstructed from the chat template. One tokenizer now produces the loss mask, the target hidden states, and the serving prompt — no separate local tokenization that can drift from serving. Design write-up (worked examples, rationale): https://claude.ai/code/artifact/6cd3bf57-3d20-4a06-b27d-5b08e5fce6b4 # What changes - **Deleted:** the two-tier local masking path — the HF `{% generation %}` tag-mask probe and the regex span-detector — plus the local render adapters and `--assistant-pattern`. - **Added:** a vLLM `/render` client and the boundary derivation. - **Unchanged:** on-policy pre-tokenized rows (#729) still pass straight through. **Why `token_ids`, not the server mask:** `/render` can also return `assistant_tokens_mask`, but it's only populated when the template carries `{% generation %}` tags, which the common models don't ship. We request only `token_ids` and compute the mask from the boundary — so it works on every template, and isn't blocked on the unreleased `return_loss_mask` (`/render` itself shipped in v0.15.0). # Requirement `--render-endpoint` is now required for off-policy data (pre-tokenized input still works without it). Point it at the vLLM instance already used for hidden-state extraction, or a GPU-less `vllm launch render`. Single source of truth holds only if that instance's `tokenizer_mode` matches serving's. # How to review ~68% of the diff is pure deletion; the real surface is ~440 lines of source: | Bucket | What | How to read | |---|---|---| | **Deletion** (~1,560) | the old two-tier mask functions and the tests that covered them | skim | | **New** (~530) | `render_client.py`, the boundary functions, `test_render_boundary.py` | review in isolation | | **Rewiring** (~200) | `_preprocess_batch`, `build_eagle3_dataset`, `load_and_preprocess_dataset` switching to the boundary, plus the CLI swap | **the logic change** | The other preprocessing helpers are byte-identical to `main`. The delete-plus-rewire core is atomic; `render_client.py` is the one piece reviewable on its own. # Tests ```bash # Serve any model with a chat template vllm serve Qwen/Qwen3-0.6B --port 8000 # Derive masks from its render boundaries, on real data python scripts/prepare_data.py \ --model Qwen/Qwen3-0.6B \ --data sharegpt \ --render-endpoint http://localhost:8000 \ --output /tmp/prepared \ --max-samples 20 ``` It prints the derived mask inline (blue = trainable, grey = masked): ``` <|im_start|>user I have a location independent business that generates 1 miljon per year...<|im_end|> <|im_start|>assistant 1/5: A clear target market and understanding of their needs is essential...<|im_end|> ``` The boundary is the thing to check: everything from `` on is blue, while the user turn and the `<|im_start|>assistant` header stay grey. Or automated — launches its own vLLM and asserts the masks: ```bash pytest tests/e2e/smoke/test_render_boundary.py -v ``` # Deferred / follow-ups - **Concurrency** — the boundary does 2–3 renders per turn, today riding `datasets.map(num_proc)`. A batched/async client if throughput needs it. - **Multimodal** — prefix-stability not yet validated on an image model. - **Reasoning emission** — fan-out (on-policy, more rows) vs. a future single-row mode (cheaper, off-policy at turn seams); undecided. - **Render-failure handling** — a dead endpoint currently drains the dataset to empty per-row instead of aborting fast; should distinguish endpoint failure from a single bad conversation. - **Shared tool-call parsing** with regen (#750), once both paths speak vLLM messages. ## Checklist - [x] Purpose of the PR - [x] Test plan / results - [x] Documentation update (CLI doc updated; user-guide tutorial pending) - [x] I (a human) have reviewed this code to the best of my ability. --------- Signed-off-by: Ranran Haoran Zhang --- docs/cli/prepare_data.md | 37 +- docs/cli/response_regeneration.md | 10 +- .../tutorials/response_regeneration.md | 8 +- docs/user_guide/tutorials/train.md | 24 +- pyproject.toml | 1 + scripts/prepare_data.py | 36 +- scripts/response_regeneration/script.py | 17 +- .../data_generation/preprocessing.py | 718 +++++----- .../data_generation/render_client.py | 71 + tests/e2e/smoke/test_render_boundary.py | 91 ++ .../integration/datagen/test_preprocessing.py | 1168 ----------------- .../datagen/test_regex_patterns.py | 144 -- .../datagen/test_render_boundary.py | 204 +++ .../scripts/test_response_regeneration.py | 21 +- tests/unit/train/test_prepare_data.py | 2 +- 15 files changed, 768 insertions(+), 1784 deletions(-) create mode 100644 src/speculators/data_generation/render_client.py create mode 100644 tests/e2e/smoke/test_render_boundary.py delete mode 100644 tests/integration/datagen/test_regex_patterns.py create mode 100644 tests/integration/datagen/test_render_boundary.py diff --git a/docs/cli/prepare_data.md b/docs/cli/prepare_data.md index 90f35f41f..f0830540f 100644 --- a/docs/cli/prepare_data.md +++ b/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,11 +42,11 @@ 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. @@ -42,7 +54,7 @@ python scripts/prepare_data.py \ - **`--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 \ diff --git a/docs/cli/response_regeneration.md b/docs/cli/response_regeneration.md index cfbf6deb0..2a708623b 100644 --- a/docs/cli/response_regeneration.md +++ b/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. diff --git a/docs/user_guide/tutorials/response_regeneration.md b/docs/user_guide/tutorials/response_regeneration.md index d40994dd6..a352b9def 100644 --- a/docs/user_guide/tutorials/response_regeneration.md +++ b/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 { diff --git a/docs/user_guide/tutorials/train.md b/docs/user_guide/tutorials/train.md index de6f43e3e..c32f1428b 100644 --- a/docs/user_guide/tutorials/train.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 365c11a76..52a566dc9 100644 --- a/pyproject.toml +++ b/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", diff --git a/scripts/prepare_data.py b/scripts/prepare_data.py index cac9b5048..22ac9dec4 100644 --- a/scripts/prepare_data.py +++ b/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, diff --git a/scripts/response_regeneration/script.py b/scripts/response_regeneration/script.py index 80a65d798..51cf23cb1 100644 --- a/scripts/response_regeneration/script.py +++ b/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: diff --git a/src/speculators/data_generation/preprocessing.py b/src/speculators/data_generation/preprocessing.py index 45d3b3c3e..fd130daf1 100644 --- a/src/speculators/data_generation/preprocessing.py +++ b/src/speculators/data_generation/preprocessing.py @@ -1,32 +1,26 @@ -import bisect import json -import re from collections.abc import Callable from contextlib import nullcontext from pathlib import Path -from re import Pattern -from typing import cast +from typing import Literal, TypedDict import torch from datasets import Dataset as HFDataset from datasets import concatenate_datasets, load_dataset -from packaging.version import Version from transformers import ( AutoProcessor, - BatchEncoding, - BatchFeature, PreTrainedTokenizerBase, ProcessorMixin, ) -from transformers import __version__ as TRANSFORMERS_VERSION # noqa: N812 from speculators.data_generation.configs import DATASET_CONFIGS from speculators.data_generation.logging_utils import PipelineLogger +from speculators.data_generation.render_client import render_conversation from speculators.data_generation.torch_utils import set_default_torch_num_threads from speculators.train.vocab_mapping import save_token_frequency_distribution __all__ = [ - "build_eagle3_dataset", + "build_speculator_training_dataset", "load_and_preprocess_dataset", "load_raw_dataset", ] @@ -125,29 +119,6 @@ def _normalize_conversation( return normalized -def _adapt_part_for_hf(part: str | dict, processor: ProcessorLike): - if isinstance(part, str) and isinstance(processor, ProcessorMixin): - return {"type": "text", "text": part} - - return part - - -def _adapt_turn_for_hf(turn: dict, processor: ProcessorLike): - if isinstance(turn["content"], str): - if isinstance(processor, ProcessorMixin): - return turn | {"content": [_adapt_part_for_hf(turn["content"], processor)]} - - return turn - - return turn | { - "content": [_adapt_part_for_hf(part, processor) for part in turn["content"]] - } - - -def _adapt_conv_for_hf(normalized_conv: list[dict], processor: ProcessorLike): - return [_adapt_turn_for_hf(turn, processor) for turn in normalized_conv] - - def _adapt_part_for_vllm(part: str | dict): if isinstance(part, str): return {"type": "text", "text": part} @@ -198,283 +169,117 @@ def _adapt_conv_for_vllm(normalized_conv: list[dict]): return [_adapt_turn_for_vllm(turn) for turn in normalized_conv] -def _supports_assistant_mask(processor: ProcessorLike) -> bool: - """Check if processor truly supports HF assistant token mask. - - Must return a non-zero mask for a conversation containing an assistant message. - """ - # NOTE: Some models (e.g. Qwen3.5) require a user message in the conversation, - # even though this check only looks at the assistant turn - test_conv = _adapt_conv_for_hf( - [ - {"role": "user", "content": "test"}, - {"role": "assistant", "content": "test"}, - ], - processor, - ) - - try: - res_any = processor.apply_chat_template( - test_conv, - tokenize=True, - return_assistant_tokens_mask=True, - return_dict=True, - ) - res = cast("BatchEncoding | BatchFeature", res_any) - - # Check both singular and plural key names - mask = res.get("assistant_masks", res.get("assistant_mask")) - if mask is None: - return False - - # Verify the mask is not all zeros - return any(m == 1 for m in mask) - except (TypeError, ValueError, KeyError, AttributeError) as e: - log.warning(f"An error occurred when trying to return assistant mask: {e}") - return False +class BoundaryUnstableError(ValueError): + """The chat template is not prefix-stable at an assistant turn boundary.""" -def _detect_assistant_pattern(processor: ProcessorLike) -> str: - """Auto-detect the assistant message pattern from the processor's chat template. +class BoundaryRow(TypedDict): + input_ids: list[int] + loss_mask: list[int] + conv: list[dict] # prefix through this turn; multimodal rows re-send it - Uses multi-turn conversation but extracts pattern from the LAST assistant - message only. - """ - test_conv = _adapt_conv_for_hf( - [ - {"role": "user", "content": "USER_MSG_1"}, - {"role": "assistant", "content": "ASSISTANT_MSG_1"}, - {"role": "user", "content": "USER_MSG_2"}, - {"role": "assistant", "content": "ASSISTANT_MSG_2"}, - ], - processor, - ) - formatted = processor.apply_chat_template( - test_conv, tokenize=False, add_generation_prompt=False - ) - assert isinstance(formatted, str), "Expected string from apply_chat_template" - - # Find the START and END of both assistant messages - first_start = formatted.find("ASSISTANT_MSG_1") - first_end = first_start + len("ASSISTANT_MSG_1") - second_start = formatted.find("ASSISTANT_MSG_2") - second_end = second_start + len("ASSISTANT_MSG_2") - - if first_start == -1 or second_start == -1: - raise ValueError("Could not detect assistant messages in chat template") - - # Extract role marker from before the second assistant message - second_user_end = formatted.find("USER_MSG_2") + len("USER_MSG_2") - prefix = formatted[second_user_end:second_start] - - # Find where the assistant role marker starts - assistant_pos = prefix.rfind("assistant") - if assistant_pos != -1: - # Search for a tag start ('<' or '[') before 'assistant' - role_start = -1 - for char in ["<", "["]: - pos = prefix.rfind(char, 0, assistant_pos) - role_start = max(role_start, pos) - if role_start != -1: - role_marker = prefix[role_start:] - else: - role_marker = prefix[assistant_pos:] - else: - role_marker = prefix - - # Strip ... blocks from the role marker. Thinking model - # templates wrap assistant content in these tags, but the test messages - # can produce empty blocks (e.g. "\n\n\n") with reasoning models, - # which then get baked into the regex as literals. Removing them ensures - # that reasoning stays within the assistant content group. - role_marker = re.sub(r".*?\s*", "", role_marker, flags=re.DOTALL) - - # Determine the stable TURN-LEVEL suffix - suffix1 = formatted[first_end : formatted.find("USER_MSG_2")] - suffix2 = formatted[second_end:] - - # The stable suffix is the common prefix of these two tails - common_len = 0 - for c1, c2 in zip(suffix1, suffix2, strict=False): - if c1 == c2: - common_len += 1 - else: - break - suffix = suffix1[:common_len] - - if not suffix: - suffix = suffix1 if suffix1 else "\n" - - # Extract dynamic boundary marker from role_marker - boundary_match = re.search( - r"((<\|?[a-zA-Z0-9_]+[\|>]?)|(\[[a-zA-Z0-9_]+\]))", role_marker - ) - if boundary_match: - boundary = re.escape(boundary_match.group(1)) - lookahead_pattern = f"(?!{boundary})" - else: - # Fallback to hardcoded if no clear tag found - lookahead_pattern = r"(?!<\|start\|)" - - return ( - re.escape(role_marker) - + r"((?:" - + lookahead_pattern - + r".)*?)" - + re.escape(suffix) - ) - - -def _create_loss_mask_from_offsets( - text: str, - offsets: list[tuple[int, int]], - assistant_pattern: str | Pattern[str], +def _encode_render( + conv_prefix: list[dict], + render_endpoint: str, *, - # For logging - conv_idx: int | None = None, - max_length: int | None = None, -) -> torch.Tensor: - """Create loss mask by finding assistant response spans in formatted text.""" - loss_mask = torch.zeros(len(offsets), dtype=torch.bool) - - matches_found = 0 - token_starts = [offset[0] for offset in offsets] - - for match in re.finditer(assistant_pattern, text, re.DOTALL): - matches_found += 1 - - # Use group(1) to get only the assistant message content, - # excluding prefix/suffix markers - span_start_char = match.start(1) - span_end_char = match.end(1) - - start_idx = bisect.bisect_left(token_starts, span_start_char) - - for idx in range(max(0, start_idx - 1), len(offsets)): - token_start, token_end = offsets[idx] - if token_start >= span_end_char: - break - # Mark token as trainable if it overlaps with assistant span - if token_end > span_start_char and token_start < span_end_char: - loss_mask[idx] = 1 - - if matches_found == 0: - warning_msg = "No assistant response spans found in conversation" - if conv_idx is not None: - warning_msg += f" {conv_idx}" - - suggestion_msg = "" - if max_length is not None and len(offsets) == max_length: - suggestion_msg += ( - "Consider increasing --seq-length to avoid truncating " - "the assistant response." - ) + add_generation_prompt: bool, + tools: list[dict] | None = None, +) -> list[int]: + """Render a conversation prefix via the vLLM ``/render`` endpoint; return ids.""" + messages = _adapt_conv_for_vllm(conv_prefix) + return render_conversation( + render_endpoint, + messages, + add_generation_prompt=add_generation_prompt, + tools=tools, + ) - log.warning(f"{warning_msg}. {suggestion_msg}") - return loss_mask +def _common_prefix_len(a: list[int], b: list[int]) -> int: + length = 0 + for x, y in zip(a, b, strict=False): + if x != y: + break + length += 1 + return length -def _get_input_ids_loss_mask( +def _render_boundary_rows( normalized_conv: list[dict], - processor: ProcessorLike, + render_endpoint: str, max_length: int, - assistant_pattern: str | Pattern[str] | None, *, tools: list[dict] | None = None, - # For logging - conv_idx: int | None = None, -): - hf_conv = _adapt_conv_for_hf(normalized_conv, processor) - - if assistant_pattern is None: - # HF assistant token mask - encoded_any = processor.apply_chat_template( - hf_conv, - tokenize=True, - tools=tools, # type: ignore[arg-type] - add_generation_prompt=False, - return_assistant_tokens_mask=True, - return_dict=True, - ) - encoded = cast("BatchEncoding | BatchFeature", encoded_any) +) -> list[BoundaryRow]: + """Build one training row per assistant turn, masked at its render boundary. - # input IDs and loss mask - input_ids = encoded["input_ids"] - # HF uses 'assistant_masks' in recent versions - mask_key = ( - "assistant_masks" if "assistant_masks" in encoded else "assistant_mask" - ) - loss_mask = torch.tensor(encoded[mask_key], dtype=torch.long) + For assistant turn ``j``, the boundary is where the ``conv[:j+1]`` full render + extends the ``conv[:j]`` generation-prompt render: earlier tokens are context + (mask 0), later ones supervised (mask 1). If the generation prompt itself + diverges -- a pre-filled ```` scaffold vs recorded reasoning, as in + DeepSeek-R1 distills and Qwen3.5 with reasoning content -- the boundary falls + back to the common prefix, valid only if history agrees. - return input_ids, loss_mask + Every turn gets its own row, carrying the history re-rendered the way + inference would see it. Trailing non-assistant messages are dropped, and a + turn whose context alone fills ``max_length`` is skipped -- only that turn, + since a later one can fit again once the template drops history reasoning. - # Fallback: regex-based detection - assert assistant_pattern is not None, "Assistant pattern required for fallback" + Raises: + BoundaryUnstableError: the renders diverge inside history. + """ + rows: list[BoundaryRow] = [] - processor_kwargs: dict = { - "return_offsets_mapping": True, - "max_length": max_length, - "truncation": True, - "add_special_tokens": False, - } + for j, turn in enumerate(normalized_conv): + # j == 0 has no preceding context to bound against; keep it as context only. + if turn["role"] != "assistant" or j == 0: + continue - if isinstance(processor, ProcessorMixin): - if Version(TRANSFORMERS_VERSION) >= Version("5.4.0"): - encoded_any = processor.apply_chat_template( - hf_conv, - tokenize=True, - tools=tools, - add_generation_prompt=False, - return_dict=True, - processor_kwargs=processor_kwargs, - ) + prompt_ids = _encode_render( + normalized_conv[:j], + render_endpoint, + add_generation_prompt=True, + tools=tools, + ) + if len(prompt_ids) >= max_length: + # Not a break: templates that strip history reasoning (Qwen3, + # DeepSeek-R1) shrink the context, so a later turn can fit again. + continue + + full_ids = _encode_render( + normalized_conv[: j + 1], + render_endpoint, + add_generation_prompt=False, + tools=tools, + ) + if full_ids[: len(prompt_ids)] == prompt_ids: + boundary = len(prompt_ids) else: - encoded_any = processor.apply_chat_template( - hf_conv, - tokenize=True, - tools=tools, + # Generation prompt diverges (scaffold vs recorded reasoning): use + # the common prefix, valid only if history itself agrees (below). + boundary = _common_prefix_len(prompt_ids, full_ids) + hist_ids = _encode_render( + normalized_conv[:j], + render_endpoint, add_generation_prompt=False, - return_dict=True, - **processor_kwargs, + tools=tools, ) + if full_ids[: len(hist_ids)] != hist_ids or boundary < len(hist_ids): + raise BoundaryUnstableError( + f"prompt and full renders diverge inside history at " + f"assistant turn {j}; cannot derive a boundary loss mask" + ) - encoded = cast("BatchFeature", encoded_any) - - # Remove batch dimension - (input_ids,) = encoded["input_ids"] - (offsets,) = encoded["offset_mapping"] - - # MM placeholder tokens are inserted separate from chat template - formatted_text = processor.decode(input_ids) - assert isinstance(formatted_text, str) - else: - # More optimized flow for text-only processors (i.e. tokenizers) - formatted_text = processor.apply_chat_template( - hf_conv, - tokenize=False, - tools=tools, # type: ignore[arg-type] - add_generation_prompt=False, + rows.append( + { + "input_ids": full_ids, + "loss_mask": [0] * boundary + [1] * (len(full_ids) - boundary), + "conv": normalized_conv[: j + 1], + } ) - assert isinstance(formatted_text, str) - - # Tokenize and get offsets - encoded_any = processor(formatted_text, **processor_kwargs) - encoded = cast("BatchEncoding", encoded_any) - - input_ids = encoded["input_ids"] - offsets = encoded["offset_mapping"] - loss_mask = _create_loss_mask_from_offsets( - formatted_text, - offsets, - assistant_pattern, - conv_idx=conv_idx, - max_length=max_length, - ) - - return input_ids, loss_mask + return rows def _parse_conv_tools(conv_tools: object, idx: int) -> list | None: @@ -500,56 +305,165 @@ def _parse_conv_tools(conv_tools: object, idx: int) -> list | None: return None +def _render_conversation_rows( + conv: list[dict], + conv_tools: object, + idx: int, + render_endpoint: str, + max_length: int, +) -> list[BoundaryRow] | None: + """Render one valid conversation; return ``None`` when it is unusable.""" + if not conv or not isinstance(conv, list): + return None + + normalized_conv = _normalize_conversation(conv) + if not normalized_conv: + return None + + parsed_tools = _parse_conv_tools(conv_tools, idx) + try: + return _render_boundary_rows( + normalized_conv, + render_endpoint, + max_length, + tools=parsed_tools, + ) + # One row the render endpoint or boundary derivation can't handle must + # not kill the run. The failure modes can't be enumerated -- templates + # are swappable and raise arbitrary types -- so catch broadly and skip. + except Exception as e: + log.error(f"Failed to process conversation {idx}: {type(e).__name__}: {e}") + return [] + + +def _append_row( + results: dict[str, list], + input_ids: list[int], + loss_mask: list[int], + max_length: int, + minimum_valid_tokens: int | None, +) -> Literal["kept", "unsupervised", "filtered"]: + """Clip to the window, filter, and tensorize a row into ``results``. + + Returns "unsupervised" (no supervised tokens in-window), "filtered" (below + ``minimum_valid_tokens``), or "kept". + """ + input_ids = input_ids[:max_length] + loss_mask = loss_mask[:max_length] + num_valid_tokens = sum(loss_mask) + if num_valid_tokens == 0: + return "unsupervised" + if minimum_valid_tokens is not None and num_valid_tokens < minimum_valid_tokens: + return "filtered" + results["input_ids"].append(torch.tensor(input_ids, dtype=torch.long)) + results["loss_mask"].append(torch.tensor(loss_mask, dtype=torch.long)) + results["seq_len"].append(len(input_ids)) + return "kept" + + +def _append_boundary_rows( + results: dict[str, list], + rows: list[BoundaryRow], + max_length: int, + minimum_valid_tokens: int | None, +) -> tuple[int, int, int]: + """Append rendered rows and return kept, unsupervised, and clipped counts.""" + num_kept = 0 + num_unsupervised = 0 + num_clipped = 0 + + for row in rows: + status = _append_row( + results, + row["input_ids"], + row["loss_mask"], + max_length, + minimum_valid_tokens, + ) + num_unsupervised += status == "unsupervised" + # Kept-but-truncated only: a row clipped past its boundary reports + # as unsupervised above, and would otherwise be counted twice. + num_clipped += status == "kept" and len(row["input_ids"]) > max_length + if status == "kept": + num_kept += 1 + if "messages" in results: + results["messages"].append(_adapt_conv_for_vllm(row["conv"])) + + return num_kept, num_unsupervised, num_clipped + + +def _warn_seq_length(num_unsupervised: int, num_clipped: int) -> None: + """Warn when ``--seq-length`` cost supervision: all of it, or just the tail.""" + if num_unsupervised: + log.warning( + f"Dropped {num_unsupervised} rows with no supervised tokens. " + f"If unexpected, consider increasing --seq-length to avoid " + f"truncating assistant responses." + ) + if num_clipped: + log.warning( + f"Clipped {num_clipped} rows at --seq-length: the assistant turn is " + f"cut mid-response, so its tail and closing tokens are never " + f"supervised. These rows are kept and look healthy. Raise " + f"--seq-length to supervise responses whole -- reasoning traces " + f"routinely exceed 2048 tokens." + ) + + 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. + """Carry speculator-format ``(input_ids, loss_mask)`` rows through. - On-policy regeneration already applied the boundary as the mask, so these rows - need no chat-template rendering or regex span detection. + The producer already recorded which target-model tokens are supervised, so + these rows only need truncation and filtering. """ results: dict[str, list] = {"input_ids": [], "loss_mask": [], "seq_len": []} + num_unsupervised = 0 + num_clipped = 0 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. + # A per-row length skew survives strict= column pairing; the collator + # packs each key independently and would shift the mask silently. if len(ids) != len(mask): raise ValueError( - f"Pre-tokenized row shape mismatch: " + f"Speculator-format 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)) + status = _append_row(results, ids, mask, max_length, minimum_valid_tokens) + num_unsupervised += status == "unsupervised" + # Kept-but-truncated only: a row clipped past its boundary reports as + # unsupervised above, and would otherwise be counted twice. + num_clipped += status == "kept" and len(ids) > max_length + _warn_seq_length(num_unsupervised, num_clipped) return results def _preprocess_batch( examples: dict, - processor: ProcessorLike, + is_multimodal: bool, + render_endpoint: str | None, max_length: int, - assistant_pattern: str | Pattern[str] | None, minimum_valid_tokens: int | None = None, ) -> dict[str, list]: - """Process a batch of conversations into tokenized format with loss masks.""" + """Convert on-policy conversations or speculator-format rows for training.""" - # On-policy regeneration rows are already masked (boundary); pass them through - # instead of re-tokenizing and re-masking. + # Speculator-format rows already carry their supervision mask; pass them + # through instead of re-rendering. if "input_ids" in examples and "loss_mask" in examples: return _passthrough_pretokenized(examples, max_length, minimum_valid_tokens) + if render_endpoint is None: + raise ValueError( + "render_endpoint is required to convert natural-language " + "conversations to speculator training rows" + ) + results: dict[str, list] = {"input_ids": [], "loss_mask": [], "seq_len": []} - conversations: list[dict] = examples.get("conversations", []) + conversations: list[list[dict]] = examples.get("conversations", []) - # MM inputs must use Chat Completions API - if isinstance(processor, ProcessorMixin): + # MM inputs are extracted via the Chat Completions API, which needs the + # original messages -- token ids alone cannot carry the images. + if is_multimodal: results["messages"] = [] if not conversations: @@ -564,122 +478,105 @@ def _preprocess_batch( ) tools_col = None + num_unsupervised = 0 + num_clipped = 0 + num_convs_in = 0 + num_convs_empty = 0 + for idx, conv in enumerate(conversations): conv_tools = tools_col[idx] if tools_col is not None else None - - if not conv or not isinstance(conv, list): - continue - - # Normalize to standard format - normalized_conv = _normalize_conversation(conv) - if not normalized_conv: - continue - - parsed_tools = _parse_conv_tools(conv_tools, idx) - - try: - input_ids, loss_mask = _get_input_ids_loss_mask( - normalized_conv, - processor, - max_length=max_length, - assistant_pattern=assistant_pattern, - tools=parsed_tools, - conv_idx=idx, - ) - # Templates reject rows they cannot render with arbitrary types -- Mistral - # and Gemma raise jinja2's TemplateError, which subclasses Exception - # directly. One unrenderable row must not kill the run. - except Exception as e: - log.error( - f"Failed to process conversation {idx} " - f"(assistant_pattern={assistant_pattern is not None}): " - f"{type(e).__name__}: {e}" - ) + rows = _render_conversation_rows( + conv, + conv_tools, + idx, + render_endpoint, + max_length, + ) + if rows is None: continue - # Assert shapes match - assert len(input_ids) == len(loss_mask), ( - f"Shape mismatch: input_ids={len(input_ids)}, loss_mask={len(loss_mask)}" + num_convs_in += 1 + num_kept, row_unsupervised, row_clipped = _append_boundary_rows( + results, + rows, + max_length, + minimum_valid_tokens, ) + num_unsupervised += row_unsupervised + num_clipped += row_clipped + num_convs_empty += num_kept == 0 - # Bound both to max_length: a turn running past the window keeps only its - # in-window tokens, and input_ids/loss_mask stay aligned and bounded. - input_ids = input_ids[:max_length] - loss_mask = loss_mask[:max_length] - - # Filtering samples out with too few valid tokens - if minimum_valid_tokens is not None: - num_valid_tokens = int(loss_mask.sum().item()) - if num_valid_tokens < minimum_valid_tokens: - continue - - # Append to results - results["input_ids"].append(torch.tensor(input_ids, dtype=torch.long)) - results["loss_mask"].append(loss_mask) - results["seq_len"].append(len(input_ids)) - - if "messages" in results: - results["messages"].append(_adapt_conv_for_vllm(normalized_conv)) + _warn_seq_length(num_unsupervised, num_clipped) + if num_convs_empty: + log.warning( + f"{num_convs_empty}/{num_convs_in} conversations produced no training " + f"rows (no assistant turn with context, unstable template, or fully " + f"truncated)" + ) + num_rows = len(results["input_ids"]) + if num_rows > num_convs_in: + log.info(f"Per-turn fan-out: {num_convs_in} conversations -> {num_rows} rows") return results -def build_eagle3_dataset( +def build_speculator_training_dataset( dataset: HFDataset, processor: ProcessorLike, max_length: int = 2048, num_proc: int = 8, - assistant_pattern: str | Pattern[str] | None = None, + *, + render_endpoint: str | None = None, minimum_valid_tokens: int | None = None, ) -> HFDataset: - """Build EAGLE3 dataset by tokenizing conversations and creating loss masks. + """Build a speculator training dataset with render-boundary loss masks. - Uses the processor's built-in chat template via apply_chat_template. + Both accepted representations contain responses produced by the target + model. Natural-language conversations are tokenized by the vLLM ``/render`` + endpoint and masked at each assistant-turn boundary, fanning out to one row + per assistant turn. Rendering only converts representation; it does not + generate responses or make arbitrary data on-policy. Speculator-format rows + already carry ``input_ids`` and ``loss_mask`` and pass straight through. Args: - dataset: Raw dataset with conversations - processor: Processor with chat template support - max_length: Maximum sequence length - num_proc: Number of processes for parallel processing - assistant_pattern: Optional custom regex pattern for matching assistant - responses. If None, pattern will be auto-detected from - chat template. - minimum_valid_tokens: Number of tokens to consider for a valid sample + dataset: On-policy natural-language conversations, or speculator-format + rows containing ``input_ids`` and ``loss_mask``. + processor: Processor, used to detect multimodal inputs and to decode. + max_length: Maximum sequence length. + num_proc: Number of worker processes; each renders concurrently. + render_endpoint: Base URL of a vLLM server. Required unless the dataset + is already in speculator format. + minimum_valid_tokens: Minimum supervised tokens for a row to be kept. """ 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. + # These rows carry their supervision mask, so _preprocess_batch passes them + # through without rendering or boundary derivation. pretokenized = {"input_ids", "loss_mask"} <= set(original_cols) + # Multimodal rows keep their `messages` so the images survive to hidden-state + # extraction. Compute once here rather than pickling the heavyweight processor + # into every map worker just to recheck it. + is_multimodal = isinstance(processor, ProcessorMixin) if pretokenized: - log.info("Pre-tokenized rows: using their loss mask, skipping chat template") - 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 - 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 - log.info("Using HF assistant token mask for loss masking") + log.info("Speculator-format rows: using their loss mask, skipping render") + elif render_endpoint is None: + raise ValueError( + "render_endpoint is required to convert natural-language " + "conversations to speculator training rows. Pass --render-endpoint " + "pointing at the target model's vLLM server." + ) else: - assistant_pattern = _detect_assistant_pattern(processor) - log.info(f"Detected assistant pattern: {str(assistant_pattern)[:80]}...") + log.info("Deriving loss masks from vLLM render boundaries") # Avoid CPU contention for MM processing: # https://github.com/vllm-project/vllm/pull/31879 - with ( - set_default_torch_num_threads() - if isinstance(processor, ProcessorMixin) - else nullcontext() - ): + with set_default_torch_num_threads() if is_multimodal else nullcontext(): dataset = dataset.map( lambda examples: _preprocess_batch( examples, - processor, + is_multimodal, + render_endpoint, max_length, - assistant_pattern, minimum_valid_tokens, ), batched=True, @@ -835,14 +732,17 @@ def load_and_preprocess_dataset( seed: int = 0, max_samples: int | None = None, token_freq_path: Path | str = "./token_freq.pt", # noqa: S107 - assistant_pattern: str | None = None, + render_endpoint: str | None = None, minimum_valid_tokens: int | None = None, allow_empty_output: bool = False, trust_remote_code: bool = False, ) -> tuple[HFDataset, ProcessorLike]: - """Load, tokenize, and preprocess a dataset for EAGLE3 training. + """Load, tokenize, and preprocess a dataset for speculator training. - Uses the processor's built-in chat template via apply_chat_template. + Natural-language conversations containing target-model responses are + tokenized by a vLLM ``/render`` endpoint and masked at each assistant-turn + boundary. Speculator-format rows pass straight through. Rendering converts + representation; it does not generate or validate response provenance. Caching is handled automatically by HuggingFace datasets. Args: @@ -854,9 +754,9 @@ def load_and_preprocess_dataset( max_samples: Optional limit on number of samples token_freq_path: Path to save token frequency distribution cache_dir: Directory to cache HuggingFace datasets (optional) - assistant_pattern: Optional custom regex pattern for matching assistant - responses. If None, pattern will be auto-detected from - chat template. + render_endpoint: Base URL of a running vLLM server (e.g. + ``http://localhost:8000``) used to render conversations. Required + unless every dataset is already in speculator format. minimum_valid_tokens: Number of tokens to consider for a valid sample allow_empty_output: If True, allow returning an empty dataset instead of raising when no samples survive preprocessing. @@ -876,11 +776,8 @@ def load_and_preprocess_dataset( log.subsection("Loading processor") processor = load_processor(target_model_path, trust_remote_code=trust_remote_code) - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - raise ValueError( - f"Processor for {target_model_path} does not support chat templates. " - "Please use a model with a pre-configured chat template." - ) + if render_endpoint is not None: + log.info(f"Rendering conversations via vLLM endpoint: {render_endpoint}") processed_datasets = [] for train_data_path in train_data_paths: @@ -903,19 +800,16 @@ def load_and_preprocess_dataset( log.info(f"Loaded {len(raw_dataset)} samples") - preprocessed_dataset = build_eagle3_dataset( + preprocessed_dataset = build_speculator_training_dataset( dataset=raw_dataset, processor=processor, max_length=seq_length, num_proc=build_dataset_num_proc, - assistant_pattern=assistant_pattern, + render_endpoint=render_endpoint, minimum_valid_tokens=minimum_valid_tokens, ) - dropped = len(raw_dataset) - len(preprocessed_dataset) - if dropped: - log.warning( - f"Dropped {dropped}/{len(raw_dataset)} samples during preprocessing" - ) + if minimum_valid_tokens is not None: + log.info(f"Kept {len(preprocessed_dataset)} samples after filtering") processed_datasets.append(preprocessed_dataset) combined_dataset = concatenate_datasets(processed_datasets) diff --git a/src/speculators/data_generation/render_client.py b/src/speculators/data_generation/render_client.py new file mode 100644 index 000000000..822f82217 --- /dev/null +++ b/src/speculators/data_generation/render_client.py @@ -0,0 +1,71 @@ +"""Client for vLLM's ``/v1/chat/completions/render`` endpoint. + +Only ``token_ids`` are requested: the loss mask is derived from the boundary +between two renders, not the server's ``assistant_tokens_mask`` (which needs +``{% generation %}`` tags). Renders come from the vLLM instance the pipeline +already runs, so one tokenizer feeds the mask, hidden states, and serving. +""" + +from http import HTTPStatus + +import httpx + +from speculators.data_generation.vllm_client import InvalidResponseError, with_retries + +DEFAULT_RENDER_TIMEOUT = 30 + +# 4xx that mean "retry", not "your request is wrong". +TRANSIENT_STATUSES = frozenset( + {HTTPStatus.REQUEST_TIMEOUT, HTTPStatus.TOO_MANY_REQUESTS} +) + + +class RenderError(Exception): + """Non-200, retry-eligible response from the render endpoint.""" + + +@with_retries +def render_conversation( + endpoint: str, + messages: list[dict], + *, + add_generation_prompt: bool, + tools: list[dict] | None = None, + chat_template_kwargs: dict | None = None, + timeout: float = DEFAULT_RENDER_TIMEOUT, +) -> list[int]: + """POST to ``/v1/chat/completions/render`` and return the token ids. + + No truncation: boundary detection needs full lengths; over-length rows are + clipped downstream. + """ + url = f"{endpoint.rstrip('/')}/v1/chat/completions/render" + + body = { + "messages": messages, + "add_generation_prompt": add_generation_prompt, + } + if tools is not None: + body["tools"] = tools + if chat_template_kwargs is not None: + body["chat_template_kwargs"] = chat_template_kwargs + + resp = httpx.post(url, json=body, timeout=timeout) + + if ( + HTTPStatus.BAD_REQUEST <= resp.status_code < HTTPStatus.INTERNAL_SERVER_ERROR + and resp.status_code not in TRANSIENT_STATUSES + ): + # Deterministic client error (bad request, wrong URL) -- retrying wastes + # requests without changing the outcome. InvalidResponseError short- + # circuits @with_retries (see vllm_client._handle_retry_error). + raise InvalidResponseError( + f"Render endpoint returned {resp.status_code}: {resp.text}" + ) + if resp.status_code != HTTPStatus.OK: + raise RenderError(f"Render endpoint returned {resp.status_code}: {resp.text}") + + data = resp.json() + if "token_ids" not in data: + raise RenderError(f"Render endpoint response missing 'token_ids': {data}") + return data["token_ids"] diff --git a/tests/e2e/smoke/test_render_boundary.py b/tests/e2e/smoke/test_render_boundary.py new file mode 100644 index 000000000..6b9cb436d --- /dev/null +++ b/tests/e2e/smoke/test_render_boundary.py @@ -0,0 +1,91 @@ +"""E2E test: derive loss masks from a live vLLM ``/render`` endpoint. + +The only test that exercises the render transport for real -- the URL, the +``token_ids`` response key, and the boundary derived from what vLLM actually +returns. Everything else stubs ``_encode_render``, which cannot catch a change +in the endpoint contract or a chat template that stops being prefix-stable. + +Qwen3-0.6B is the model under test because its template rewrites history: it +injects a ```` scaffold into the current assistant turn and strips it +from past ones. That breaks the append-only chain, so every multi-turn +conversation fans out to one row per assistant turn -- the production default +for any thinking model. +""" + +import pytest +from datasets import Dataset as HFDataset +from transformers import AutoTokenizer + +from speculators.data_generation.preprocessing import build_speculator_training_dataset +from tests.conftest import requires_cuda +from tests.e2e.utils import launch_vllm_server_context + +MODEL = "Qwen/Qwen3-0.6B" +PORT = 8106 + +# One conversation covering three behaviours at once: a leading assistant turn +# (index 0 has no context to bound against, so it yields no row), two bounded +# assistant turns (one row each), and a trailing user turn (dropped). +CONVERSATION = [ + {"role": "assistant", "content": "Here is the first batch of data."}, + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "It is 4."}, + {"role": "user", "content": "And 3+3?"}, + {"role": "assistant", "content": "It is 6."}, + {"role": "user", "content": "thanks"}, +] + + +@pytest.mark.e2e +@pytest.mark.smoke +@requires_cuda +def test_render_boundary_masks_against_live_vllm(tmp_path): + tokenizer = AutoTokenizer.from_pretrained(MODEL) + + with launch_vllm_server_context( + MODEL, + PORT, + str(tmp_path / "hidden_states"), + max_model_len=1024, + gpu_memory_utilization=0.25, + ): + dataset = build_speculator_training_dataset( + HFDataset.from_dict({"conversations": [CONVERSATION]}), + tokenizer, + max_length=1024, + num_proc=1, + render_endpoint=f"http://localhost:{PORT}", + ) + + # Fan-out: the two bounded assistant turns each get a row. The leading + # assistant turn and the trailing user turn contribute neither. + assert len(dataset) == 2 + + supervised_per_row, context_per_row = [], [] + for row in dataset: + ids = row["input_ids"].tolist() + mask = row["loss_mask"].tolist() + + supervised = tokenizer.decode([t for t, m in zip(ids, mask, strict=True) if m]) + context = tokenizer.decode([t for t, m in zip(ids, mask, strict=True) if not m]) + supervised_per_row.append(supervised) + context_per_row.append(context) + + # The invariant the mask exists to hold: supervise the assistant, never + # the prompt. A leaked user turn shows up as its role header. + assert "<|im_start|>user" not in supervised + assert "<|im_start|>assistant" in context + + # The turn terminator is supervised -- the model must learn to stop. + assert "<|im_end|>" in supervised + + # Each row supervises its own turn and not the other's. + assert "It is 4." in supervised_per_row[0] + assert "It is 6." not in supervised_per_row[0] + assert "It is 6." in supervised_per_row[1] + + # Row 1 carries turn 1 as context, re-rendered the way inference sees it: + # Qwen3 strips the scaffold from history, which is exactly why the + # append-only chain breaks and this conversation fans out. + assert "It is 4." in context_per_row[1] + assert "" not in context_per_row[1] diff --git a/tests/integration/datagen/test_preprocessing.py b/tests/integration/datagen/test_preprocessing.py index df63db775..6a55d559c 100644 --- a/tests/integration/datagen/test_preprocessing.py +++ b/tests/integration/datagen/test_preprocessing.py @@ -3,33 +3,21 @@ """ import json -import re from typing import Any from unittest.mock import patch import pytest -import torch from datasets import Dataset as HFDataset from PIL import Image -from transformers import AutoTokenizer from speculators.data_generation.configs import ( DATASET_CONFIGS, _normalize_nemotron, ) from speculators.data_generation.preprocessing import ( - _adapt_conv_for_hf, _adapt_conv_for_vllm, - _create_loss_mask_from_offsets, - _detect_assistant_pattern, _load_hf_dataset, _normalize_conversation, - _preprocess_batch, - _supports_assistant_mask, - build_eagle3_dataset, - get_tokenizer, - load_and_preprocess_dataset, - load_processor, load_raw_dataset, ) @@ -160,68 +148,6 @@ def test_normalize_conversation_tool_calls_not_leaked(): assert "tool_call_id" not in result[1] -# Tests for _adapt_conv_for_hf -@pytest.mark.sanity -def test_adapt_conv_for_hf_text_only_processor(): - """ - Test converting from normalized conversation to HF format - with a text-only processor (i.e. tokenizer). - """ - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - - conv: list[dict] = [ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": ["Hello"]}, - {"role": "assistant", "content": "Hi!"}, - ] - result = _adapt_conv_for_hf(conv, processor) - - assert result == conv - - -@pytest.mark.sanity -def test_adapt_conv_for_hf_multimodal_processor(): - """ - Test converting from normalized conversation to HF format - with a multi-modal processor. - """ - processor = load_processor(MM_MODEL_REPO, trust_remote_code=True) - - conv: list[dict] = [ - { - "role": "system", - "content": "You are a helpful assistant.", - }, - { - "role": "user", - "content": ["Hello", {"type": "image", "path": "/path/to/img"}], - }, - { - "role": "assistant", - "content": "Hi!", - }, - ] - result = _adapt_conv_for_hf(conv, processor) - - assert result == [ - { - "role": "system", - "content": [{"type": "text", "text": "You are a helpful assistant."}], - }, - { - "role": "user", - "content": [ - {"type": "text", "text": "Hello"}, - {"type": "image", "path": "/path/to/img"}, - ], - }, - { - "role": "assistant", - "content": [{"type": "text", "text": "Hi!"}], - }, - ] - - # Tests for _adapt_conv_for_vllm @pytest.mark.sanity def test_adapt_conv_for_vllm_all_content_formats(): @@ -297,906 +223,6 @@ def test_adapt_conv_for_vllm_invalid_content_formats(): ) -# Tests for _detect_assistant_pattern -@pytest.mark.sanity -def test_detect_assistant_pattern_structure(): - """Test that the detected pattern has the correct regex structure.""" - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - pytest.skip("Processor does not support chat templates") - - pattern = _detect_assistant_pattern(processor) - - # Pattern should be a valid regex string - assert isinstance(pattern, str) - assert len(pattern) > 0 - - # Pattern should compile without errors - compiled = re.compile(pattern, re.DOTALL) - assert compiled is not None - - # Pattern should contain balanced parentheses - assert pattern.count("(") == pattern.count(")") - # Pattern should have at least one capture group (may use negative lookahead) - assert "(" in pattern, "Pattern should have a capture group for content" - - -@pytest.mark.sanity -def test_detect_assistant_pattern_correctly_identifies_assistant_vs_user(): - """Test that pattern correctly distinguishes assistant from user content.""" - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - pytest.skip("Processor does not support chat templates") - - # Get the pattern - pattern = _detect_assistant_pattern(processor) - - # Format a conversation manually to test the pattern - test_conv = [ - {"role": "user", "content": "USER_MSG"}, - {"role": "assistant", "content": "ASSISTANT_MSG"}, - ] - formatted: str = processor.apply_chat_template( # type: ignore[assignment] - test_conv, tokenize=False, add_generation_prompt=False - ) - - # Apply the pattern - matches = list(re.finditer(pattern, formatted, re.DOTALL)) - - # Should find exactly 1 match (the assistant message) - assert len(matches) == 1, f"Expected 1 match, got {len(matches)}" - - # The match should capture only ASSISTANT_MSG, not USER_MSG - captured_content = ( - matches[0].group(1) if matches[0].lastindex else matches[0].group(0) - ) - assert "ASSISTANT_MSG" in captured_content, ( - "Pattern should capture assistant content" - ) - assert "USER_MSG" not in captured_content, "Pattern should NOT capture user content" - - -@pytest.mark.sanity -def test_detect_assistant_pattern_extracts_correct_content(): - """Test that the pattern's capture group extracts only assistant message content.""" - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - pytest.skip("Processor does not support chat templates") - - pattern = _detect_assistant_pattern(processor) - - # Test with a multi-turn conversation - test_conv = [ - {"role": "user", "content": "First question"}, - {"role": "assistant", "content": "First answer"}, - {"role": "user", "content": "Second question"}, - {"role": "assistant", "content": "Second answer"}, - ] - - formatted: str = processor.apply_chat_template( # type: ignore [assignment] - test_conv, tokenize=False, add_generation_prompt=False - ) - - matches = list(re.finditer(pattern, formatted, re.DOTALL)) - - # Should match exactly 2 assistant messages - assert len(matches) == 2, f"Expected 2 assistant matches, got {len(matches)}" - - # First match should contain "First answer" but not questions - first_match = matches[0].group(0) - assert "First answer" in first_match - assert "First question" not in first_match - assert "Second question" not in first_match - - # Second match should contain "Second answer" but not questions - second_match = matches[1].group(0) - assert "Second answer" in second_match - assert "First question" not in second_match - assert "Second question" not in second_match - - -# Tests for _create_loss_mask_from_offsets - - -@pytest.mark.sanity -def test_create_loss_mask_simple(): - """Test creating loss mask for a simple case.""" - text = "User: Hello\nAssistant: Hi there!\nUser: How are you?\nAssistant: Good!" - pattern = r"Assistant: (.*?)(?=\n|$)" - - # Simulate token offsets (character positions) - offsets = [ - (0, 4), # "User" - (4, 5), # ":" - (6, 11), # "Hello" - (11, 12), # "\n" - (12, 21), # "Assistant" - (21, 22), # ":" - (23, 25), # "Hi" - (26, 31), # "there" - (31, 32), # "!" - (32, 33), # "\n" - (33, 37), # "User" - (37, 38), # ":" - (39, 42), # "How" - (43, 46), # "are" - (47, 51), # "you?" - (51, 52), # "\n" - (52, 61), # "Assistant" - (61, 62), # ":" - (63, 67), # "Good" - (67, 68), # "!" - ] - - mask = _create_loss_mask_from_offsets(text, offsets, pattern) - - assert len(mask) == len(offsets) - assert mask.dtype == torch.bool - - # Tokens in assistant responses should have mask = 1 - # "Hi there!" is at positions 6-8 (indices in offsets) - # "Good!" is at positions 18-19 - assert mask[6].item() == 1 # "Hi" - assert mask[7].item() == 1 # "there" - assert mask[8].item() == 1 # "!" - assert mask[18].item() == 1 # "Good" - assert mask[19].item() == 1 # "!" - - # User messages should have mask = 0 - assert mask[0].item() == 0 # "User" - assert mask[2].item() == 0 # "Hello" - - -@pytest.mark.sanity -def test_create_loss_mask_no_matches(): - """Test creating loss mask when no assistant patterns match.""" - text = "User: Hello\nUser: How are you?" - pattern = r"Assistant: (.*?)(?=\n|$)" - - offsets = [(0, 4), (4, 5), (6, 11)] - - mask = _create_loss_mask_from_offsets(text, offsets, pattern) - - # All zeros when no matches - assert torch.all(mask == 0) - - -@pytest.mark.sanity -def test_create_loss_mask_empty_offsets(): - """Test creating loss mask with empty offsets.""" - text = "User: Hello\nAssistant: Hi!" - pattern = r"Assistant: (.*?)(?=\n|$)" - - mask = _create_loss_mask_from_offsets(text, [], pattern) - - assert len(mask) == 0 - - -# Tests for _preprocess_batch - - -@pytest.mark.sanity -def test_preprocess_batch_basic(): - """Test preprocessing a basic batch of conversations.""" - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - pytest.skip("Processor does not support chat templates") - - examples = { - "conversations": [ - [ - {"role": "user", "content": "Hello"}, - {"role": "assistant", "content": "Hi there!"}, - ], - [ - {"role": "user", "content": "How are you?"}, - {"role": "assistant", "content": "I'm doing well!"}, - ], - ] - } - - assistant_pattern = _detect_assistant_pattern(processor) - results = _preprocess_batch( - examples, processor, max_length=512, assistant_pattern=assistant_pattern - ) - - assert "input_ids" in results - assert "loss_mask" in results - assert len(results["input_ids"]) == 2 - assert len(results["loss_mask"]) == 2 - - # Check that input_ids and loss_mask have same length for each example - for i in range(2): - assert len(results["input_ids"][i]) == len(results["loss_mask"][i]) - assert isinstance(results["input_ids"][i], torch.Tensor) - assert isinstance(results["loss_mask"][i], torch.Tensor) - - -@pytest.mark.sanity -def test_preprocess_batch_multimodal(tmp_path): - """Test preprocessing a batch of multimodal conversations.""" - processor = load_processor(MM_MODEL_REPO, trust_remote_code=True) - - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - pytest.skip("Processor does not support chat templates") - - img_path = str(tmp_path / "blank.png") - Image.new("RGB", (256, 256)).save(img_path) - - examples = { - "conversations": [ - [ - { - "role": "user", - "content": "Hello, how are you?", - }, - { - "role": "assistant", - "content": "I am a helpful assistant.", - }, - { - "role": "user", - "content": "What is the capital of France?", - }, - { - "role": "assistant", - "content": "The capital of France is Paris.", - }, - ], - [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "What is the difference between these two images?", - }, - {"type": "image", "path": img_path}, - {"type": "image", "path": img_path}, - ], - }, - { - "role": "assistant", - "content": [ - {"type": "text", "text": "They are the exact same image."}, - ], - }, - { - "role": "user", - "content": [ - {"type": "text", "text": "Why?"}, - ], - }, - { - "role": "assistant", - "content": [ - {"type": "text", "text": "They are both blank."}, - ], - }, - ], - ] - } - - assistant_pattern = _detect_assistant_pattern(processor) - results = _preprocess_batch( - examples, processor, max_length=2048, assistant_pattern=assistant_pattern - ) - - assert "input_ids" in results - assert "loss_mask" in results - assert len(results["input_ids"]) == 2 - assert len(results["loss_mask"]) == 2 - - # Check that input_ids and loss_mask have same length for each example - for i in range(2): - assert len(results["input_ids"][i]) == len(results["loss_mask"][i]) - assert isinstance(results["input_ids"][i], torch.Tensor) - assert isinstance(results["loss_mask"][i], torch.Tensor) - - -@pytest.mark.sanity -def test_preprocess_batch_empty_conversations(): - """Test preprocessing batch with no conversations.""" - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - pytest.skip("Processor does not support chat templates") - - examples: dict[str, list] = {"conversations": []} - assistant_pattern = _detect_assistant_pattern(processor) - results = _preprocess_batch( - examples, processor, max_length=512, assistant_pattern=assistant_pattern - ) - - assert results["input_ids"] == [] - assert results["loss_mask"] == [] - - -@pytest.mark.sanity -def test_preprocess_batch_skips_rows_a_strict_template_rejects(): - """A template that rejects a role sequence must skip the row, not kill the run. - - Mistral and Gemma templates validate role order and call ``raise_exception``, - which surfaces as ``jinja2.TemplateError``. Reproduced here with a minimal - strict template so the test needs no gated download. - """ - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - processor.chat_template = ( - "{% for m in messages %}" - "{% if (m['role'] == 'user') != (loop.index0 % 2 == 0) %}" - "{{ raise_exception('roles must alternate user/assistant/...') }}" - "{% endif %}" - "{{ m['role'] }}: {{ m['content'] }}\n" - "{% endfor %}" - ) - - examples = { - "conversations": [ - # Opens on an assistant turn, the shape sharegpt rows carry. - [ - {"role": "assistant", "content": "Sure"}, - {"role": "user", "content": "Hello"}, - {"role": "assistant", "content": "Hi there!"}, - ], - [ - {"role": "user", "content": "Hello"}, - {"role": "assistant", "content": "Hi there!"}, - ], - ] - } - - results = _preprocess_batch( - examples, processor, max_length=512, assistant_pattern=r"assistant: (.*?)\n" - ) - - # The clean row survives, and survives usable -- proving the skip is selective - # rather than the batch collapsing. - assert len(results["input_ids"]) == 1 - assert results["loss_mask"][0].sum() > 0 - - -@pytest.mark.sanity -def test_preprocess_batch_invalid_conversation(): - """Test preprocessing batch with invalid conversations.""" - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - pytest.skip("Processor does not support chat templates") - - examples = { - "conversations": [ - None, # Invalid - [], # Empty - [{"role": "user", "content": "Valid"}], # Valid - ] - } - - assistant_pattern = _detect_assistant_pattern(processor) - results = _preprocess_batch( - examples, processor, max_length=512, assistant_pattern=assistant_pattern - ) - - # Should only process the valid conversation - assert len(results["input_ids"]) <= 1 - assert len(results["loss_mask"]) <= 1 - - -@pytest.mark.sanity -def test_preprocess_batch_truncation(): - """Test that long sequences are truncated to max_length.""" - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - pytest.skip("Processor does not support chat templates") - - # Create a very long message - long_content = "word " * 1000 - - examples = { - "conversations": [ - [ - {"role": "user", "content": long_content}, - {"role": "assistant", "content": "Short reply"}, - ] - ] - } - - max_length = 100 - assistant_pattern = _detect_assistant_pattern(processor) - results = _preprocess_batch( - examples, processor, max_length=max_length, assistant_pattern=assistant_pattern - ) - - if len(results["input_ids"]) > 0: - # Should be truncated to max_length - assert len(results["input_ids"][0]) <= max_length - assert len(results["loss_mask"][0]) <= max_length - - -@pytest.mark.sanity -def test_preprocess_batch_uses_hf_assistant_mask(): - """Test that HF assistant token mask is used when supported.""" - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - pytest.skip("Processor does not support chat templates") - - # Skip test if assistant mask is not supported/functional for this processor - if not _supports_assistant_mask(processor): - pytest.skip("Processor does not support assistant token mask") - - examples = { - "conversations": [ - [ - {"role": "user", "content": "Hello"}, - {"role": "assistant", "content": "Hi there!"}, - ] - ] - } - - # Pass None to trigger masking path - results = _preprocess_batch( - examples, - processor, - max_length=128, - assistant_pattern=None, - ) - - assert "input_ids" in results - assert "loss_mask" in results - assert len(results["input_ids"]) == 1 - assert len(results["loss_mask"]) == 1 - - # Ensure at least some assistant tokens are trainable - assert torch.any(results["loss_mask"][0] == 1) - - -@pytest.mark.sanity -def test_preprocess_batch_falls_back_to_regex(): - """Test that preprocessing falls back to regex-based detection - when HF mask is unavailable. - """ - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - pytest.skip("Processor does not support chat templates") - - # Monkeypatch apply_chat_template to force HF mask failure - original_apply_chat_template = processor.apply_chat_template - - def patched_apply_chat_template(*args, **kwargs): - if kwargs.get("return_assistant_tokens_mask", False): - raise ValueError("Forcing fallback to regex path") - return original_apply_chat_template(*args, **kwargs) - - processor.apply_chat_template = patched_apply_chat_template # type: ignore [method-assign] - - examples = { - "conversations": [ - [ - {"role": "user", "content": "Hello"}, - {"role": "assistant", "content": "Hi!"}, - ] - ] - } - - assistant_pattern = _detect_assistant_pattern(processor) - - results = _preprocess_batch( - examples, - processor, - max_length=128, - assistant_pattern=assistant_pattern, - ) - - assert "input_ids" in results - assert "loss_mask" in results - assert len(results["input_ids"]) == 1 - assert len(results["loss_mask"]) == 1 - - # Regex path should still mark assistant tokens - assert torch.any(results["loss_mask"][0] == 1) - - -@pytest.mark.sanity -def test_preprocess_batch_minimum_valid_tokens_filters_regex_path(): - """Test that minimum_valid_tokens drops short samples on regex path.""" - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - pytest.skip("Processor does not support chat templates") - - examples = { - "conversations": [ - [ - {"role": "user", "content": "Hello"}, - {"role": "assistant", "content": "OK"}, - ] - ] - } - - assistant_pattern = _detect_assistant_pattern(processor) - - baseline = _preprocess_batch( - examples, - processor, - max_length=128, - assistant_pattern=assistant_pattern, - ) - - assert len(baseline["loss_mask"]) == 1 - valid_count = int(baseline["loss_mask"][0].sum().item()) - assert valid_count > 0 - - filtered = _preprocess_batch( - examples, - processor, - max_length=128, - assistant_pattern=assistant_pattern, - minimum_valid_tokens=valid_count + 1, - ) - - assert filtered["input_ids"] == [] - assert filtered["loss_mask"] == [] - assert filtered["seq_len"] == [] - - -@pytest.mark.sanity -def test_preprocess_batch_minimum_valid_tokens_keeps_boundary_case(): - """Test that a sample is kept when valid tokens equal the threshold.""" - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - pytest.skip("Processor does not support chat templates") - - examples = { - "conversations": [ - [ - {"role": "user", "content": "Explain speculative decoding."}, - { - "role": "assistant", - "content": ( - "Speculative decoding uses a draft model to propose tokens " - "that a verifier model then checks." - ), - }, - ] - ] - } - - assistant_pattern = _detect_assistant_pattern(processor) - - baseline = _preprocess_batch( - examples, - processor, - max_length=256, - assistant_pattern=assistant_pattern, - ) - - assert len(baseline["loss_mask"]) == 1 - valid_count = int(baseline["loss_mask"][0].sum().item()) - assert valid_count > 0 - - kept = _preprocess_batch( - examples, - processor, - max_length=256, - assistant_pattern=assistant_pattern, - minimum_valid_tokens=valid_count, - ) - - assert len(kept["input_ids"]) == 1 - assert len(kept["loss_mask"]) == 1 - assert int(kept["loss_mask"][0].sum().item()) == valid_count - - -# Tests for build_eagle3_dataset - - -@pytest.mark.sanity -def test_build_eagle3_dataset_basic(): - """Test building EAGLE3 dataset from a simple HuggingFace dataset.""" - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - pytest.skip("Processor does not support chat templates") - - # Create a simple dataset - data = { - "conversations": [ - [ - {"role": "user", "content": "Hello"}, - {"role": "assistant", "content": "Hi!"}, - ], - [ - {"role": "user", "content": "Goodbye"}, - {"role": "assistant", "content": "Bye!"}, - ], - ] - } - - dataset = HFDataset.from_dict(data) - result = build_eagle3_dataset(dataset, processor, max_length=512, num_proc=1) - - assert isinstance(result, HFDataset) - assert len(result) <= len(dataset) - - # Check that the dataset has the expected columns - if len(result) > 0: - assert "input_ids" in result.column_names - assert "loss_mask" in result.column_names - - -@pytest.mark.sanity -def test_build_eagle3_dataset_preserves_format(): - """Test that build_eagle3_dataset sets the correct format.""" - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - pytest.skip("Processor does not support chat templates") - - data = { - "conversations": [ - [ - {"role": "user", "content": "Test"}, - {"role": "assistant", "content": "Response"}, - ] - ] - } - - dataset = HFDataset.from_dict(data) - result = build_eagle3_dataset(dataset, processor, max_length=512, num_proc=1) - - # Dataset should be in torch format - assert result.format["type"] == "torch" - - -@pytest.mark.sanity -def test_build_eagle3_dataset_removes_original_columns(): - """Test that original columns are removed after processing.""" - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - pytest.skip("Processor does not support chat templates") - - data = { - "conversations": [ - [ - {"role": "user", "content": "Test"}, - {"role": "assistant", "content": "Response"}, - ] - ], - "extra_column": ["extra_data"], - } - - dataset = HFDataset.from_dict(data) - result = build_eagle3_dataset(dataset, processor, max_length=512, num_proc=1) - - # Original columns should be removed - if len(result) > 0: - assert "conversations" not in result.column_names - assert "extra_column" not in result.column_names - - -@pytest.mark.sanity -def test_build_eagle3_dataset_minimum_valid_tokens_filters_short_samples(): - """Test that build_eagle3_dataset removes samples below the token threshold.""" - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - pytest.skip("Processor does not support chat templates") - - short_conv = [ - {"role": "user", "content": "Hi"}, - {"role": "assistant", "content": "OK"}, - ] - long_conv = [ - {"role": "user", "content": "Explain speculative decoding."}, - { - "role": "assistant", - "content": ( - "Speculative decoding uses a draft model to propose multiple " - "candidate tokens that a stronger verifier then checks." - ), - }, - ] - - assistant_pattern = _detect_assistant_pattern(processor) - - short_baseline = _preprocess_batch( - {"conversations": [short_conv]}, - processor, - max_length=256, - assistant_pattern=assistant_pattern, - ) - long_baseline = _preprocess_batch( - {"conversations": [long_conv]}, - processor, - max_length=256, - assistant_pattern=assistant_pattern, - ) - - assert len(short_baseline["loss_mask"]) == 1 - assert len(long_baseline["loss_mask"]) == 1 - - short_count = int(short_baseline["loss_mask"][0].sum().item()) - long_count = int(long_baseline["loss_mask"][0].sum().item()) - - assert short_count > 0 - assert long_count > short_count - - threshold = short_count + 1 - - dataset = HFDataset.from_dict({"conversations": [short_conv, long_conv]}) - result = build_eagle3_dataset( - dataset, - processor, - max_length=256, - num_proc=1, - assistant_pattern=assistant_pattern, - minimum_valid_tokens=threshold, - ) - - assert isinstance(result, HFDataset) - assert len(result) == 1 - - remaining_valid_count = int(result[0]["loss_mask"].sum().item()) - assert remaining_valid_count >= threshold - - -# Tests for custom assistant pattern feature - - -@pytest.mark.sanity -def test_detect_assistant_pattern_thinking_model(): - """Test pattern detection with a real thinking model (Qwen3). - - Thinking templates wrap assistant content in ... tags. - The detection uses simple test messages that produce empty think blocks, - but the pattern must still match real conversations where the think block - contains substantial content. - """ - processor = load_processor("Qwen/Qwen3-8B", trust_remote_code=True) - pattern = _detect_assistant_pattern(processor) - - # Format a multi-turn conversation with thinking content injected - # directly into the formatted string (as it would appear in real data) - test_conv = [ - {"role": "user", "content": "What is 2+2?"}, - { - "role": "assistant", - "content": "The answer is 4.", - "reasoning_content": "We are adding 2 and 2.", - }, - {"role": "user", "content": "What is 3+3?"}, - { - "role": "assistant", - "content": "The answer is 6.", - "reasoning_content": "We are adding 3 and 3.", - }, - ] - formatted: str = processor.apply_chat_template( # type: ignore[assignment] - test_conv, tokenize=False, add_generation_prompt=False, enable_thinking=True - ) - - matches = list(re.finditer(pattern, formatted, re.DOTALL)) - assert len(matches) == 2, ( - f"Expected 2 matches, got {len(matches)}.\n" - f"Pattern: {pattern}\nText: {formatted}" - ) - - # Each match should capture its own assistant content, not the other's - assert "answer is 4" in matches[0].group(1) - assert "answer is 6" not in matches[0].group(1) - assert "answer is 6" in matches[1].group(1) - - # Neither match should contain user content - for m in matches: - assert "What is" not in m.group(1) - - # Reasoning content should be stripped from context turns - assert "2 and 2" not in matches[0].group(1) - - # Reasoning content should be present in the final turn - assert "3 and 3" in matches[1].group(1) - - -@pytest.mark.sanity -@pytest.mark.parametrize( - "thinking_content", - [ - "", - "Let me think step by step.\nThe user asked about France.", - ], - ids=["no_thinking", "with_thinking"], -) -def test_create_loss_mask_thinking_model(thinking_content): - """Test _create_loss_mask_from_offsets with Qwen3's thinking template. - - Verifies correct masking both with and without thinking content in the - block. - """ - processor = load_processor("Qwen/Qwen3-8B", trust_remote_code=True) - pattern = _detect_assistant_pattern(processor) - - # Build formatted text using the real chat template - conv = [ - {"role": "user", "content": "What is the capital of France?"}, - {"role": "assistant", "content": "Paris is the capital."}, - ] - if thinking_content: - conv[-1]["reasoning_content"] = thinking_content - formatted: str = processor.apply_chat_template( # type: ignore[assignment] - conv, - tokenize=False, - add_generation_prompt=False, - enable_thinking=bool(thinking_content), - ) - - # Tokenize with offsets - encoding = processor( - formatted, - return_offsets_mapping=True, - add_special_tokens=False, - ) - offsets = encoding["offset_mapping"] - - mask = _create_loss_mask_from_offsets(formatted, offsets, pattern) - - assert len(mask) == len(offsets) - assert mask.sum() > 0, "Loss mask should not be all zeros" - - # Decode masked vs unmasked regions - input_ids = torch.tensor(encoding["input_ids"]) - trainable_text = processor.decode(input_ids[mask == 1]) - masked_text = processor.decode(input_ids[mask == 0]) - - # Assistant response must be in the trainable region - assert "Paris is the capital" in trainable_text - - # User message must NOT be in the trainable region - assert "What is the capital of France" not in trainable_text - assert "What is the capital of France" in masked_text - - # Thinking content should be in the trainable region (part of assistant turn) - if thinking_content: - assert "step by step" in trainable_text - - -@pytest.mark.sanity -def test_build_eagle3_dataset_with_custom_pattern(): - """Test building dataset with custom assistant pattern.""" - processor = load_processor(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(processor, "apply_chat_template") or processor.chat_template is None: - pytest.skip("Processor does not support chat templates") - - data = { - "conversations": [ - [ - {"role": "user", "content": "Hello"}, - {"role": "assistant", "content": "Hi!"}, - ] - ] - } - - # Use a simple custom pattern - custom_pattern = r"<\|im_start\|>assistant\s*(.*?)<\|im_end\|>" - - dataset = HFDataset.from_dict(data) - result = build_eagle3_dataset( - dataset, processor, max_length=512, num_proc=1, assistant_pattern=custom_pattern - ) - - # Should successfully build dataset with custom pattern - assert isinstance(result, HFDataset) - assert len(result) > 0 - - # Tests for tool role and tool_calls / thinking field preservation @@ -1235,161 +261,6 @@ def test_normalize_conversation_preserves_tool_calls_field(): assert assistant_turn["tool_calls"] == tool_calls -@pytest.mark.sanity -def test_preprocess_batch_with_tools(): - """Test that tools from the dataset are forwarded to apply_chat_template. - - tools must be a list of JSON strings (one per conversation in the batch), - matching the HuggingFace datasets batched-column convention. - """ - tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(tokenizer, "apply_chat_template") or tokenizer.chat_template is None: - pytest.skip("Tokenizer does not support chat templates") - - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - - example_tools = [ - { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get the current weather", - "parameters": { - "type": "object", - "properties": {"location": {"type": "string"}}, - "required": ["location"], - }, - }, - } - ] - - conv = [ - {"role": "user", "content": "What's the weather in Prague?"}, - {"role": "assistant", "content": "Let me check."}, - ] - # tools must be a list — one JSON string per conversation - examples_with_tools = { - "conversations": [conv], - "tools": [json.dumps(example_tools)], - } - examples_without_tools = { - "conversations": [conv], - } - - assistant_pattern = _detect_assistant_pattern(tokenizer) - results_with = _preprocess_batch( - examples_with_tools, - tokenizer, - max_length=512, - assistant_pattern=assistant_pattern, - ) - results_without = _preprocess_batch( - examples_without_tools, - tokenizer, - max_length=512, - assistant_pattern=assistant_pattern, - ) - - assert "input_ids" in results_with - assert "loss_mask" in results_with - assert len(results_with["input_ids"]) == 1 - - # When the template renders tool definitions the sequence must be strictly - # longer than without tools. Skip the length check if the template silently - # ignores the tools kwarg (tool name absent from decoded output). - decoded_with = tokenizer.decode(results_with["input_ids"][0]) - if "get_weather" in decoded_with: - assert len(results_with["input_ids"][0]) > len( - results_without["input_ids"][0] - ), "Token sequence should be longer when tool definitions are included" - - -@pytest.mark.sanity -def test_preprocess_batch_with_invalid_tools_json(): - """Test that invalid JSON in the tools column is handled gracefully. - - The pipeline should warn and continue without tools rather than raising. - """ - tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(tokenizer, "apply_chat_template") or tokenizer.chat_template is None: - pytest.skip("Tokenizer does not support chat templates") - - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - - examples = { - "conversations": [ - [ - {"role": "user", "content": "Hello"}, - {"role": "assistant", "content": "Hi!"}, - ] - ], - "tools": ["this is not valid json"], - } - - assistant_pattern = _detect_assistant_pattern(tokenizer) - # Must not raise; the bad JSON entry is skipped with a warning - results = _preprocess_batch( - examples, tokenizer, max_length=512, assistant_pattern=assistant_pattern - ) - - assert "input_ids" in results - assert len(results["input_ids"]) == 1 - - -@pytest.mark.sanity -def test_preprocess_batch_tools_with_hf_assistant_mask(): - """Test that tools are forwarded when using the HF assistant token mask path.""" - tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_REPO, trust_remote_code=True) - - if not hasattr(tokenizer, "apply_chat_template") or tokenizer.chat_template is None: - pytest.skip("Tokenizer does not support chat templates") - - if not _supports_assistant_mask(tokenizer): - pytest.skip("Tokenizer does not support HF assistant token mask") - - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - - example_tools = [ - { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get the current weather", - "parameters": { - "type": "object", - "properties": {"location": {"type": "string"}}, - "required": ["location"], - }, - }, - } - ] - - examples = { - "conversations": [ - [ - {"role": "user", "content": "What's the weather in Prague?"}, - {"role": "assistant", "content": "Let me check."}, - ] - ], - "tools": [json.dumps(example_tools)], - } - - # assistant_pattern=None selects the HF mask path - results = _preprocess_batch( - examples, tokenizer, max_length=512, assistant_pattern=None - ) - - assert "input_ids" in results - assert "loss_mask" in results - assert len(results["input_ids"]) == 1 - assert torch.any(results["loss_mask"][0] == 1) - - @pytest.mark.sanity def test_normalize_conversation_tool_calls_with_empty_content(): """Test that an assistant turn with tool_calls and no text content is normalized.""" @@ -1593,42 +464,3 @@ def test_normalize_nemotron_builds_conversations(): # Tests for load_and_preprocess_dataset - - -@pytest.mark.sanity -def test_load_and_preprocess_dataset_shuffles_combined_datasets(tmp_path): - """Combined datasets must be shuffled before truncating to max_samples, - otherwise only samples from the first dataset are kept.""" - paths = [] - for marker in ("ALPHA", "BETA"): - path = tmp_path / f"{marker.lower()}.jsonl" - path.write_text( - "\n".join( - json.dumps( - { - "conversations": [ - {"role": "user", "content": f"Question {i}?"}, - {"role": "assistant", "content": f"{marker} answer {i}."}, - ] - } - ) - for i in range(12) - ) - ) - paths.append(str(path)) - - dataset, processor = load_and_preprocess_dataset( - TEXT_MODEL_REPO, - paths, - seq_length=128, - build_dataset_num_proc=1, - seed=0, - max_samples=12, - token_freq_path=tmp_path / "token_freq.pt", - trust_remote_code=True, - ) - - assert len(dataset) == 12 - tokenizer = get_tokenizer(processor) - decoded = [tokenizer.decode(row["input_ids"]) for row in dataset] - assert any("BETA" in text for text in decoded) diff --git a/tests/integration/datagen/test_regex_patterns.py b/tests/integration/datagen/test_regex_patterns.py deleted file mode 100644 index 3c68026ff..000000000 --- a/tests/integration/datagen/test_regex_patterns.py +++ /dev/null @@ -1,144 +0,0 @@ -""" -Tests for dynamic regex assistant pattern detection across different model families. -""" - -from re import Pattern - -import pytest -from loguru import logger as log -from packaging.version import Version -from PIL import Image -from transformers import ProcessorMixin -from transformers import __version__ as TRANSFORMERS_VERSION # noqa: N812 - -from speculators.data_generation.preprocessing import ( - _detect_assistant_pattern, - _preprocess_batch, - get_tokenizer, - load_processor, -) - -# Test models covering major template families -MODELS = [ - # Qwen/ChatML style - "Qwen/Qwen2-0.5B-Instruct", - # Llama-3 style (<|begin_of_text|>...) - "unsloth/llama-3-8b-Instruct", - # Mistral family ([INST] ... [/INST]) - "mistralai/Mistral-7B-Instruct-v0.2", - # Gemma style - "unsloth/gemma-2b-it", - # Phi-3 style - "microsoft/Phi-3-mini-4k-instruct", - # GPT-OSS - "openai/gpt-oss-20b", -] - -if Version(TRANSFORMERS_VERSION) >= Version("5.5.0"): - # Multimodal - MODELS.append("google/gemma-4-E2B-it") - - -@pytest.fixture(scope="module", params=MODELS) -def processor(request): - model_id = request.param - try: - # Using trust_remote_code=True for variety of templates - return load_processor(model_id, trust_remote_code=True) - except (TypeError, ValueError, KeyError, AttributeError, RuntimeError) as e: - pytest.skip(f"Failed to load processor for {model_id}: {e}") - - -def test_regex_detection_across_models(tmp_path, processor): - """ - Verify that _detect_assistant_pattern and _preprocess_batch (regex path) - work correctly for a variety of model families. - """ - tokenizer = get_tokenizer(processor) - model_name = tokenizer.name_or_path - log.info(f"Testing family: {model_name}") - - # 1. Detect pattern - try: - pattern = _detect_assistant_pattern(processor) - except (ValueError, RuntimeError) as e: - pytest.fail(f"Failed to detect assistant pattern for {model_name}: {e}") - - log.info(f"Detected pattern: {pattern}") - assert isinstance(pattern, (str, Pattern)), "Pattern must be str or regex object" - - # 2. Preprocess a simple multi-turn conversation using REGEX path - if isinstance(processor, ProcessorMixin): - img_path = str(tmp_path / "blank.png") - Image.new("RGB", (256, 256)).save(img_path) - - conversation = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "Hello, how are you?"}, - {"type": "image", "path": img_path}, - ], - }, - { - "role": "assistant", - "content": [ - {"type": "text", "text": "I am a helpful assistant."}, - ], - }, - { - "role": "user", - "content": [ - {"type": "text", "text": "What is the capital"}, - {"type": "image", "path": img_path}, - {"type": "text", "text": "of France?"}, - ], - }, - { - "role": "assistant", - "content": [ - { - "type": "text", - "text": "The capital of France is Paris.", - }, - ], - }, - ] - else: - conversation = [ - {"role": "user", "content": "Hello, how are you?"}, - {"role": "assistant", "content": "I am a helpful assistant."}, - {"role": "user", "content": "What is the capital of France?"}, - {"role": "assistant", "content": "The capital of France is Paris."}, - ] - - examples = {"conversations": [conversation]} - - # Regex path by passing the explicit pattern - results = _preprocess_batch( - examples, processor, max_length=2048, assistant_pattern=pattern - ) - - assert len(results["input_ids"]) == 1 - assert len(results["loss_mask"]) == 1 - - input_ids = results["input_ids"][0] - loss_mask = results["loss_mask"][0] - - # Verify basic properties - assert len(input_ids) == len(loss_mask) - assert loss_mask.sum() > 0, "Loss mask should not be all zeros" - - # 3. Qualitative check: Assistant content should be masked as 1 - trainable_tokens = input_ids[loss_mask == 1] - decoded_assistant = processor.decode(trainable_tokens) - - log.info(f"Decoded trainable regions: {decoded_assistant}") - - # It should at least contain parts of our assistant messages - assert "helpful assistant" in decoded_assistant - assert "Paris" in decoded_assistant - - # It should NOT contain user message content - assert "Hello" not in decoded_assistant - assert "France?" not in decoded_assistant diff --git a/tests/integration/datagen/test_render_boundary.py b/tests/integration/datagen/test_render_boundary.py new file mode 100644 index 000000000..fafcc8efb --- /dev/null +++ b/tests/integration/datagen/test_render_boundary.py @@ -0,0 +1,204 @@ +"""Unit tests for branches a live render endpoint cannot reach, and the client. + +The happy path -- fan-out and boundary derivation against a real chat template -- +is covered against a live vLLM server in ``tests/e2e/smoke/test_render_boundary``. +What is left here is only what a real server cannot produce on demand: the +scaffold fallback (needs a template that pre-fills ````), the unstable +guard (needs a template that rewrites history), and the client's error paths. +""" + +import time +from typing import cast + +import pytest +from datasets import Dataset as HFDataset + +from speculators.data_generation import preprocessing, render_client +from speculators.data_generation.preprocessing import ( + ProcessorLike, + build_speculator_training_dataset, +) +from speculators.data_generation.vllm_client import InvalidResponseError + +# Neither build path below reads the processor: the missing-endpoint guard +# raises before it is used, and speculator-format rows skip rendering entirely. +NO_PROCESSOR = cast("ProcessorLike", None) + + +def _conv(n: int) -> list[dict]: + """A conversation of ``n`` turns alternating user/assistant from user.""" + roles = ["user", "assistant"] * ((n + 1) // 2) + return [{"role": roles[i], "content": f"m{i}"} for i in range(n)] + + +def _patch_encode(monkeypatch, renders: dict[tuple[int, bool], list[int]]): + """Stub ``_encode_render`` to return crafted ids keyed by (prefix_len, gen).""" + + def fake(conv_prefix, render_endpoint, *, add_generation_prompt, tools=None): + return renders[(len(conv_prefix), add_generation_prompt)] + + monkeypatch.setattr(preprocessing, "_encode_render", fake) + + +# --------------------------------------------------------------------------- # +# _render_boundary_rows -- branches no real template reaches # +# --------------------------------------------------------------------------- # +def test_scaffold_lcp_fallback(monkeypatch): + # Load-bearing, not hypothetical: DeepSeek-R1 distills pre-fill `\n` + # in the generation prompt, and Qwen3.5 pre-fills an empty `` + # that recorded reasoning then contradicts. Both break the prefix and land + # here. Qwen3-0.6B (the e2e model) does not, so this stays a unit test. + # The generation prompt ends in a scaffold token the full render replaces; + # boundary falls back to the common prefix, valid because history agrees. + _patch_encode( + monkeypatch, + { + (1, True): [1, 2, 3, 77], # prompt with scaffold 77 + (1, False): [1, 2, 3], # history render + (2, False): [1, 2, 3, 4, 5], # full: diverges from prompt at idx 3 + }, + ) + rows = preprocessing._render_boundary_rows(_conv(2), "http://x", 100) + assert len(rows) == 1 + assert rows[0]["loss_mask"] == [0, 0, 0, 1, 1] + + +def test_boundary_unstable_raises(monkeypatch): + # Renders diverge inside history (not just the generation-prompt tail). + _patch_encode( + monkeypatch, + { + (1, True): [1, 2, 3], + (1, False): [1, 9, 9], # history disagrees with the full render + (2, False): [1, 5, 6, 7], + }, + ) + with pytest.raises(preprocessing.BoundaryUnstableError): + preprocessing._render_boundary_rows(_conv(2), "http://x", 100) + + +def test_over_length_turn_does_not_drop_later_turns(monkeypatch): + # Qwen3 strips `` from history once a later user turn arrives, so + # turn 3 can exceed the window while turn 5 fits again. + _patch_encode( + monkeypatch, + { + (1, True): [1, 2], # turn 1 context: fits + (2, False): [1, 2, 8, 9], + (3, True): [1] * 12, # turn 3 context: over max_length=10 + (5, True): [1, 2, 3, 4], # turn 5: reasoning stripped, fits again + (6, False): [1, 2, 3, 4, 7, 7], + }, + ) + rows = preprocessing._render_boundary_rows(_conv(6), "http://x", 10) + assert len(rows) == 2 # turns 1 and 5; only turn 3 is skipped + assert rows[0]["loss_mask"] == [0, 0, 1, 1] + assert rows[1]["loss_mask"] == [0, 0, 0, 0, 1, 1] + + +def test_over_length_first_turn_yields_no_rows(monkeypatch): + # No assistant message in the first turn's context, so nothing can be + # stripped: it is the smallest the conversation ever gets. + _patch_encode( + monkeypatch, + { + (1, True): [1] * 12, + (3, True): [1] * 15, + }, + ) + assert preprocessing._render_boundary_rows(_conv(4), "http://x", 10) == [] + + +# --------------------------------------------------------------------------- # +# _append_row -- clip / filter / keep # +# --------------------------------------------------------------------------- # +def test_append_row_statuses(): + results: dict[str, list] = {"input_ids": [], "loss_mask": [], "seq_len": []} + assert ( + preprocessing._append_row(results, [1, 2, 3], [0, 0, 0], 10, None) + == "unsupervised" + ) + assert preprocessing._append_row(results, [1, 2, 3], [0, 1, 1], 10, 3) == "filtered" + assert preprocessing._append_row(results, [1, 2, 3], [0, 1, 1], 10, 1) == "kept" + assert len(results["input_ids"]) == 1 + assert results["seq_len"] == [3] + + +# --------------------------------------------------------------------------- # +# render_client # +# --------------------------------------------------------------------------- # +class _Resp: + def __init__(self, status_code, payload=None, text=""): + self.status_code = status_code + self._payload = payload + self.text = text + + def json(self): + return self._payload + + +def test_render_conversation_missing_token_ids_raises(monkeypatch): + monkeypatch.setattr(render_client.httpx, "post", lambda *a, **k: _Resp(200, {})) + with pytest.raises(render_client.RenderError): + render_client.render_conversation( + "http://x", [], add_generation_prompt=False, max_retries=0 + ) + + +def test_render_conversation_client_error_not_retried(monkeypatch): + calls = [] + + def post(*a, **k): + calls.append(1) + return _Resp(400, {}, "bad request") + + monkeypatch.setattr(render_client.httpx, "post", post) + with pytest.raises(InvalidResponseError): + render_client.render_conversation("http://x", [], add_generation_prompt=False) + assert len(calls) == 1 # 4xx is deterministic: no retry + + +@pytest.mark.parametrize("status", [408, 429]) +def test_render_conversation_transient_status_is_retried(monkeypatch, status): + calls = [] + + def post(*a, **k): + calls.append(1) + return _Resp(status, {}, "slow down") + + monkeypatch.setattr(render_client.httpx, "post", post) + monkeypatch.setattr(time, "sleep", lambda _: None) # skip the backoff + with pytest.raises(render_client.RenderError): + render_client.render_conversation( + "http://x", [], add_generation_prompt=False, max_retries=2 + ) + assert len(calls) == 3 # initial attempt + 2 retries + + +# --------------------------------------------------------------------------- # +# build_speculator_training_dataset -- contracts that need no render at all # +# --------------------------------------------------------------------------- # +def test_build_speculator_training_dataset_requires_render_endpoint(): + data = { + "conversations": [ + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "yo"}, + ] + ] + } + with pytest.raises(ValueError, match="render_endpoint is required"): + build_speculator_training_dataset( + HFDataset.from_dict(data), NO_PROCESSOR, num_proc=1 + ) + + +def test_pretokenized_dataset_skips_render(): + # The load-bearing contract: speculator-format rows build without a render + # endpoint. Passthrough content (ids/mask) is covered by the regeneration + # tests in test_response_regeneration.py. + data = {"input_ids": [[1, 2, 3, 4]], "loss_mask": [[0, 0, 1, 1]]} + ds = build_speculator_training_dataset( + HFDataset.from_dict(data), NO_PROCESSOR, num_proc=1 + ) + assert len(ds) == 1 diff --git a/tests/unit/scripts/test_response_regeneration.py b/tests/unit/scripts/test_response_regeneration.py index 41a003264..8c426c2a9 100644 --- a/tests/unit/scripts/test_response_regeneration.py +++ b/tests/unit/scripts/test_response_regeneration.py @@ -189,7 +189,7 @@ def test_extract_conversation_no_usable_input_returns_empty(): # --------------------------------------------------------------------------- -# 2. The generation boundary is the loss mask; pre-tokenized rows pass through. +# 2. The generation boundary is the loss mask; speculator-format rows pass through. # --------------------------------------------------------------------------- @@ -200,8 +200,9 @@ def test_build_boundary_sample_is_the_mask(): def test_pretokenized_rows_pass_through_preprocessing(): - # A regen row reaches training already masked: no processor, no re-masking, - # and the review-only `conversations` field is dropped. + # A speculator-format regeneration row reaches training already masked: no + # processor, no re-masking, and the review-only `conversations` field is + # dropped. input_ids, loss_mask = regen.build_boundary_sample([10, 11, 12], [20, 21]) out = _preprocess_batch( { @@ -209,9 +210,9 @@ def test_pretokenized_rows_pass_through_preprocessing(): "loss_mask": [loss_mask], "conversations": [[{"role": "user", "content": "2+2?"}]], }, - processor=None, # type: ignore[arg-type] # passthrough never touches it + is_multimodal=False, # passthrough returns before this is read max_length=2048, - assistant_pattern=None, + render_endpoint=None, ) assert out["input_ids"][0].tolist() == input_ids assert out["loss_mask"][0].tolist() == loss_mask @@ -225,9 +226,9 @@ def test_pretokenized_passthrough_truncates_and_filters(): cut = regen.build_boundary_sample([1, 2, 3, 4], [5, 6]) # completion truncated off out = _preprocess_batch( {"input_ids": [kept[0], cut[0]], "loss_mask": [kept[1], cut[1]]}, - processor=None, # type: ignore[arg-type] # passthrough never touches it + is_multimodal=False, # passthrough returns before this is read max_length=4, - assistant_pattern=None, + render_endpoint=None, minimum_valid_tokens=1, ) assert [t.tolist() for t in out["input_ids"]] == [[1, 2, 3, 4]] @@ -241,9 +242,9 @@ def test_pretokenized_passthrough_rejects_length_mismatch(): with pytest.raises(ValueError, match="shape mismatch"): _preprocess_batch( {"input_ids": [[1, 2, 3, 4, 5]], "loss_mask": [[0, 0, 1]]}, - processor=None, # type: ignore[arg-type] # passthrough never touches it + is_multimodal=False, # passthrough returns before this is read max_length=2048, - assistant_pattern=None, + render_endpoint=None, ) @@ -854,7 +855,7 @@ def test_regenerate_truncates_on_tool_name_mismatch(): # --------------------------------------------------------------------------- -# 6. Every shared-registry preset works on-policy (off-policy parity). +# 6. Every text-only shared-registry preset works in on-policy regeneration. # --------------------------------------------------------------------------- diff --git a/tests/unit/train/test_prepare_data.py b/tests/unit/train/test_prepare_data.py index b0b5df9cf..5970b1e63 100644 --- a/tests/unit/train/test_prepare_data.py +++ b/tests/unit/train/test_prepare_data.py @@ -90,7 +90,7 @@ def _patch_empty_pipeline(monkeypatch: pytest.MonkeyPatch) -> None: lambda _path: (HFDataset.from_dict({"conversations": []}), None), ) monkeypatch.setattr( - preprocessing_module, "build_eagle3_dataset", lambda *a, **k: empty + preprocessing_module, "build_speculator_training_dataset", lambda *a, **k: empty ) monkeypatch.setattr( preprocessing_module, "save_token_frequency_distribution", lambda **k: None From fee672c4f2538b905f8663e8b66fa3068866ad4e Mon Sep 17 00:00:00 2001 From: Orestis Zambounis Date: Tue, 4 Aug 2026 12:12:10 +0000 Subject: [PATCH 3/4] feat(dflash): remap fused QKV weights during conversion for Laguna warm-start Laguna-style DFlash checkpoints use fused qkv_proj + per-head g_proj gating, while DFlashDraftModel expects separate q/k/v_proj (Qwen3-style). Add a _remap_weights step that splits fused QKV, drops incompatible keys (g_proj, aux_hidden_norms), and slices fc.weight when the source has more target layers than the model expects. Enables warm-starting DSpark training from published Laguna DFlash checkpoints. Signed-off-by: Claude Co-Authored-By: Claude Opus 4.6 Signed-off-by: Orestis Zambounis Signed-off-by: shotsan --- src/speculators/convert/dflash/converter.py | 52 ++++++++++++ tests/unit/convert/test_dflash_converter.py | 94 +++++++++++++++++++-- 2 files changed, 138 insertions(+), 8 deletions(-) diff --git a/src/speculators/convert/dflash/converter.py b/src/speculators/convert/dflash/converter.py index 5d76f1514..47837bdaf 100644 --- a/src/speculators/convert/dflash/converter.py +++ b/src/speculators/convert/dflash/converter.py @@ -142,6 +142,57 @@ def _build_config( speculators_config=speculators_config, ) + def _remap_weights( + self, + weights: dict[str, torch.Tensor], + config: DFlashSpeculatorConfig, + model: DFlashDraftModel, + ) -> dict[str, torch.Tensor]: + """Remap checkpoint weights to match DFlashDraftModel's state dict. + + Handles Laguna-style fused ``qkv_proj`` → separate ``q/k/v_proj``, + drops ``g_proj`` and ``aux_hidden_norms`` (no DFlash equivalent), and + slices ``fc.weight`` when the source has more target layers than needed. + """ + has_fused_qkv = any("qkv_proj" in k for k in weights) + if not has_fused_qkv: + return weights + + tl = config.transformer_layer_config + q_dim = tl.num_attention_heads * tl.head_dim + kv_dim = tl.num_key_value_heads * tl.head_dim + + remapped: dict[str, torch.Tensor] = {} + dropped: list[str] = [] + + for key, tensor in weights.items(): + if "qkv_proj" in key: + q, k, v = tensor.split([q_dim, kv_dim, kv_dim], dim=0) + remapped[key.replace("qkv_proj", "q_proj")] = q + remapped[key.replace("qkv_proj", "k_proj")] = k + remapped[key.replace("qkv_proj", "v_proj")] = v + elif ".g_proj." in key or key.startswith("aux_hidden_norms."): + dropped.append(key) + elif key == "fc.weight": + model_fc_dim = model.fc.in_features + if tensor.shape[1] > model_fc_dim: + logger.info( + f"Slicing fc.weight from {tensor.shape[1]} to {model_fc_dim}" + ) + remapped[key] = tensor[:, :model_fc_dim] + else: + remapped[key] = tensor + else: + remapped[key] = tensor + + if dropped: + logger.info(f"Dropped {len(dropped)} incompatible keys: {dropped}") + logger.info( + f"Remapped {sum(1 for k in weights if 'qkv_proj' in k)} fused qkv_proj " + f"→ separate q/k/v_proj" + ) + return remapped + def _save( self, config: DFlashSpeculatorConfig, @@ -151,6 +202,7 @@ def _save( model = DFlashDraftModel(config=config) body = {k: v for k, v in weights.items() if k not in ("t2d", "d2t")} + body = self._remap_weights(body, config, model) missing, unexpected = model.load_state_dict(body, strict=False) if unexpected: raise ValueError( diff --git a/tests/unit/convert/test_dflash_converter.py b/tests/unit/convert/test_dflash_converter.py index 959a56bb9..fbcbef281 100644 --- a/tests/unit/convert/test_dflash_converter.py +++ b/tests/unit/convert/test_dflash_converter.py @@ -1,31 +1,39 @@ -"""Unit tests for DFlashConverter config building.""" +"""Unit tests for DFlashConverter config building and weight remapping.""" from unittest.mock import patch import pytest +import torch from transformers import Qwen3Config from speculators.config import SpeculatorsConfig, VerifierConfig from speculators.convert.dflash.converter import DFlashConverter -from speculators.models.dflash import DFlashSpeculatorConfig +from speculators.models.dflash import DFlashDraftModel, DFlashSpeculatorConfig from speculators.proposals.greedy import GreedyTokenProposalConfig +_HIDDEN = 16 +_NUM_HEADS = 2 +_NUM_KV_HEADS = 1 +_HEAD_DIM = 8 +_Q_DIM = _NUM_HEADS * _HEAD_DIM # 16 +_KV_DIM = _NUM_KV_HEADS * _HEAD_DIM # 8 -def _tiny_dflash_config(): + +def _tiny_dflash_config(num_aux_layers=1): return DFlashSpeculatorConfig( transformer_layer_config=Qwen3Config( vocab_size=32, - hidden_size=16, + hidden_size=_HIDDEN, intermediate_size=32, num_hidden_layers=1, - num_attention_heads=2, - num_key_value_heads=1, - head_dim=8, + num_attention_heads=_NUM_HEADS, + num_key_value_heads=_NUM_KV_HEADS, + head_dim=_HEAD_DIM, max_position_embeddings=32, ), draft_vocab_size=32, block_size=4, - aux_hidden_state_layer_ids=[0], + aux_hidden_state_layer_ids=list(range(num_aux_layers)), mask_token_id=1, speculators_config=SpeculatorsConfig( algorithm="dflash", @@ -104,6 +112,76 @@ def test_missing_target_layer_ids_raises(self, mock_get_config): DFlashConverter()._build_config(source, "Qwen/Qwen3-8B", None) +class TestRemapWeights: + def _make_fused_weights(self): + qkv = torch.randn(_Q_DIM + 2 * _KV_DIM, _HIDDEN) + return { + "layers.0.self_attn.qkv_proj.weight": qkv, + "layers.0.self_attn.g_proj.weight": torch.randn(2, _HIDDEN), + "layers.0.self_attn.o_proj.weight": torch.randn(_HIDDEN, _Q_DIM), + "aux_hidden_norms.0.weight": torch.randn(_HIDDEN), + "fc.weight": torch.randn(_HIDDEN, _HIDDEN * 2), + "norm.weight": torch.randn(_HIDDEN), + } + + def test_splits_fused_qkv(self): + config = _tiny_dflash_config() + model = DFlashDraftModel(config=config) + weights = self._make_fused_weights() + qkv = weights["layers.0.self_attn.qkv_proj.weight"] + + remapped = DFlashConverter()._remap_weights(weights, config, model) + + assert "layers.0.self_attn.qkv_proj.weight" not in remapped + assert torch.equal(remapped["layers.0.self_attn.q_proj.weight"], qkv[:_Q_DIM]) + assert torch.equal( + remapped["layers.0.self_attn.k_proj.weight"], + qkv[_Q_DIM : _Q_DIM + _KV_DIM], + ) + assert torch.equal( + remapped["layers.0.self_attn.v_proj.weight"], + qkv[_Q_DIM + _KV_DIM :], + ) + + def test_drops_g_proj_and_aux_hidden_norms(self): + config = _tiny_dflash_config() + model = DFlashDraftModel(config=config) + remapped = DFlashConverter()._remap_weights( + self._make_fused_weights(), config, model + ) + assert not any("g_proj" in k for k in remapped) + assert not any("aux_hidden_norms" in k for k in remapped) + + def test_slices_fc_weight(self): + config = _tiny_dflash_config(num_aux_layers=1) + model = DFlashDraftModel(config=config) + weights = self._make_fused_weights() + # fc from checkpoint is wider than model expects + wide_fc = torch.randn(_HIDDEN, _HIDDEN * 3) + weights["fc.weight"] = wide_fc + + remapped = DFlashConverter()._remap_weights(weights, config, model) + assert remapped["fc.weight"].shape[1] == model.fc.in_features + assert torch.equal( + remapped["fc.weight"], wide_fc[:, : model.fc.in_features] + ) + + def test_passthrough_when_no_fused_qkv(self): + config = _tiny_dflash_config() + model = DFlashDraftModel(config=config) + weights = {"norm.weight": torch.randn(_HIDDEN)} + remapped = DFlashConverter()._remap_weights(weights, config, model) + assert remapped is weights + + def test_preserves_other_keys(self): + config = _tiny_dflash_config() + model = DFlashDraftModel(config=config) + weights = self._make_fused_weights() + remapped = DFlashConverter()._remap_weights(weights, config, model) + assert "layers.0.self_attn.o_proj.weight" in remapped + assert "norm.weight" in remapped + + class TestSave: def test_missing_draft_weights_raise(self, tmp_path): # No source weights: every draft-body weight (fc, norm, hidden_norm, From f13d60e64edc44c158f5b736a51b8c90f962b63d Mon Sep 17 00:00:00 2001 From: shotsan Date: Wed, 5 Aug 2026 12:34:32 -0700 Subject: [PATCH 4/4] fix(dflash): reject mixed fused and separate QKV projections Addresses PR review feedback to validate that separated q/k/v_proj keys don't already exist when splitting a fused qkv_proj. Also includes formatting fixes to resolve the quality checks failure. Signed-off-by: shotsan --- src/speculators/convert/dflash/converter.py | 15 ++++++++++++--- tests/unit/convert/test_dflash_converter.py | 14 +++++++++++--- 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/src/speculators/convert/dflash/converter.py b/src/speculators/convert/dflash/converter.py index 47837bdaf..841ce1b26 100644 --- a/src/speculators/convert/dflash/converter.py +++ b/src/speculators/convert/dflash/converter.py @@ -167,10 +167,19 @@ def _remap_weights( for key, tensor in weights.items(): if "qkv_proj" in key: + q_key = key.replace("qkv_proj", "q_proj") + k_key = key.replace("qkv_proj", "k_proj") + v_key = key.replace("qkv_proj", "v_proj") + for dest_key in (q_key, k_key, v_key): + if dest_key in remapped or dest_key in weights: + raise ValueError( + f"Mixed fused and separate projections: " + f"target key {dest_key} already exists." + ) q, k, v = tensor.split([q_dim, kv_dim, kv_dim], dim=0) - remapped[key.replace("qkv_proj", "q_proj")] = q - remapped[key.replace("qkv_proj", "k_proj")] = k - remapped[key.replace("qkv_proj", "v_proj")] = v + remapped[q_key] = q + remapped[k_key] = k + remapped[v_key] = v elif ".g_proj." in key or key.startswith("aux_hidden_norms."): dropped.append(key) elif key == "fc.weight": diff --git a/tests/unit/convert/test_dflash_converter.py b/tests/unit/convert/test_dflash_converter.py index fbcbef281..ad4d88b9f 100644 --- a/tests/unit/convert/test_dflash_converter.py +++ b/tests/unit/convert/test_dflash_converter.py @@ -162,9 +162,7 @@ def test_slices_fc_weight(self): remapped = DFlashConverter()._remap_weights(weights, config, model) assert remapped["fc.weight"].shape[1] == model.fc.in_features - assert torch.equal( - remapped["fc.weight"], wide_fc[:, : model.fc.in_features] - ) + assert torch.equal(remapped["fc.weight"], wide_fc[:, : model.fc.in_features]) def test_passthrough_when_no_fused_qkv(self): config = _tiny_dflash_config() @@ -181,6 +179,16 @@ def test_preserves_other_keys(self): assert "layers.0.self_attn.o_proj.weight" in remapped assert "norm.weight" in remapped + def test_rejects_mixed_fused_and_separate_projections(self): + config = _tiny_dflash_config() + model = DFlashDraftModel(config=config) + weights = self._make_fused_weights() + # Add a conflicting separate projection + weights["layers.0.self_attn.q_proj.weight"] = torch.randn(_Q_DIM, _HIDDEN) + + with pytest.raises(ValueError, match="Mixed fused and separate projections"): + DFlashConverter()._remap_weights(weights, config, model) + class TestSave: def test_missing_draft_weights_raise(self, tmp_path):