update preprocess.py to use vllm render endpoint - #975
Conversation
Replace the two-tier local masking path (HF {% generation %} assistant-mask
probe + regex span detector) with a single mechanism: render each conversation
through the training vLLM instance's /v1/chat/completions/render endpoint and
take the loss mask from the boundary between the prompt render and the full
render of each assistant turn. Append-only templates pack to one row per
conversation; history-rewriting (reasoning) templates fan out to one row per
turn. On-policy pre-tokenized rows (vllm-project#729) still pass straight through.
One tokenizer now produces the mask, the target hidden states, and the serving
prompt. Deletes ~270 lines of regex/tag masking and the local HF render path;
adds render_client.py (token_ids only, no truncation) and boundary derivation.
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
- drop redundant early-return in _render_boundary_rows (subsumed by the no-turns guard) - pass an is_multimodal bool into _preprocess_batch instead of the heavyweight processor, so the map closure no longer pickles it into every worker Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
render_client.py imports httpx; it was only present transitively, so a clean install would fail to resolve it. Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Tests: - drop test_render_conversation_returns_token_ids (a one-line passthrough; the 200 path is covered by the missing-token_ids test) - drop test_context_filling_window_yields_no_rows (single-guard edge) - trim test_pretokenized_dataset_skips_render to its load-bearing assertion Prose: - condense the ported _render_boundary_rows docstring and the verbose comments/docstrings in preprocessing.py and render_client.py Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
… vLLM
Drop the packed branch. Every assistant turn now gets its own row carrying
the history re-rendered the way inference sees it, for every template.
Packing was an optimization, never a correctness requirement: it only held
for append-only templates, and thinking models (Qwen3 strips <think> from
history) already fanned out. Removing it deletes _Turn, the append_only
chain check, and the two-pass structure -- one path instead of three.
The cost lands on append-only templates, measured on 10 real multi-turn
ShareGPT conversations rendered by Qwen2.5-0.5B-Instruct:
packed 10 rows, 8606 tokens
fan-out 38 rows, 20449 tokens (3.8x rows, 2.4x tokens)
Qwen3-0.6B output is byte-identical before and after (37 rows), since it
took the fan-out path already.
The <think>-scaffold LCP fallback stays. It is load-bearing, not defensive:
DeepSeek-R1-Distill pre-fills `<think>\n` in the generation prompt, and
Qwen3.5 pre-fills an empty `<think></think>` that recorded reasoning then
contradicts. Both break prefix-stability and land in that branch.
Also warn when --seq-length silently truncates supervision. A row clipped
mid-response keeps thousands of supervised tokens, so it passes the
existing zero-supervision check and looks healthy while its tail and
closing tokens are never learned. On 40 real gsm8k rows with recorded
reasoning at --seq-length 2048, 13 are clipped, discarding a median 34%
of the assistant turn (worst: 1482 of 3530 tokens) -- previously with no
output at all. num_clipped counts only kept-but-truncated rows; rows
clipped past their boundary already report as unsupervised and would
otherwise be counted twice.
Tests:
- add tests/e2e/smoke/test_render_boundary.py -- the first test to exercise
the render transport for real (URL, token_ids key, boundary from what vLLM
actually returns). Qwen3-0.6B, live server, one conversation covering
fan-out, the leading-assistant exclusion, and the trailing-user drop.
- drop the tests that faked the render and are now covered there: the packed
and fan-out routing tests, first-turn-context-only, trailing-dropped, and
test_build_eagle3_dataset_packed_end_to_end (which named itself end-to-end
while pointing at http://fake, and used a template that structurally could
not reach fan-out).
- keep what a live server cannot produce on demand: the scaffold fallback,
the unstable guard, and the client's error paths.
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Both examples pass `--data sharegpt`, which is off-policy, so as written they
now fail immediately:
ValueError: render_endpoint is required to derive loss masks for
off-policy conversations. Pass --render-endpoint pointing at a vLLM server.
The flag was documented in the argument reference but never reached the
copy-pasteable commands. Note the requirement in Basic Usage as well, since
it is conditional -- pre-tokenized input still needs no endpoint.
Reported by CodeRabbit on vllm-project#794.
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
- `assert sum(mask) > 0`: _append_row returns "unsupervised" and does not append when num_valid_tokens == 0, so every row reaching the dataset satisfies this by construction. - `assert len(ids) == len(mask)`: the mask is built as [0]*boundary + [1]*(len(full_ids) - boundary) and both are clipped by the same slice, so they cannot diverge -- and `zip(..., strict=True)` two lines below already raises on skew. - unused `logging` import and module logger. Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
408 and 429 fall inside 4xx, so they raised InvalidResponseError, which short-circuits @with_retries: a proxy shedding load dropped the conversation on the first 429 instead of backing off. Carve them out so they reach the retry-eligible RenderError; the rest of 4xx still short-circuits. vLLM emits neither, but --render-endpoint takes any URL and deployments front it with proxies that do. Reported by CodeRabbit on vllm-project#794. Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
…ports Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
…sation Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
…ing_dataset The prepared dataset is not eagle3-specific: prepare_data.py takes no speculator type, and train.py builds eagle3, dflash, dspark, peagle, and mtp from the same rows. Rename the function and drop EAGLE3 from the two docstrings that described its output. Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Point --render-endpoint at the launch_vllm.py instance rather than a second server, and flag the URL-form difference: data_generation_offline takes an /v1-suffixed endpoint, while this one is appended to and 404s on that form. Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
The argparse help still carried the pre-change wording. Also scope the requirement correctly: --data is repeatable and the check runs per dataset, so a single raw input among pre-tokenized ones still needs the endpoint. Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
mypy rejects passing None where ProcessorLike is expected. Neither test reads the processor: the missing-endpoint guard raises before it is used, and pre-tokenized rows skip preprocessing. Name that with a cast sentinel rather than an ignore. Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: shanjiaz <zsjwpianpian@gmail.com>
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe PR replaces processor-based conversation preprocessing with vLLM render-boundary processing. It adds pretokenized-row passthrough, updates the preparation CLI and response-regeneration workflow, and revises documentation and end-to-end coverage. ChangesRender-based speculator data pipeline
Possibly related issues
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require approval from approved reviewers listWaiting for any of
This rule is failing.All pull requests must have at least one approving review from a member of the approved reviewers list before merging.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (2)
src/speculators/data_generation/preprocessing.py (1)
234-267: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winCache renders per conversation to remove redundant round trips.
Each assistant turn issues 2-3 blocking HTTP requests, and every request re-renders the whole prefix. The history render for turn
jis identical to the full render already computed for turnj-1when turnj-1is an assistant turn. A small per-conversation memo removes those duplicate requests and keeps the boundary logic unchanged.♻️ Proposed memoization inside `_render_boundary_rows`
rows: list[BoundaryRow] = [] + cache: dict[tuple[int, bool], list[int]] = {} + + def render(prefix_len: int, *, add_generation_prompt: bool) -> list[int]: + key = (prefix_len, add_generation_prompt) + if key not in cache: + cache[key] = _encode_render( + normalized_conv[:prefix_len], + render_endpoint, + add_generation_prompt=add_generation_prompt, + tools=tools, + ) + return cache[key] 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 - prompt_ids = _encode_render( - normalized_conv[:j], - render_endpoint, - add_generation_prompt=True, - tools=tools, - ) + prompt_ids = render(j, add_generation_prompt=True) 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, - ) + full_ids = render(j + 1, add_generation_prompt=False) if full_ids[: len(prompt_ids)] == prompt_ids: boundary = len(prompt_ids) else: # 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, - tools=tools, - ) + hist_ids = render(j, add_generation_prompt=False)The stub in tests/integration/datagen/test_render_boundary.py keys on
(len(conv_prefix), add_generation_prompt), so it stays compatible.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/speculators/data_generation/preprocessing.py` around lines 234 - 267, Update _render_boundary_rows to memoize _encode_render results per conversation, keyed by the conversation-prefix length and add_generation_prompt value. Reuse the previous assistant turn’s full render as the current turn’s history render, while caching all other renders through the same memo; preserve the existing boundary and common-prefix logic unchanged.tests/integration/datagen/test_render_boundary.py (1)
181-204: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a test for the render-failure skip path.
No test reaches the broad catch in
_render_conversation_rowsat preprocessing.py line 334. That branch implements the stated behavior "rendering failures skip only the affected conversation", and the e2e test uses a healthy server, so it cannot cover it. The existing_encode_renderseam makes the test cheap.💚 Proposed test
def test_render_failure_skips_only_that_conversation(monkeypatch): # A conversation the endpoint cannot render must not abort the batch. def fake(conv_prefix, render_endpoint, *, add_generation_prompt, tools=None): if conv_prefix[0]["content"] == "boom": raise RuntimeError("template exploded") return [1, 2, 3] if add_generation_prompt else [1, 2, 3, 4, 5] monkeypatch.setattr(preprocessing, "_encode_render", fake) out = preprocessing._preprocess_batch( { "conversations": [ [{"role": "user", "content": "boom"}, {"role": "assistant", "content": "x"}], [{"role": "user", "content": "ok"}, {"role": "assistant", "content": "y"}], ] }, is_multimodal=False, render_endpoint="http://x", max_length=100, ) assert len(out["input_ids"]) == 1 assert out["loss_mask"][0].tolist() == [0, 0, 0, 1, 1]As per path instructions: "Check that new code paths introduced in the PR are covered."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/datagen/test_render_boundary.py` around lines 181 - 204, Add a test named test_render_failure_skips_only_that_conversation that monkeypatches preprocessing._encode_render to raise for one conversation and return token sequences for the other, then calls preprocessing._preprocess_batch with a render endpoint and asserts only the successful conversation remains with the expected input_ids and loss_mask.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/cli/response_regeneration.md`:
- Line 143: Update docs/user_guide/tutorials/response_regeneration.md at line 81
to document the review-only text field emitted by
scripts/response_regeneration/script.py, replacing conversations in both the
output JSON example and its description so the tutorial matches the generated
schema. The anchor in docs/cli/response_regeneration.md at line 143 requires no
direct change.
In `@docs/user_guide/tutorials/train.md`:
- Around line 131-133: Update the --data option description to use the complete
sentence “You can supply this option multiple times to combine datasets.” while
preserving the existing explanation of supported data formats.
In `@scripts/prepare_data.py`:
- Around line 139-151: Add a CLI validator for the --render-endpoint argument
that accepts only well-formed HTTP(S) base URLs and rejects any URL whose path
ends with /v1, including trailing-slash variants. Wire the validator into the
argument definition so invalid values are rejected during argument parsing,
before preprocessing or render_client.py URL construction begins.
In `@src/speculators/data_generation/preprocessing.py`:
- Around line 413-421: Update
build_speculator_training_dataset/_preprocess_batch so pretokenized rows
containing multimodal content are not silently passed through
_passthrough_pretokenized without their messages; either reject this combination
explicitly or preserve the source messages for downstream hidden-state
extraction and vLLM processing. Ensure the processor’s multimodal path is not
ignored for rows identified by input_ids and loss_mask.
- Around line 324-336: Update the exception handling around
_render_boundary_rows to re-raise transport-level endpoint failures, including
connection/unavailability errors and persistent InvalidResponseError, while
retaining the broad catch that logs and skips conversation-specific template or
boundary failures. Ensure endpoint outages propagate to
load_and_preprocess_dataset instead of producing empty output.
In `@src/speculators/data_generation/render_client.py`:
- Around line 42-53: The render request path in render_conversation should reuse
an httpx.Client rather than calling httpx.post for every render. Add a pooled
client per worker/process and use it for the POST while preserving the existing
URL, payload, and timeout; update
tests/integration/datagen/test_render_boundary.py to patch the client instance
or its post method instead of render_client.httpx.post.
---
Nitpick comments:
In `@src/speculators/data_generation/preprocessing.py`:
- Around line 234-267: Update _render_boundary_rows to memoize _encode_render
results per conversation, keyed by the conversation-prefix length and
add_generation_prompt value. Reuse the previous assistant turn’s full render as
the current turn’s history render, while caching all other renders through the
same memo; preserve the existing boundary and common-prefix logic unchanged.
In `@tests/integration/datagen/test_render_boundary.py`:
- Around line 181-204: Add a test named
test_render_failure_skips_only_that_conversation that monkeypatches
preprocessing._encode_render to raise for one conversation and return token
sequences for the other, then calls preprocessing._preprocess_batch with a
render endpoint and asserts only the successful conversation remains with the
expected input_ids and loss_mask.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: c9b8d4a1-27d7-445a-b6c5-d1f9acb6e12d
📒 Files selected for processing (25)
docs/cli/prepare_data.mddocs/cli/response_regeneration.mddocs/user_guide/tutorials/response_regeneration.mddocs/user_guide/tutorials/train.mdpyproject.tomlscripts/prepare_data.pyscripts/response_regeneration/script.pysrc/speculators/data_generation/preprocessing.pysrc/speculators/data_generation/render_client.pytests/e2e/regression/test_eagle3_offline_acceptance.pytests/e2e/regression/test_eagle3_online_acceptance.pytests/e2e/smoke/test_finetuning_sanity.pytests/e2e/smoke/test_mooncake_online_training.pytests/e2e/smoke/test_mtp_finetuning.pytests/e2e/smoke/test_offline_training.pytests/e2e/smoke/test_online_training.pytests/e2e/smoke/test_render_boundary.pytests/e2e/smoke/test_resume_optimizer.pytests/e2e/smoke/test_training_only.pytests/e2e/utils.pytests/integration/datagen/test_preprocessing.pytests/integration/datagen/test_regex_patterns.pytests/integration/datagen/test_render_boundary.pytests/unit/scripts/test_response_regeneration.pytests/unit/train/test_prepare_data.py
💤 Files with no reviewable changes (2)
- tests/integration/datagen/test_regex_patterns.py
- tests/integration/datagen/test_preprocessing.py
speculatorsbot
left a comment
speculatorsbot
left a comment
There was a problem hiding this comment.
Review Summary
The non-test changes (source code, docs, CLI) carry over from the previously approved PR #794 and are not reviewed again here. This review focuses on the test-only changes authored in this PR.
What the test changes do
-
E2e tests switch from off-policy
"sharegpt"to pre-tokenized on-policy datasets (hf:inference-optimization/speculators-ci-datasets:smoke_regen/tutorial_regen). This aligns the test data with the new pipeline contract that requires target-model responses. -
Data preparation moves inside the vLLM server context in
test_offline_training.py,test_online_training.py, andtest_mtp_finetuning.py. This is necessary because conversation-format datasets now need the render endpoint, which is served by the same vLLM instance used for hidden-state extraction. Pre-tokenized datasets pass through without rendering, so the endpoint is harmless when supplied but unnecessary. -
test_resume_optimizer.pycorrectly omitsrender_endpointsincesmoke_regenis pre-tokenized and no server is running at that point in the test. -
test_training_only.pyandtest_finetuning_sanity.pyconsolidate ontoinference-optimization/speculators-ci-datasets, replace--legacy-datawith--on-missing raise, and add--hidden-states-pathto match the current training script contract. The removal of_generate_t2d_d2tand the vocab mapping imports is consistent with the new data format bundling these. -
New
test_render_boundary.py(e2e smoke) exercises the live vLLM/renderendpoint with Qwen3-0.6B. Good choice of model -- its template rewrites history (strips<think>from past turns), which forces the fan-out code path. The conversation is well-designed: a leading assistant turn (no row), two bounded assistant turns (one row each), and a trailing user turn (dropped). The assertions check the core invariant (user tokens never supervised, assistant tokens are) and verify per-row isolation and context re-rendering. -
New
test_render_boundary.py(integration) covers the edge cases a live server cannot produce on demand: scaffold LCP fallback (DeepSeek-R1 / Qwen3.5), the unstable-boundary guard, over-length turn skipping without dropping later turns, and the render client's error/retry paths. The_patch_encodehelper is clean and the keying by(prefix_len, add_generation_prompt)makes the stubs easy to trace. -
Deleted tests (
test_regex_patterns.py, regex/HF-mask tests intest_preprocessing.py) are correct removals -- these tested_detect_assistant_pattern,_create_loss_mask_from_offsets,_supports_assistant_mask, and_preprocess_batchwith the old signature, all of which are removed in the source changes. -
test_response_regeneration.pyupdates_preprocess_batchcall sites to use the new signature (is_multimodal=False, render_endpoint=Noneinstead ofprocessor=None, assistant_pattern=None). The passthrough contract is preserved. -
tests/e2e/utils.pyadds the optionalrender_endpointparameter torun_prepare_data, forwarded as--render-endpointwhen set. Clean and minimal.
Observations
-
The
hf:spechf:inference-optimization/speculators-ci-datasets:smoke_regenis parsed by_load_hf_datasetassplit="smoke_regen"(no subset). CI confirms this resolves correctly. If the HF dataset is later reorganized to use named configs instead of split-level partitioning, these specs would need to become three-part (hf:id:subset:train). Not a blocker -- just noting the coupling. -
The deleted
test_load_and_preprocess_dataset_shuffles_combined_datasetstest verified that combined multi-dataset inputs are shuffled beforemax_samplestruncation. That shuffle logic is unchanged in the new code, so coverage is only reduced, not broken. Worth noting if future refactors touch the combine-and-truncate path.
LGTM. The test updates are consistent, well-commented, and properly exercise both the new render-boundary pipeline and the speculator-format passthrough path.
4 participants
Purpose
This was originally PR 794. Changes are previously reviewed and approved. This PR updated e2e tests to use updated data flow & format.
Tests
Ran smoke tests locally, all passed, all e2e test run in progress.
run_prepare_dataintests/e2e/utils.pynow has an optionalrender_endpointparametertest_regex_patterns.py(regex masking removed)llm-compressor-testing can't pick up Ranran's branch. Ran tests locally:
nightly (13 passed 16 skipped):
Checklist
I have filled in: