From 8389610cb27233927eaf0155c139c51bb9e90371 Mon Sep 17 00:00:00 2001 From: Ranran Haoran Zhang Date: Fri, 26 Jun 2026 00:41:48 -0500 Subject: [PATCH 01/12] implement renderer endpoint contract. blocked by vllm PR Signed-off-by: Ranran Haoran Zhang --- scripts/prepare_data.py | 12 ++ .../data_generation/preprocessing.py | 125 ++++++++++++------ .../data_generation/render_client.py | 48 +++++++ .../data_generation/test_render_client.py | 97 ++++++++++++++ 4 files changed, 244 insertions(+), 38 deletions(-) create mode 100644 src/speculators/data_generation/render_client.py create mode 100644 tests/unit/data_generation/test_render_client.py diff --git a/scripts/prepare_data.py b/scripts/prepare_data.py index 5c4624103..95d92a90f 100644 --- a/scripts/prepare_data.py +++ b/scripts/prepare_data.py @@ -105,6 +105,17 @@ def parse_args(): "per conversation for data augmentation." ), ) + parser.add_argument( + "--render-endpoint", + type=str, + default=None, + help=( + "URL of a vLLM render service " + "(e.g. http://localhost:8000). When set, tokenization is " + "delegated to the render endpoint instead of local HF " + "apply_chat_template. See RFC #652." + ), + ) # Output arguments parser.add_argument( @@ -190,6 +201,7 @@ def main(): turn_dropout=args.turn_dropout, minimum_valid_tokens=args.minimum_valid_tokens, trust_remote_code=args.trust_remote_code, + render_endpoint=args.render_endpoint, ) log.info("Done preparing data") diff --git a/src/speculators/data_generation/preprocessing.py b/src/speculators/data_generation/preprocessing.py index da9c3729c..49515b51f 100644 --- a/src/speculators/data_generation/preprocessing.py +++ b/src/speculators/data_generation/preprocessing.py @@ -23,11 +23,13 @@ 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_dataset_from_render", "load_and_preprocess_dataset", "load_raw_dataset", ] @@ -651,6 +653,62 @@ def build_eagle3_dataset( return dataset +def build_dataset_from_render( + dataset: HFDataset, + render_endpoint: str, + max_length: int = 2048, + chat_template_kwargs: dict | None = None, + minimum_valid_tokens: int | None = None, +) -> HFDataset: + """Tokenize via vLLM render endpoint instead of local HF apply_chat_template.""" + results: dict[str, list] = {"input_ids": [], "loss_mask": [], "seq_len": []} + + has_tools = "tools" in dataset.column_names + + for idx in range(len(dataset)): + row = dataset[idx] + conv = row.get("conversations") + if not conv or not isinstance(conv, list): + continue + + normalized = _normalize_conversation(conv) + if not normalized: + continue + + vllm_messages = _adapt_conv_for_vllm(normalized) + + conv_tools = _parse_conv_tools(row.get("tools"), idx) if has_tools else None + + rendered = render_conversation( + render_endpoint, + vllm_messages, + tools=conv_tools, + chat_template_kwargs=chat_template_kwargs, + max_length=max_length, + ) + + token_ids = rendered["token_ids"] + loss_mask = rendered["loss_mask"] + + if loss_mask is None: + raise NotImplementedError( + "render endpoint does not yet return loss_mask; pending vLLM PR" + ) + + if minimum_valid_tokens is not None: + num_valid = sum(1 for m in loss_mask if m == 1) + if num_valid < minimum_valid_tokens: + continue + + results["input_ids"].append(torch.tensor(token_ids, dtype=torch.long)) + results["loss_mask"].append(torch.tensor(loss_mask, dtype=torch.long)) + results["seq_len"].append(len(token_ids)) + + out = HFDataset.from_dict(results) + out.set_format(type="torch") + return out + + def load_raw_dataset( train_data_path: str, ) -> tuple[HFDataset, Callable[[dict], dict] | None]: @@ -709,32 +767,10 @@ def load_and_preprocess_dataset( turn_dropout: bool = False, minimum_valid_tokens: int | None = None, trust_remote_code: bool = False, + render_endpoint: str | None = None, + chat_template_kwargs: dict | None = None, ) -> tuple[HFDataset, ProcessorLike]: - """Load, tokenize, and preprocess a dataset for EAGLE3 training. - - Uses the processor's built-in chat template via apply_chat_template. - Caching is handled automatically by HuggingFace datasets. - - Args: - target_model_path: HuggingFace model ID or local path - train_data_path: Dataset name or path to JSON/JSONL file - seq_length: Maximum sequence length - build_dataset_num_proc: Number of processes for dataset building - seed: Random seed for shuffling - 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. - turn_dropout: If True, randomly keeps first N consecutive turns per - conversation - minimum_valid_tokens: Number of tokens to consider for a valid sample - trust_remote_code: If True, allows executing code from HF Hub. - - Returns: - Tuple of (preprocessed_dataset, processor) - """ + """Load, tokenize, and preprocess a dataset for EAGLE3 training.""" if minimum_valid_tokens is not None and minimum_valid_tokens < 0: raise ValueError("minimum_valid_tokens must be >= 0") log.section("Starting dataset preprocessing") @@ -746,7 +782,11 @@ 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: + if render_endpoint is not None: + log.info(f"Using render endpoint: {render_endpoint}") + elif ( + 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." @@ -773,18 +813,27 @@ def load_and_preprocess_dataset( log.info(f"Loaded {len(raw_dataset)} samples") - if turn_dropout: - log.info("Turn dropout enabled: randomly keeping N consecutive turns") - - preprocessed_dataset = build_eagle3_dataset( - dataset=raw_dataset, - processor=processor, - max_length=seq_length, - num_proc=build_dataset_num_proc, - assistant_pattern=assistant_pattern, - turn_dropout=turn_dropout, - minimum_valid_tokens=minimum_valid_tokens, - ) + if render_endpoint is not None: + preprocessed_dataset = build_dataset_from_render( + dataset=raw_dataset, + render_endpoint=render_endpoint, + max_length=seq_length, + chat_template_kwargs=chat_template_kwargs, + minimum_valid_tokens=minimum_valid_tokens, + ) + else: + if turn_dropout: + log.info("Turn dropout enabled: randomly keeping N consecutive turns") + + preprocessed_dataset = build_eagle3_dataset( + dataset=raw_dataset, + processor=processor, + max_length=seq_length, + num_proc=build_dataset_num_proc, + assistant_pattern=assistant_pattern, + turn_dropout=turn_dropout, + minimum_valid_tokens=minimum_valid_tokens, + ) if minimum_valid_tokens is not None: log.info(f"Kept {len(preprocessed_dataset)} samples after filtering") processed_datasets.append(preprocessed_dataset) diff --git a/src/speculators/data_generation/render_client.py b/src/speculators/data_generation/render_client.py new file mode 100644 index 000000000..ecef389d8 --- /dev/null +++ b/src/speculators/data_generation/render_client.py @@ -0,0 +1,48 @@ +"""Client for vLLM's /v1/chat/completions/render endpoint (RFC #652).""" + +from http import HTTPStatus + +import httpx + +from speculators.data_generation.vllm_client import with_retries + +DEFAULT_RENDER_TIMEOUT = 30 + + +class RenderError(Exception): + """Non-200 response from the render endpoint.""" + + +@with_retries +def render_conversation( + endpoint: str, + messages: list[dict], + *, + tools: list[dict] | None = None, + chat_template_kwargs: dict | None = None, + max_length: int | None = None, + timeout: float = DEFAULT_RENDER_TIMEOUT, +) -> dict: + """POST to /v1/chat/completions/render and return token_ids + loss_mask.""" + url = f"{endpoint.rstrip('/')}/v1/chat/completions/render" + + body: dict = { + "messages": messages, + "add_generation_prompt": False, + } + if tools is not None: + body["tools"] = tools + if chat_template_kwargs is not None: + body["chat_template_kwargs"] = chat_template_kwargs + if max_length is not None: + body["max_tokens"] = max_length + + resp = httpx.post(url, json=body, timeout=timeout) + + if resp.status_code != HTTPStatus.OK: + raise RenderError( + f"Render endpoint returned {resp.status_code}: {resp.text[:500]}" + ) + + data = resp.json() + return {"token_ids": data["token_ids"], "loss_mask": data.get("loss_mask")} diff --git a/tests/unit/data_generation/test_render_client.py b/tests/unit/data_generation/test_render_client.py new file mode 100644 index 000000000..335793e47 --- /dev/null +++ b/tests/unit/data_generation/test_render_client.py @@ -0,0 +1,97 @@ +"""Tests for render_client.py.""" + +import json + +import httpx +import pytest + +import speculators.data_generation.render_client as render_mod +from speculators.data_generation.render_client import ( + RenderError, + render_conversation, +) + +MOCK_TOKEN_IDS = [ + 151644, + 872, + 198, + 9707, + 151645, + 198, + 151644, + 77091, + 198, + 13048, + 0, + 151645, +] + +MESSAGES = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi!"}, +] + + +def _make_transport(*, token_ids=None, loss_mask=None, status_code=200): + def handler(request: httpx.Request) -> httpx.Response: + if status_code != 200: + return httpx.Response(status_code, text="error from mock") + body = {"token_ids": token_ids or MOCK_TOKEN_IDS} + if loss_mask is not None: + body["loss_mask"] = loss_mask + return httpx.Response(200, json=body) + + return httpx.MockTransport(handler) + + +def _capturing_transport(): + """Return (transport, captured_body_dict).""" + captured = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured.update(json.loads(request.content)) + return httpx.Response(200, json={"token_ids": MOCK_TOKEN_IDS}) + + return httpx.MockTransport(handler), captured + + +def _call_render(transport, **kwargs): + original_post = httpx.post + + def mock_post(url, **kw): + return httpx.Client(transport=transport).request("POST", url, **kw) + + render_mod.httpx.post = mock_post # type: ignore[attr-defined] + try: + return render_conversation( + "http://localhost:8000", MESSAGES, max_retries=0, **kwargs + ) + finally: + render_mod.httpx.post = original_post # type: ignore[attr-defined] + + +class TestRenderConversation: + def test_basic_render(self): + result = _call_render(_make_transport()) + assert result["token_ids"] == MOCK_TOKEN_IDS + assert result["loss_mask"] is None # not yet returned by endpoint + + def test_loss_mask_passthrough(self): + mask = [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0] + result = _call_render(_make_transport(loss_mask=mask)) + assert result["loss_mask"] == mask + + def test_error_raises(self): + with pytest.raises(RenderError, match="500"): + _call_render(_make_transport(status_code=500)) + + def test_request_body_forwarding(self): + transport, body = _capturing_transport() + tools = [{"type": "function", "function": {"name": "f", "parameters": {}}}] + _call_render( + transport, tools=tools, chat_template_kwargs={"enable_thinking": True} + ) + + assert body["tools"] == tools + assert body["chat_template_kwargs"] == {"enable_thinking": True} + assert body["add_generation_prompt"] is False From e82b37f9f668d8097af7c4e4af19b1c1ed2ac73a Mon Sep 17 00:00:00 2001 From: Ranran Haoran Zhang Date: Fri, 26 Jun 2026 11:28:52 -0500 Subject: [PATCH 02/12] wire vllm pr e2e Signed-off-by: Ranran Haoran Zhang --- .../data_generation/preprocessing.py | 45 +++++++--- .../data_generation/render_client.py | 1 + .../test_build_dataset_from_render.py | 90 +++++++++++++++++++ 3 files changed, 125 insertions(+), 11 deletions(-) create mode 100644 tests/unit/data_generation/test_build_dataset_from_render.py diff --git a/src/speculators/data_generation/preprocessing.py b/src/speculators/data_generation/preprocessing.py index 49515b51f..9a1875e56 100644 --- a/src/speculators/data_generation/preprocessing.py +++ b/src/speculators/data_generation/preprocessing.py @@ -28,8 +28,8 @@ from speculators.train.vocab_mapping import save_token_frequency_distribution __all__ = [ - "build_eagle3_dataset", "build_dataset_from_render", + "build_eagle3_dataset", "load_and_preprocess_dataset", "load_raw_dataset", ] @@ -656,27 +656,32 @@ def build_eagle3_dataset( def build_dataset_from_render( dataset: HFDataset, render_endpoint: str, + processor: ProcessorLike, max_length: int = 2048, chat_template_kwargs: dict | None = None, + turn_dropout: bool = False, minimum_valid_tokens: int | None = None, ) -> HFDataset: """Tokenize via vLLM render endpoint instead of local HF apply_chat_template.""" results: dict[str, list] = {"input_ids": [], "loss_mask": [], "seq_len": []} - has_tools = "tools" in dataset.column_names + # Regex fallback for when the render endpoint returns loss_mask=null + # (template lacks {% generation %} tags). + assistant_pattern = _detect_assistant_pattern(processor) + tokenizer = get_tokenizer(processor) + for idx in range(len(dataset)): row = dataset[idx] conv = row.get("conversations") if not conv or not isinstance(conv, list): continue - normalized = _normalize_conversation(conv) + normalized = _normalize_conversation(conv, turn_dropout) if not normalized: continue vllm_messages = _adapt_conv_for_vllm(normalized) - conv_tools = _parse_conv_tools(row.get("tools"), idx) if has_tools else None rendered = render_conversation( @@ -688,20 +693,36 @@ def build_dataset_from_render( ) token_ids = rendered["token_ids"] - loss_mask = rendered["loss_mask"] + loss_mask_raw = rendered["loss_mask"] - if loss_mask is None: - raise NotImplementedError( - "render endpoint does not yet return loss_mask; pending vLLM PR" + if loss_mask_raw is not None: + loss_mask = torch.tensor(loss_mask_raw, dtype=torch.long) + else: + text = tokenizer.decode(token_ids) + encoded = tokenizer( + text, + return_offsets_mapping=True, + max_length=max_length, + truncation=True, + add_special_tokens=False, + ) + loss_mask = _create_loss_mask_from_offsets( + text, + encoded["offset_mapping"], + assistant_pattern, + conv_idx=idx, + max_length=max_length, ) + if len(loss_mask) != len(token_ids): + loss_mask = loss_mask[: len(token_ids)] + if minimum_valid_tokens is not None: - num_valid = sum(1 for m in loss_mask if m == 1) - if num_valid < minimum_valid_tokens: + if int(loss_mask.sum().item()) < minimum_valid_tokens: continue results["input_ids"].append(torch.tensor(token_ids, dtype=torch.long)) - results["loss_mask"].append(torch.tensor(loss_mask, dtype=torch.long)) + results["loss_mask"].append(loss_mask) results["seq_len"].append(len(token_ids)) out = HFDataset.from_dict(results) @@ -817,8 +838,10 @@ def load_and_preprocess_dataset( preprocessed_dataset = build_dataset_from_render( dataset=raw_dataset, render_endpoint=render_endpoint, + processor=processor, max_length=seq_length, chat_template_kwargs=chat_template_kwargs, + turn_dropout=turn_dropout, minimum_valid_tokens=minimum_valid_tokens, ) else: diff --git a/src/speculators/data_generation/render_client.py b/src/speculators/data_generation/render_client.py index ecef389d8..eadc27145 100644 --- a/src/speculators/data_generation/render_client.py +++ b/src/speculators/data_generation/render_client.py @@ -29,6 +29,7 @@ def render_conversation( body: dict = { "messages": messages, "add_generation_prompt": False, + "return_loss_mask": True, } if tools is not None: body["tools"] = tools diff --git a/tests/unit/data_generation/test_build_dataset_from_render.py b/tests/unit/data_generation/test_build_dataset_from_render.py new file mode 100644 index 000000000..6814af507 --- /dev/null +++ b/tests/unit/data_generation/test_build_dataset_from_render.py @@ -0,0 +1,90 @@ +"""Tests for build_dataset_from_render: server-provided and regex-fallback masks.""" + +from unittest.mock import patch + +import pytest +from datasets import Dataset as HFDataset +from transformers import AutoTokenizer + +from speculators.data_generation.preprocessing import ( + build_dataset_from_render, + load_processor, +) + +MODEL = "Qwen/Qwen2.5-0.5B-Instruct" + +CONV = [ + {"role": "user", "content": "What is the capital of France?"}, + {"role": "assistant", "content": "The capital of France is Paris."}, +] + + +@pytest.fixture(scope="module") +def processor(): + return load_processor(MODEL) + + +@pytest.fixture(scope="module") +def tokenizer(): + return AutoTokenizer.from_pretrained(MODEL) + + +def _make_dataset(): + return HFDataset.from_dict({"conversations": [CONV]}) + + +def _mock_render(token_ids, loss_mask=None): + """Patch render_conversation to return fixed token_ids + loss_mask.""" + + def _fake(endpoint, messages, **kwargs): + return {"token_ids": token_ids, "loss_mask": loss_mask} + + return patch( + "speculators.data_generation.preprocessing.render_conversation", + side_effect=_fake, + ) + + +@pytest.mark.sanity +def test_server_provided_mask_flows_through(processor, tokenizer): + """When the render endpoint returns a real loss_mask, use it directly.""" + ids = tokenizer.apply_chat_template( + CONV, tokenize=True, add_generation_prompt=False, return_dict=False + ) + mask = [0] * len(ids) + # Mark the last 5 tokens as trainable (simulates assistant span). + for i in range(len(ids) - 5, len(ids)): + mask[i] = 1 + + with _mock_render(list(ids), loss_mask=mask): + ds = build_dataset_from_render( + _make_dataset(), "http://fake", processor, max_length=4096 + ) + + assert len(ds) == 1 + assert ds[0]["input_ids"].tolist() == list(ids) + assert ds[0]["loss_mask"].tolist() == mask + assert ds[0]["loss_mask"].sum().item() == 5 + + +@pytest.mark.sanity +def test_regex_fallback_when_loss_mask_is_none(processor, tokenizer): + """When loss_mask=None, fall back to regex-based assistant detection.""" + ids = tokenizer.apply_chat_template( + CONV, tokenize=True, add_generation_prompt=False, return_dict=False + ) + + with _mock_render(list(ids), loss_mask=None): + ds = build_dataset_from_render( + _make_dataset(), "http://fake", processor, max_length=4096 + ) + + assert len(ds) == 1 + assert ds[0]["input_ids"].tolist() == list(ids) + mask = ds[0]["loss_mask"] + # The regex fallback should mark assistant tokens as trainable. + assert mask.sum().item() > 0, "regex fallback produced all-zero mask" + # Verify the masked-in tokens decode to the assistant content. + masked_ids = [t for t, m in zip(ids, mask.tolist(), strict=False) if m] + decoded = tokenizer.decode(masked_ids) + assert "Paris" in decoded From 3a3e260da6dcb080a9c0336ae7a1308aab892208 Mon Sep 17 00:00:00 2001 From: Ranran Haoran Zhang Date: Fri, 26 Jun 2026 11:48:01 -0500 Subject: [PATCH 03/12] add truncation Signed-off-by: Ranran Haoran Zhang --- src/speculators/data_generation/render_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/speculators/data_generation/render_client.py b/src/speculators/data_generation/render_client.py index eadc27145..4a0b844f1 100644 --- a/src/speculators/data_generation/render_client.py +++ b/src/speculators/data_generation/render_client.py @@ -36,7 +36,7 @@ def render_conversation( if chat_template_kwargs is not None: body["chat_template_kwargs"] = chat_template_kwargs if max_length is not None: - body["max_tokens"] = max_length + body["truncate_prompt_tokens"] = max_length resp = httpx.post(url, json=body, timeout=timeout) From a0e86b4c86c991430361f7de615741601161dd0d Mon Sep 17 00:00:00 2001 From: Ranran Haoran Zhang Date: Fri, 26 Jun 2026 11:57:30 -0500 Subject: [PATCH 04/12] add e2e tests on truncation Signed-off-by: Ranran Haoran Zhang --- .../data_generation/render_client.py | 1 + .../datagen/test_render_truncation.py | 122 ++++++++++++++++++ .../data_generation/test_render_client.py | 6 +- 3 files changed, 128 insertions(+), 1 deletion(-) create mode 100644 tests/integration/datagen/test_render_truncation.py diff --git a/src/speculators/data_generation/render_client.py b/src/speculators/data_generation/render_client.py index 4a0b844f1..2d09f411f 100644 --- a/src/speculators/data_generation/render_client.py +++ b/src/speculators/data_generation/render_client.py @@ -37,6 +37,7 @@ def render_conversation( body["chat_template_kwargs"] = chat_template_kwargs if max_length is not None: body["truncate_prompt_tokens"] = max_length + body["truncation_side"] = "right" resp = httpx.post(url, json=body, timeout=timeout) diff --git a/tests/integration/datagen/test_render_truncation.py b/tests/integration/datagen/test_render_truncation.py new file mode 100644 index 000000000..f231c2ec7 --- /dev/null +++ b/tests/integration/datagen/test_render_truncation.py @@ -0,0 +1,122 @@ +"""Test that render endpoint truncation matches HF truncation. + +Heavy/local test — spins up a real vLLM render server. Skips in CI +(requires vllm installed in a sibling venv). + +Proves the render path is a drop-in replacement for the HF path: +same conversation + same seq_length → same truncated token_ids. +""" + +import os +import shutil +import socket +import subprocess +import time + +import pytest + +try: + import httpx +except ImportError: + pytest.skip("httpx not available", allow_module_level=True) + +VLLM_PYTHON = "/home/ran/workspace/vllm/.venv/bin/python" +MODEL = "Qwen/Qwen3-0.6B" +# Short enough to force truncation on a multi-turn conversation. +SEQ_LENGTH = 30 + + +def _find_free_port(): + with socket.socket() as s: + s.bind(("", 0)) + return s.getsockname()[1] + + +@pytest.fixture(scope="module") +def render_server(): + if shutil.which(VLLM_PYTHON) is None: + pytest.skip(f"vLLM python not found at {VLLM_PYTHON}") + + port = _find_free_port() + proc = subprocess.Popen( # noqa: S603 + [ + VLLM_PYTHON, + "-m", + "vllm.entrypoints.cli.main", + "launch", + "render", + "--model", + MODEL, + "--port", + str(port), + "--host", + "127.0.0.1", + "--max-model-len", + "4096", + ], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + env={**os.environ, "HF_HUB_OFFLINE": "1"}, + ) + url = f"http://127.0.0.1:{port}" + for _ in range(60): + try: + if httpx.get(f"{url}/health", timeout=2).status_code == 200: + break + except httpx.ConnectError: + pass + time.sleep(1) + else: + proc.kill() + pytest.skip("render server failed to start") + + yield url + proc.terminate() + proc.wait(timeout=10) + + +@pytest.mark.sanity +def test_render_truncation_matches_hf(render_server): + """Truncated token_ids from the render endpoint must equal + HF apply_chat_template with the same max_length.""" + from transformers import AutoTokenizer # noqa: PLC0415 + + tok = AutoTokenizer.from_pretrained(MODEL) + + conv = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Tell me about the history of Paris."}, + { + "role": "assistant", + "content": "Paris has a long and fascinating history.", + }, + {"role": "user", "content": "What about Rome?"}, + {"role": "assistant", "content": "Rome is the Eternal City."}, + ] + + # HF path: what speculators does today. + full_ids = tok.apply_chat_template( + conv, + tokenize=True, + add_generation_prompt=False, + return_dict=False, + ) + assert len(full_ids) > SEQ_LENGTH, "conversation too short to test truncation" + hf_truncated = list(full_ids[:SEQ_LENGTH]) + + # Render path: POST with truncate_prompt_tokens. + resp = httpx.post( + f"{render_server}/v1/chat/completions/render", + json={ + "model": MODEL, + "messages": conv, + "add_generation_prompt": False, + "truncate_prompt_tokens": SEQ_LENGTH, + "truncation_side": "right", + }, + timeout=30, + ) + assert resp.status_code == 200, resp.text + render_ids = resp.json()["token_ids"] + + assert render_ids == hf_truncated diff --git a/tests/unit/data_generation/test_render_client.py b/tests/unit/data_generation/test_render_client.py index 335793e47..e1a79a870 100644 --- a/tests/unit/data_generation/test_render_client.py +++ b/tests/unit/data_generation/test_render_client.py @@ -89,9 +89,13 @@ def test_request_body_forwarding(self): transport, body = _capturing_transport() tools = [{"type": "function", "function": {"name": "f", "parameters": {}}}] _call_render( - transport, tools=tools, chat_template_kwargs={"enable_thinking": True} + transport, + tools=tools, + chat_template_kwargs={"enable_thinking": True}, + max_length=4096, ) assert body["tools"] == tools assert body["chat_template_kwargs"] == {"enable_thinking": True} assert body["add_generation_prompt"] is False + assert body["truncate_prompt_tokens"] == 4096 From cc9b72e65a35baacd688b9858b995fb30f9fb4d6 Mon Sep 17 00:00:00 2001 From: Ranran Haoran Zhang Date: Sat, 27 Jun 2026 00:07:26 -0500 Subject: [PATCH 05/12] clarify help doc, and rename var Signed-off-by: Ranran Haoran Zhang --- scripts/prepare_data.py | 9 +++++---- src/speculators/data_generation/render_client.py | 7 +++++-- tests/unit/data_generation/test_render_client.py | 8 ++++---- 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/scripts/prepare_data.py b/scripts/prepare_data.py index 95d92a90f..045ba0420 100644 --- a/scripts/prepare_data.py +++ b/scripts/prepare_data.py @@ -110,10 +110,11 @@ def parse_args(): type=str, default=None, help=( - "URL of a vLLM render service " - "(e.g. http://localhost:8000). When set, tokenization is " - "delegated to the render endpoint instead of local HF " - "apply_chat_template. See RFC #652." + "URL of the vLLM server, typically the same instance used " + "for hidden-state extraction (e.g. http://localhost:8000). " + "When set, tokenization is delegated to the server's " + "/v1/chat/completions/render endpoint instead of local HF " + "apply_chat_template." ), ) diff --git a/src/speculators/data_generation/render_client.py b/src/speculators/data_generation/render_client.py index 2d09f411f..d63c7a4b2 100644 --- a/src/speculators/data_generation/render_client.py +++ b/src/speculators/data_generation/render_client.py @@ -29,7 +29,7 @@ def render_conversation( body: dict = { "messages": messages, "add_generation_prompt": False, - "return_loss_mask": True, + "return_assistant_tokens_mask": True, } if tools is not None: body["tools"] = tools @@ -47,4 +47,7 @@ def render_conversation( ) data = resp.json() - return {"token_ids": data["token_ids"], "loss_mask": data.get("loss_mask")} + return { + "token_ids": data["token_ids"], + "loss_mask": data.get("assistant_tokens_mask"), + } diff --git a/tests/unit/data_generation/test_render_client.py b/tests/unit/data_generation/test_render_client.py index e1a79a870..88aa4fdd3 100644 --- a/tests/unit/data_generation/test_render_client.py +++ b/tests/unit/data_generation/test_render_client.py @@ -32,13 +32,13 @@ ] -def _make_transport(*, token_ids=None, loss_mask=None, status_code=200): +def _make_transport(*, token_ids=None, assistant_tokens_mask=None, status_code=200): def handler(request: httpx.Request) -> httpx.Response: if status_code != 200: return httpx.Response(status_code, text="error from mock") body = {"token_ids": token_ids or MOCK_TOKEN_IDS} - if loss_mask is not None: - body["loss_mask"] = loss_mask + if assistant_tokens_mask is not None: + body["assistant_tokens_mask"] = assistant_tokens_mask return httpx.Response(200, json=body) return httpx.MockTransport(handler) @@ -78,7 +78,7 @@ def test_basic_render(self): def test_loss_mask_passthrough(self): mask = [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0] - result = _call_render(_make_transport(loss_mask=mask)) + result = _call_render(_make_transport(assistant_tokens_mask=mask)) assert result["loss_mask"] == mask def test_error_raises(self): From 4b14fa43651470018af52c41a7ff165140336744 Mon Sep 17 00:00:00 2001 From: Ranran Haoran Zhang Date: Mon, 29 Jun 2026 00:32:56 -0500 Subject: [PATCH 06/12] add log to which path of regex, hf, renderer Signed-off-by: Ranran Haoran Zhang --- .../data_generation/preprocessing.py | 20 +++++++++++--- .../test_build_dataset_from_render.py | 26 +++++++++++++++++++ 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/src/speculators/data_generation/preprocessing.py b/src/speculators/data_generation/preprocessing.py index 9a1875e56..6970ec106 100644 --- a/src/speculators/data_generation/preprocessing.py +++ b/src/speculators/data_generation/preprocessing.py @@ -665,12 +665,15 @@ def build_dataset_from_render( """Tokenize via vLLM render endpoint instead of local HF apply_chat_template.""" results: dict[str, list] = {"input_ids": [], "loss_mask": [], "seq_len": []} has_tools = "tools" in dataset.column_names - - # Regex fallback for when the render endpoint returns loss_mask=null - # (template lacks {% generation %} tags). - assistant_pattern = _detect_assistant_pattern(processor) tokenizer = get_tokenizer(processor) + # Regex fallback for conversations the render endpoint returns no mask for + # (template lacks {% generation %} tags). Detected lazily, so a run the + # server masks completely never fails here. + assistant_pattern: str | None = None + server_mask_count = 0 + regex_fallback_count = 0 + for idx in range(len(dataset)): row = dataset[idx] conv = row.get("conversations") @@ -697,7 +700,11 @@ def build_dataset_from_render( if loss_mask_raw is not None: loss_mask = torch.tensor(loss_mask_raw, dtype=torch.long) + server_mask_count += 1 else: + if assistant_pattern is None: + assistant_pattern = _detect_assistant_pattern(processor) + regex_fallback_count += 1 text = tokenizer.decode(token_ids) encoded = tokenizer( text, @@ -725,6 +732,11 @@ def build_dataset_from_render( results["loss_mask"].append(loss_mask) results["seq_len"].append(len(token_ids)) + log.info( + f"Render masking: {server_mask_count} conversations used the vLLM server " + f"mask, {regex_fallback_count} used the local regex fallback." + ) + out = HFDataset.from_dict(results) out.set_format(type="torch") return out diff --git a/tests/unit/data_generation/test_build_dataset_from_render.py b/tests/unit/data_generation/test_build_dataset_from_render.py index 6814af507..5ba9522d1 100644 --- a/tests/unit/data_generation/test_build_dataset_from_render.py +++ b/tests/unit/data_generation/test_build_dataset_from_render.py @@ -88,3 +88,29 @@ def test_regex_fallback_when_loss_mask_is_none(processor, tokenizer): masked_ids = [t for t, m in zip(ids, mask.tolist(), strict=False) if m] decoded = tokenizer.decode(masked_ids) assert "Paris" in decoded + + +@pytest.mark.sanity +def test_full_server_mask_skips_local_pattern_detection(processor, tokenizer): + """When the server masks every conversation, the local regex pattern is + never detected -- so a model whose detection would raise still succeeds.""" + ids = tokenizer.apply_chat_template( + CONV, tokenize=True, add_generation_prompt=False, return_dict=False + ) + mask = [0] * len(ids) + mask[-1] = 1 + + with ( + _mock_render(list(ids), loss_mask=mask), + patch( + "speculators.data_generation.preprocessing._detect_assistant_pattern", + side_effect=ValueError("cannot detect"), + ) as detect, + ): + ds = build_dataset_from_render( + _make_dataset(), "http://fake", processor, max_length=4096 + ) + + detect.assert_not_called() + assert len(ds) == 1 + assert ds[0]["loss_mask"].sum().item() == 1 From 8b1938f793645fa3b224ba108dc57034a5806a59 Mon Sep 17 00:00:00 2001 From: Ranran Haoran Zhang Date: Tue, 30 Jun 2026 11:04:13 -0500 Subject: [PATCH 07/12] fix corner case Signed-off-by: Ranran Haoran Zhang --- .../data_generation/preprocessing.py | 67 +++++++-- .../datagen/test_render_multimodal.py | 130 ++++++++++++++++++ .../test_build_dataset_from_render.py | 73 ++++++++++ 3 files changed, 255 insertions(+), 15 deletions(-) create mode 100644 tests/integration/datagen/test_render_multimodal.py diff --git a/src/speculators/data_generation/preprocessing.py b/src/speculators/data_generation/preprocessing.py index 6970ec106..9a66eef1a 100644 --- a/src/speculators/data_generation/preprocessing.py +++ b/src/speculators/data_generation/preprocessing.py @@ -8,6 +8,7 @@ from re import Pattern from typing import cast +import httpx import torch from datasets import Dataset as HFDataset from datasets import concatenate_datasets, load_dataset @@ -23,7 +24,7 @@ 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.render_client import RenderError, render_conversation from speculators.data_generation.torch_utils import set_default_torch_num_threads from speculators.train.vocab_mapping import save_token_frequency_distribution @@ -667,12 +668,20 @@ def build_dataset_from_render( has_tools = "tools" in dataset.column_names tokenizer = get_tokenizer(processor) + # Multimodal models extract hidden states via the Chat Completions API, which + # needs the original messages -- token_ids alone cannot carry the images. + # Mirror the default path and keep messages for ProcessorMixin processors so + # render-path multimodal rows can be extracted downstream. + if isinstance(processor, ProcessorMixin): + results["messages"] = [] + # Regex fallback for conversations the render endpoint returns no mask for # (template lacks {% generation %} tags). Detected lazily, so a run the # server masks completely never fails here. assistant_pattern: str | None = None server_mask_count = 0 regex_fallback_count = 0 + dropped_count = 0 for idx in range(len(dataset)): row = dataset[idx] @@ -687,30 +696,46 @@ def build_dataset_from_render( vllm_messages = _adapt_conv_for_vllm(normalized) conv_tools = _parse_conv_tools(row.get("tools"), idx) if has_tools else None - rendered = render_conversation( - render_endpoint, - vllm_messages, - tools=conv_tools, - chat_template_kwargs=chat_template_kwargs, - max_length=max_length, - ) + # Resilience parity with the default path (_preprocess_batch wraps each + # conversation in try/except): one conversation that fails to render must + # not abort the whole run. + try: + rendered = render_conversation( + render_endpoint, + vllm_messages, + tools=conv_tools, + chat_template_kwargs=chat_template_kwargs, + max_length=max_length, + ) + except (RenderError, httpx.HTTPError) as exc: + log.error(f"Render failed for conversation {idx}: {exc}") + dropped_count += 1 + continue token_ids = rendered["token_ids"] loss_mask_raw = rendered["loss_mask"] + # vLLM does not truncate multimodal prompts -- it cannot split an image's + # placeholder block -- so token_ids can exceed max_length. The default + # pipeline drops these rows; match it instead of emitting an over-length + # (and, after local masking, misaligned) sample. + if len(token_ids) > max_length: + dropped_count += 1 + continue + if loss_mask_raw is not None: - loss_mask = torch.tensor(loss_mask_raw, dtype=torch.long) + loss_mask = torch.as_tensor(loss_mask_raw, dtype=torch.long) server_mask_count += 1 else: if assistant_pattern is None: assistant_pattern = _detect_assistant_pattern(processor) - regex_fallback_count += 1 + # token_ids is already within max_length (over-length rows dropped + # above), so re-tokenize without truncation -- it round-trips to the + # same tokens and the offsets line up with the server's token_ids. text = tokenizer.decode(token_ids) encoded = tokenizer( text, return_offsets_mapping=True, - max_length=max_length, - truncation=True, add_special_tokens=False, ) loss_mask = _create_loss_mask_from_offsets( @@ -720,9 +745,18 @@ def build_dataset_from_render( conv_idx=idx, max_length=max_length, ) + loss_mask = torch.as_tensor(loss_mask, dtype=torch.long) + regex_fallback_count += 1 + # Parity with the default path's `assert len(input_ids) == len(loss_mask)`: + # never emit a misaligned pair -- drop it rather than corrupt the sample. if len(loss_mask) != len(token_ids): - loss_mask = loss_mask[: len(token_ids)] + log.error( + f"Render mask/token length mismatch for conversation {idx} " + f"(input_ids={len(token_ids)}, loss_mask={len(loss_mask)}); dropping" + ) + dropped_count += 1 + continue if minimum_valid_tokens is not None: if int(loss_mask.sum().item()) < minimum_valid_tokens: @@ -731,10 +765,13 @@ def build_dataset_from_render( results["input_ids"].append(torch.tensor(token_ids, dtype=torch.long)) results["loss_mask"].append(loss_mask) results["seq_len"].append(len(token_ids)) + if "messages" in results: + results["messages"].append(vllm_messages) log.info( - f"Render masking: {server_mask_count} conversations used the vLLM server " - f"mask, {regex_fallback_count} used the local regex fallback." + f"Render masking: {server_mask_count} used the vLLM server mask, " + f"{regex_fallback_count} used the local regex fallback, " + f"{dropped_count} dropped (render error / over-length / misaligned)." ) out = HFDataset.from_dict(results) diff --git a/tests/integration/datagen/test_render_multimodal.py b/tests/integration/datagen/test_render_multimodal.py new file mode 100644 index 000000000..81dff3117 --- /dev/null +++ b/tests/integration/datagen/test_render_multimodal.py @@ -0,0 +1,130 @@ +"""Faithful multimodal test for the render fallback — real Qwen3-VL render server. + +Heavy/local test (skips in CI). Reproduces a real failure: vLLM does NOT truncate +multimodal prompts to ``truncate_prompt_tokens``, so for a conversation that exceeds +``seq_length`` the server returns more ``token_ids`` than ``max_length`` while the +local regex fallback truncates the mask to ``max_length`` -> ``input_ids`` and +``loss_mask`` end up different lengths. + +Set ``RENDER_ENDPOINT`` to reuse an already-running Qwen3-VL render server; otherwise +one is launched from the sibling vLLM venv. +""" + +import os +import shutil +import socket +import subprocess +import time + +import numpy as np +import pytest + +try: + import httpx +except ImportError: + pytest.skip("httpx not available", allow_module_level=True) + +from datasets import Dataset as HFDataset +from PIL import Image + +from speculators.data_generation.preprocessing import ( + build_dataset_from_render, + load_processor, +) + +VLLM_PYTHON = "/home/ran/workspace/vllm/.venv/bin/python" +MODEL = "Qwen/Qwen3-VL-2B-Instruct" +MEDIA_DIR = "/tmp" +IMG_PATH = "/tmp/specforge_render_mm_test.png" + + +def _find_free_port(): + with socket.socket() as s: + s.bind(("", 0)) + return s.getsockname()[1] + + +@pytest.fixture(scope="module") +def image(): + arr = np.zeros((128, 128, 3), dtype=np.uint8) + arr[:64, :, 0] = 200 + arr[64:, :, 2] = 200 + arr[32:96, 32:96, 1] = 180 + Image.fromarray(arr).save(IMG_PATH) + return IMG_PATH + + +@pytest.fixture(scope="module") +def render_server(): + env_endpoint = os.environ.get("RENDER_ENDPOINT") + if env_endpoint: + yield env_endpoint.rstrip("/") + return + + if shutil.which(VLLM_PYTHON) is None: + pytest.skip(f"vLLM python not found at {VLLM_PYTHON}") + + port = _find_free_port() + proc = subprocess.Popen( # noqa: S603 + [ + VLLM_PYTHON, "-m", "vllm.entrypoints.cli.main", "launch", "render", + "--model", MODEL, "--port", str(port), "--host", "127.0.0.1", + "--max-model-len", "4096", "--gpu-memory-utilization", "0.3", + "--allowed-local-media-path", MEDIA_DIR, + ], + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + env={**os.environ, "HF_HUB_OFFLINE": "1"}, + ) + url = f"http://127.0.0.1:{port}" + for _ in range(120): + try: + if httpx.get(f"{url}/health", timeout=2).status_code == 200: + break + except httpx.ConnectError: + pass + time.sleep(2) + else: + proc.kill() + pytest.skip("render server failed to start") + yield url + proc.terminate() + proc.wait(timeout=10) + + +@pytest.mark.sanity +def test_multimodal_over_length_not_misaligned(render_server, image): + """A multimodal conversation that exceeds max_length must never be emitted as + a misaligned (input_ids vs loss_mask) pair. vLLM won't truncate the image + prompt, so the render path drops the row -- parity with the default pipeline, + which drops over-length multimodal rows too. (Before the fix this emitted a + 86-vs-48 misaligned sample.)""" + processor = load_processor(MODEL) + conv = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe the colors."}, + {"type": "image", "path": image}, + ], + }, + { + "role": "assistant", + "content": [{"type": "text", "text": "Red, blue, green."}], + }, + ] + # max_length below the image's own token count: vLLM can't fit it. + ds = build_dataset_from_render( + HFDataset.from_dict({"conversations": [conv]}), + render_server, + processor, + max_length=48, + ) + + for i in range(len(ds)): + n_ids = len(ds[i]["input_ids"]) + n_mask = len(ds[i]["loss_mask"]) + assert n_ids == n_mask, ( + f"misaligned sample emitted: input_ids={n_ids} loss_mask={n_mask}" + ) + # over-length multimodal row is dropped (parity with the default pipeline) + assert len(ds) == 0 diff --git a/tests/unit/data_generation/test_build_dataset_from_render.py b/tests/unit/data_generation/test_build_dataset_from_render.py index 5ba9522d1..4ede9c981 100644 --- a/tests/unit/data_generation/test_build_dataset_from_render.py +++ b/tests/unit/data_generation/test_build_dataset_from_render.py @@ -10,6 +10,7 @@ build_dataset_from_render, load_processor, ) +from speculators.data_generation.render_client import RenderError MODEL = "Qwen/Qwen2.5-0.5B-Instruct" @@ -114,3 +115,75 @@ def test_full_server_mask_skips_local_pattern_detection(processor, tokenizer): detect.assert_not_called() assert len(ds) == 1 assert ds[0]["loss_mask"].sum().item() == 1 + + +@pytest.mark.sanity +def test_over_length_render_dropped(processor): + """vLLM does not truncate multimodal prompts, so the server can return + token_ids longer than max_length. Such a row is dropped (parity with the + default path), never emitted as an over-length / misaligned sample.""" + long_ids = list(range(100)) # 100 > max_length (50) below + with _mock_render(long_ids, loss_mask=None): + ds = build_dataset_from_render( + _make_dataset(), "http://fake", processor, max_length=50 + ) + assert len(ds) == 0 + + +@pytest.mark.sanity +def test_render_error_does_not_crash_the_run(processor): + """A per-conversation render failure is skipped, not propagated -- one bad + row must not abort the whole dataset build (parity with the default path).""" + with patch( + "speculators.data_generation.preprocessing.render_conversation", + side_effect=RenderError("boom"), + ): + ds = build_dataset_from_render( + _make_dataset(), "http://fake", processor, max_length=4096 + ) + assert len(ds) == 0 + + +@pytest.fixture(scope="module") +def mm_processor(): + return load_processor("Qwen/Qwen3-VL-2B-Instruct") + + +@pytest.mark.sanity +def test_multimodal_keeps_messages_for_extraction(mm_processor): + """Multimodal (ProcessorMixin) rows keep a `messages` column so hidden-state + extraction can re-send the conversation with images via the chat API -- + token_ids alone cannot carry the image.""" + mm_conv = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this image?"}, + {"type": "image", "path": "/tmp/x.png"}, + ], + }, + {"role": "assistant", "content": [{"type": "text", "text": "A cat."}]}, + ] + ids = list(range(12)) + mask = [0] * 9 + [1, 1, 1] + with _mock_render(ids, loss_mask=mask): + ds = build_dataset_from_render( + HFDataset.from_dict({"conversations": [mm_conv]}), + "http://fake", + mm_processor, + max_length=4096, + ) + assert "messages" in ds.column_names + assert len(ds) == 1 + assert ds[0]["messages"][0]["role"] == "user" + + +@pytest.mark.sanity +def test_text_model_has_no_messages_column(processor): + """Text-only processors must not add a `messages` column (parity with the + default path, which only adds it for ProcessorMixin).""" + with _mock_render(list(range(10)), loss_mask=[0] * 7 + [1, 1, 1]): + ds = build_dataset_from_render( + _make_dataset(), "http://fake", processor, max_length=4096 + ) + assert "messages" not in ds.column_names From e6292e295cf278eb4f8ac43aa3af06976cf002c1 Mon Sep 17 00:00:00 2001 From: Ranran Haoran Zhang Date: Fri, 3 Jul 2026 15:03:46 -0500 Subject: [PATCH 08/12] fix path leakage, dependencies Signed-off-by: Ranran Haoran Zhang --- pyproject.toml | 1 + tests/integration/datagen/test_render_multimodal.py | 4 ++-- tests/integration/datagen/test_render_truncation.py | 3 ++- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 1c2450862..8a45ec27e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,7 @@ keywords = [ dependencies = [ "click", "datasets>=4.0.0,<=5.0.0", + "httpx", "huggingface-hub", "loguru>=0.7.2,<=0.7.3", "numpy>=2.0.0,<=2.4.6", diff --git a/tests/integration/datagen/test_render_multimodal.py b/tests/integration/datagen/test_render_multimodal.py index 81dff3117..e0f0fad14 100644 --- a/tests/integration/datagen/test_render_multimodal.py +++ b/tests/integration/datagen/test_render_multimodal.py @@ -31,11 +31,11 @@ build_dataset_from_render, load_processor, ) +from tests.e2e.utils import VLLM_PYTHON -VLLM_PYTHON = "/home/ran/workspace/vllm/.venv/bin/python" MODEL = "Qwen/Qwen3-VL-2B-Instruct" MEDIA_DIR = "/tmp" -IMG_PATH = "/tmp/specforge_render_mm_test.png" +IMG_PATH = "/tmp/speculators_render_mm_test.png" def _find_free_port(): diff --git a/tests/integration/datagen/test_render_truncation.py b/tests/integration/datagen/test_render_truncation.py index f231c2ec7..3baec633b 100644 --- a/tests/integration/datagen/test_render_truncation.py +++ b/tests/integration/datagen/test_render_truncation.py @@ -20,7 +20,8 @@ except ImportError: pytest.skip("httpx not available", allow_module_level=True) -VLLM_PYTHON = "/home/ran/workspace/vllm/.venv/bin/python" +from tests.e2e.utils import VLLM_PYTHON + MODEL = "Qwen/Qwen3-0.6B" # Short enough to force truncation on a multi-turn conversation. SEQ_LENGTH = 30 From 4d44353f6836c28ae858a0dc7cb2b267024a5a82 Mon Sep 17 00:00:00 2001 From: Ranran Haoran Zhang Date: Fri, 3 Jul 2026 15:56:38 -0500 Subject: [PATCH 09/12] add more guards, restore and update the doc Signed-off-by: Ranran Haoran Zhang --- scripts/prepare_data.py | 23 ++++++++-- .../data_generation/preprocessing.py | 44 ++++++++++++++++++- .../test_load_and_preprocess_dataset.py | 39 ++++++++++++++++ 3 files changed, 102 insertions(+), 4 deletions(-) create mode 100644 tests/unit/data_generation/test_load_and_preprocess_dataset.py diff --git a/scripts/prepare_data.py b/scripts/prepare_data.py index 045ba0420..c86054112 100644 --- a/scripts/prepare_data.py +++ b/scripts/prepare_data.py @@ -24,6 +24,7 @@ import argparse import glob +import json import logging import shutil import sys @@ -106,17 +107,27 @@ def parse_args(): ), ) parser.add_argument( - "--render-endpoint", + "--vllm-server-url", type=str, default=None, help=( - "URL of the vLLM server, typically the same instance used " + "Base URL of the vLLM server, typically the same instance used " "for hidden-state extraction (e.g. http://localhost:8000). " "When set, tokenization is delegated to the server's " "/v1/chat/completions/render endpoint instead of local HF " "apply_chat_template." ), ) + parser.add_argument( + "--chat-template-kwargs", + type=str, + default=None, + help=( + "JSON object of extra kwargs forwarded to the chat template, " + 'e.g. \'{"enable_thinking": true}\'. Only supported on the render ' + "path (requires --vllm-server-url)." + ), + ) # Output arguments parser.add_argument( @@ -189,6 +200,11 @@ def main(): if args.token_freq_path is None else Path(args.token_freq_path) ) + chat_template_kwargs = ( + json.loads(args.chat_template_kwargs) + if args.chat_template_kwargs is not None + else None + ) dataset, _ = load_and_preprocess_dataset( target_model_path=args.model, @@ -202,7 +218,8 @@ def main(): turn_dropout=args.turn_dropout, minimum_valid_tokens=args.minimum_valid_tokens, trust_remote_code=args.trust_remote_code, - render_endpoint=args.render_endpoint, + render_endpoint=args.vllm_server_url, + chat_template_kwargs=chat_template_kwargs, ) log.info("Done preparing data") diff --git a/src/speculators/data_generation/preprocessing.py b/src/speculators/data_generation/preprocessing.py index 9a66eef1a..a01ad3583 100644 --- a/src/speculators/data_generation/preprocessing.py +++ b/src/speculators/data_generation/preprocessing.py @@ -840,9 +840,51 @@ def load_and_preprocess_dataset( render_endpoint: str | None = None, chat_template_kwargs: dict | None = None, ) -> tuple[HFDataset, ProcessorLike]: - """Load, tokenize, and preprocess a dataset for EAGLE3 training.""" + """Load, tokenize, and preprocess a dataset for EAGLE3 training. + + Uses the processor's built-in chat template via apply_chat_template, or + delegates tokenization to a running vLLM render server when render_endpoint + is set. Caching is handled automatically by HuggingFace datasets. + + Args: + target_model_path: HuggingFace model ID or local path. + train_data_paths: Dataset names or paths to JSON/JSONL files. + seq_length: Maximum sequence length. + build_dataset_num_proc: Number of processes for dataset building. + seed: Random seed for shuffling. + max_samples: Optional limit on number of samples. + token_freq_path: Path to save token frequency distribution. + assistant_pattern: Optional custom regex pattern for matching assistant + responses. If None, pattern will be auto-detected from the chat + template. Mutually exclusive with render_endpoint -- the render + path always auto-detects its own pattern. + turn_dropout: If True, randomly keeps first N consecutive turns per + conversation. + minimum_valid_tokens: Number of tokens to consider for a valid sample. + trust_remote_code: If True, allows executing code from HF Hub. + render_endpoint: URL of a running vLLM server. When set, tokenization + is delegated to the server's /v1/chat/completions/render endpoint + instead of local HF apply_chat_template. + chat_template_kwargs: Extra kwargs forwarded to the chat template + (e.g. {"enable_thinking": True}). Only supported on the render + path (requires render_endpoint). + + Returns: + Tuple of (preprocessed_dataset, processor). + """ if minimum_valid_tokens is not None and minimum_valid_tokens < 0: raise ValueError("minimum_valid_tokens must be >= 0") + if render_endpoint is not None and assistant_pattern is not None: + raise ValueError( + "assistant_pattern has no effect when render_endpoint is set -- the " + "render path always auto-detects its own pattern from the chat " + "template. Pass only one of assistant_pattern / render_endpoint." + ) + if chat_template_kwargs is not None and render_endpoint is None: + raise ValueError( + "chat_template_kwargs is only supported on the render path; pass " + "render_endpoint or drop chat_template_kwargs." + ) log.section("Starting dataset preprocessing") if minimum_valid_tokens is not None: log.info( diff --git a/tests/unit/data_generation/test_load_and_preprocess_dataset.py b/tests/unit/data_generation/test_load_and_preprocess_dataset.py new file mode 100644 index 000000000..febfe126b --- /dev/null +++ b/tests/unit/data_generation/test_load_and_preprocess_dataset.py @@ -0,0 +1,39 @@ +"""Tests for load_and_preprocess_dataset's render-path argument validation. + +Both checks fire before target_model_path / train_data_paths are ever touched +(no processor load, no dataset load), so these run with no mocking. +""" + +import pytest + +from speculators.data_generation.preprocessing import load_and_preprocess_dataset + + +@pytest.mark.sanity +def test_assistant_pattern_rejected_with_render_endpoint(): + """assistant_pattern is silently unreachable on the render path (it always + auto-detects its own pattern) -- reject the combination instead of + dropping the user's pattern without telling them.""" + with pytest.raises(ValueError, match="assistant_pattern"): + load_and_preprocess_dataset( + target_model_path="unused", + train_data_paths=["unused"], + seq_length=10, + assistant_pattern="custom-pattern", + render_endpoint="http://localhost:8000", + ) + + +@pytest.mark.sanity +def test_chat_template_kwargs_rejected_without_render_endpoint(): + """chat_template_kwargs only reaches build_dataset_from_render; the default + HF path (build_eagle3_dataset) never receives it. Reject rather than + silently ignore a value the user explicitly set.""" + with pytest.raises(ValueError, match="chat_template_kwargs"): + load_and_preprocess_dataset( + target_model_path="unused", + train_data_paths=["unused"], + seq_length=10, + chat_template_kwargs={"enable_thinking": True}, + render_endpoint=None, + ) From 3a39c78ffe1aaea9bb20687429307bd801e4cc22 Mon Sep 17 00:00:00 2001 From: Ranran Haoran Zhang Date: Fri, 3 Jul 2026 16:20:31 -0500 Subject: [PATCH 10/12] better fail retry and tests Signed-off-by: Ranran Haoran Zhang --- .../data_generation/preprocessing.py | 3 +- .../data_generation/render_client.py | 13 ++++- .../data_generation/test_render_client.py | 48 +++++++++++++++---- 3 files changed, 54 insertions(+), 10 deletions(-) diff --git a/src/speculators/data_generation/preprocessing.py b/src/speculators/data_generation/preprocessing.py index a01ad3583..f2f2d59f4 100644 --- a/src/speculators/data_generation/preprocessing.py +++ b/src/speculators/data_generation/preprocessing.py @@ -26,6 +26,7 @@ from speculators.data_generation.logging_utils import PipelineLogger from speculators.data_generation.render_client import RenderError, render_conversation from speculators.data_generation.torch_utils import set_default_torch_num_threads +from speculators.data_generation.vllm_client import InvalidResponseError from speculators.train.vocab_mapping import save_token_frequency_distribution __all__ = [ @@ -707,7 +708,7 @@ def build_dataset_from_render( chat_template_kwargs=chat_template_kwargs, max_length=max_length, ) - except (RenderError, httpx.HTTPError) as exc: + except (RenderError, InvalidResponseError, httpx.HTTPError) as exc: log.error(f"Render failed for conversation {idx}: {exc}") dropped_count += 1 continue diff --git a/src/speculators/data_generation/render_client.py b/src/speculators/data_generation/render_client.py index d63c7a4b2..5d0e3a9a0 100644 --- a/src/speculators/data_generation/render_client.py +++ b/src/speculators/data_generation/render_client.py @@ -4,7 +4,7 @@ import httpx -from speculators.data_generation.vllm_client import with_retries +from speculators.data_generation.vllm_client import InvalidResponseError, with_retries DEFAULT_RENDER_TIMEOUT = 30 @@ -41,12 +41,23 @@ def render_conversation( resp = httpx.post(url, json=body, timeout=timeout) + if HTTPStatus.BAD_REQUEST <= resp.status_code < HTTPStatus.INTERNAL_SERVER_ERROR: + # 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[:500]}" + ) if resp.status_code != HTTPStatus.OK: raise RenderError( f"Render endpoint returned {resp.status_code}: {resp.text[:500]}" ) data = resp.json() + if "token_ids" not in data: + raise RenderError( + f"Render endpoint response missing 'token_ids': {str(data)[:500]}" + ) return { "token_ids": data["token_ids"], "loss_mask": data.get("assistant_tokens_mask"), diff --git a/tests/unit/data_generation/test_render_client.py b/tests/unit/data_generation/test_render_client.py index 88aa4fdd3..09542fee7 100644 --- a/tests/unit/data_generation/test_render_client.py +++ b/tests/unit/data_generation/test_render_client.py @@ -10,6 +10,7 @@ RenderError, render_conversation, ) +from speculators.data_generation.vllm_client import InvalidResponseError MOCK_TOKEN_IDS = [ 151644, @@ -32,10 +33,8 @@ ] -def _make_transport(*, token_ids=None, assistant_tokens_mask=None, status_code=200): +def _make_transport(*, token_ids=None, assistant_tokens_mask=None): def handler(request: httpx.Request) -> httpx.Response: - if status_code != 200: - return httpx.Response(status_code, text="error from mock") body = {"token_ids": token_ids or MOCK_TOKEN_IDS} if assistant_tokens_mask is not None: body["assistant_tokens_mask"] = assistant_tokens_mask @@ -55,7 +54,21 @@ def handler(request: httpx.Request) -> httpx.Response: return httpx.MockTransport(handler), captured -def _call_render(transport, **kwargs): +def _counting_transport(*, status_code=200, body=None): + """Return (transport, call_counter) -- counter tracks requests received, + to assert whether a failure was retried.""" + calls = {"count": 0} + + def handler(request: httpx.Request) -> httpx.Response: + calls["count"] += 1 + if status_code != 200: + return httpx.Response(status_code, text="error from mock") + return httpx.Response(200, json=body or {"token_ids": MOCK_TOKEN_IDS}) + + return httpx.MockTransport(handler), calls + + +def _call_render(transport, *, max_retries=0, **kwargs): original_post = httpx.post def mock_post(url, **kw): @@ -64,7 +77,7 @@ def mock_post(url, **kw): render_mod.httpx.post = mock_post # type: ignore[attr-defined] try: return render_conversation( - "http://localhost:8000", MESSAGES, max_retries=0, **kwargs + "http://localhost:8000", MESSAGES, max_retries=max_retries, **kwargs ) finally: render_mod.httpx.post = original_post # type: ignore[attr-defined] @@ -81,9 +94,28 @@ def test_loss_mask_passthrough(self): result = _call_render(_make_transport(assistant_tokens_mask=mask)) assert result["loss_mask"] == mask - def test_error_raises(self): - with pytest.raises(RenderError, match="500"): - _call_render(_make_transport(status_code=500)) + @pytest.mark.parametrize( + ("status_code", "body", "max_retries", "expected_exc", "expected_calls"), + [ + # 4xx is deterministic (bad request, wrong URL) -- retrying wastes + # requests without changing the outcome. + (400, None, 3, InvalidResponseError, 1), + # 5xx may be transient -- goes through the normal retry policy + # (initial attempt + max_retries). + (500, None, 2, RenderError, 3), + # 200 but missing token_ids must raise RenderError, not a bare + # KeyError that build_dataset_from_render's except clause can't + # catch -- and it's not exempt from the retry policy either. + (200, {"assistant_tokens_mask": [0, 1]}, 1, RenderError, 2), + ], + ) + def test_error_handling_and_retry_policy( + self, status_code, body, max_retries, expected_exc, expected_calls + ): + transport, calls = _counting_transport(status_code=status_code, body=body) + with pytest.raises(expected_exc): + _call_render(transport, max_retries=max_retries) + assert calls["count"] == expected_calls def test_request_body_forwarding(self): transport, body = _capturing_transport() From 4297074cddf88eeb73562d96923518fe8fb3e0d8 Mon Sep 17 00:00:00 2001 From: Ranran Haoran Zhang Date: Tue, 7 Jul 2026 10:23:15 -0500 Subject: [PATCH 11/12] fix test temp path Signed-off-by: Ranran Haoran Zhang --- .../datagen/test_render_multimodal.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/tests/integration/datagen/test_render_multimodal.py b/tests/integration/datagen/test_render_multimodal.py index e0f0fad14..7bcb11144 100644 --- a/tests/integration/datagen/test_render_multimodal.py +++ b/tests/integration/datagen/test_render_multimodal.py @@ -34,8 +34,6 @@ from tests.e2e.utils import VLLM_PYTHON MODEL = "Qwen/Qwen3-VL-2B-Instruct" -MEDIA_DIR = "/tmp" -IMG_PATH = "/tmp/speculators_render_mm_test.png" def _find_free_port(): @@ -45,17 +43,25 @@ def _find_free_port(): @pytest.fixture(scope="module") -def image(): +def test_dir(tmp_path_factory): + """Module-scoped isolated temp dir: holds the generated image and doubles + as the server's --allowed-local-media-path root.""" + return tmp_path_factory.mktemp("speculators_render") + + +@pytest.fixture(scope="module") +def image(test_dir): + img_path = test_dir / "render_mm_test.png" arr = np.zeros((128, 128, 3), dtype=np.uint8) arr[:64, :, 0] = 200 arr[64:, :, 2] = 200 arr[32:96, 32:96, 1] = 180 - Image.fromarray(arr).save(IMG_PATH) - return IMG_PATH + Image.fromarray(arr).save(img_path) + return str(img_path) @pytest.fixture(scope="module") -def render_server(): +def render_server(test_dir): env_endpoint = os.environ.get("RENDER_ENDPOINT") if env_endpoint: yield env_endpoint.rstrip("/") @@ -70,7 +76,7 @@ def render_server(): VLLM_PYTHON, "-m", "vllm.entrypoints.cli.main", "launch", "render", "--model", MODEL, "--port", str(port), "--host", "127.0.0.1", "--max-model-len", "4096", "--gpu-memory-utilization", "0.3", - "--allowed-local-media-path", MEDIA_DIR, + "--allowed-local-media-path", str(test_dir), ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, env={**os.environ, "HF_HUB_OFFLINE": "1"}, From 7c14e759e868dfe7cd3b06b31a03c254f5503a2a Mon Sep 17 00:00:00 2001 From: Ranran Haoran Zhang Date: Tue, 7 Jul 2026 21:41:13 -0500 Subject: [PATCH 12/12] style: apply ruff format to prepare_data.py and render test Fixes the failing quality-checks CI job (ruff format --check). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01WTsu9vFfCfPsTTE8PgfEyi Signed-off-by: Ranran Haoran Zhang --- scripts/prepare_data.py | 2 +- .../datagen/test_render_multimodal.py | 24 +++++++++++++++---- 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/scripts/prepare_data.py b/scripts/prepare_data.py index c86054112..3fb39e6aa 100644 --- a/scripts/prepare_data.py +++ b/scripts/prepare_data.py @@ -124,7 +124,7 @@ def parse_args(): default=None, help=( "JSON object of extra kwargs forwarded to the chat template, " - 'e.g. \'{"enable_thinking": true}\'. Only supported on the render ' + "e.g. '{\"enable_thinking\": true}'. Only supported on the render " "path (requires --vllm-server-url)." ), ) diff --git a/tests/integration/datagen/test_render_multimodal.py b/tests/integration/datagen/test_render_multimodal.py index 7bcb11144..5488cb508 100644 --- a/tests/integration/datagen/test_render_multimodal.py +++ b/tests/integration/datagen/test_render_multimodal.py @@ -73,12 +73,26 @@ def render_server(test_dir): port = _find_free_port() proc = subprocess.Popen( # noqa: S603 [ - VLLM_PYTHON, "-m", "vllm.entrypoints.cli.main", "launch", "render", - "--model", MODEL, "--port", str(port), "--host", "127.0.0.1", - "--max-model-len", "4096", "--gpu-memory-utilization", "0.3", - "--allowed-local-media-path", str(test_dir), + VLLM_PYTHON, + "-m", + "vllm.entrypoints.cli.main", + "launch", + "render", + "--model", + MODEL, + "--port", + str(port), + "--host", + "127.0.0.1", + "--max-model-len", + "4096", + "--gpu-memory-utilization", + "0.3", + "--allowed-local-media-path", + str(test_dir), ], - stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, env={**os.environ, "HF_HUB_OFFLINE": "1"}, ) url = f"http://127.0.0.1:{port}"