fix(specdec): correct resume and bound staging in the vLLM hidden-state dump - #2080
Open
yeyu-nvidia wants to merge 2 commits into
Open
fix(specdec): correct resume and bound staging in the vLLM hidden-state dump#2080yeyu-nvidia wants to merge 2 commits into
yeyu-nvidia wants to merge 2 commits into
Conversation
…te dump
Two issues in compute_hidden_states_vllm.py that only surface on large offline dumps.
Resume re-processed finished work: keep_conversation skips conversations whose .pt already
exists, but that depends on on-disk state, which is not part of the fingerprint datasets
computes from the function and the dataset. With a persistent HF cache reused across a resumed
or requeued run, the cached 'keep everything' result from an earlier run (when fewer or no .pt
existed) was replayed, so the dump re-generated and overwrote conversations it had already
finished. Observed on a 194k-conversation dump: tens of thousands of .pt rewritten over a
multi-hour window with the output count flat. Pass load_from_cache_file=False so the filter
re-checks the disk each run.
Staging exhausted /dev/shm: the whole dataset was generated before anything was saved, so every
conversation's hidden states stayed staged in the connector's shared_storage_path (/dev/shm,
i.e. RAM, by default) until the single save loop freed them. Large dumps ran out of space
partway through ('No space left on device' from the connector's safetensors write). Generate
and save in chunks of --save-chunk-size (default 256) so at most one chunk is staged. This also
makes the dump incrementally durable: an interrupted run keeps its finished conversations and
resumes from them.
Signed-off-by: Ye Yu <yeyu@nvidia.com>
Contributor
📝 WalkthroughWalkthroughThis change validates conversation IDs, disables cached resume filtering, and processes hidden-state dumps in configurable chunks. The changelog records the new filtering, chunking, and cleanup behavior. ChangesvLLM hidden-state collection flow
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 3
🤖 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
`@examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py`:
- Around line 113-121: Validate --save-chunk-size at argument parsing by
requiring a strictly positive integer, causing argparse to reject zero and
negative values. In the processing logic around the save-chunk-size usage,
remove the silent conversion to 1 and use args.save_chunk_size directly.
- Around line 318-354: Wrap all processing after hidden_states_path is acquired,
including validation and file saving, in a try block, and move
cleanup_hidden_states(hidden_states_path) to a finally block so it runs on every
path, including the short loss_mask continue. Use the existing
hidden_states_path and example_hidden_states_connector symbols without changing
the processing behavior.
- Around line 342-343: Validate and constrain conversation_id at the dataset
boundary before constructing output_file, including the resume check and write
around output_file. Reject absolute paths and traversal components, allowing
only safe identifiers that remain within output_dir, and preserve normal
processing for valid IDs.
🪄 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: Enterprise
Run ID: 4c2e0717-1d85-4d7d-bcd8-a7c8c733f9f0
📒 Files selected for processing (2)
CHANGELOG.rstexamples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py
…chunk size - Free each conversation's staged hidden states in a finally, so a conversation skipped mid-loop (short loss_mask) no longer leaks its staging file. This is the same failure this PR set out to prevent, on an error path. - Validate conversation_id as a plain filename where the dataset is read: it is used directly as the output filename, so an absolute path or one containing a separator or '..' would resolve outside --output-dir. - Reject non-positive --save-chunk-size at argument parsing instead of silently clamping. Signed-off-by: Ye Yu <yeyu@nvidia.com>
Contributor
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 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
`@examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py`:
- Around line 371-382: Update the save path in compute_hidden_states_vllm so
each conversation dump is written atomically: write the torch.save payload to a
temporary file in output_dir first, then replace the final output_file only
after serialization succeeds. Keep the existing output_file naming and payload
structure in the same save block, and anchor the change around the
open(output_file, "wb") / torch.save call site so the resume filter never sees a
partially written .pt file.
🪄 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: Enterprise
Run ID: 2adf9037-fb2d-4ed4-bfc0-0a960c9357fa
📒 Files selected for processing (2)
CHANGELOG.rstexamples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py
🚧 Files skipped from review as they are similar to previous changes (1)
- CHANGELOG.rst
Comment on lines
+371
to
+382
| output_file = output_dir / f"{conv_id}.pt" | ||
| with open(output_file, "wb") as f: | ||
| torch.save( | ||
| { | ||
| "input_ids": token_ids.cpu(), | ||
| "hidden_states": output_hidden_states, | ||
| "aux_hidden_states": aux_hidden_states, | ||
| "loss_mask": loss_mask[: output_hidden_states.shape[0]].cpu(), | ||
| "conversation_id": conv_id, | ||
| }, | ||
| f, | ||
| ) |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Write each completed dump atomically.
The resume filter at Line 193 treats any existing .pt file as complete. torch.save() writes directly to that final path. If serialization fails or the process stops during the write, a truncated file remains. A resumed run then skips that conversation permanently.
Write to a temporary file in output_dir. Replace the final path only after torch.save() succeeds.
Proposed fix
output_file = output_dir / f"{conv_id}.pt"
- with open(output_file, "wb") as f:
+ temp_output_file = output_file.with_suffix(f"{output_file.suffix}.tmp")
+ with open(temp_output_file, "wb") as f:
torch.save(
{
"input_ids": token_ids.cpu(),
"hidden_states": output_hidden_states,
"aux_hidden_states": aux_hidden_states,
"loss_mask": loss_mask[: output_hidden_states.shape[0]].cpu(),
"conversation_id": conv_id,
},
f,
)
+ os.replace(temp_output_file, output_file)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| output_file = output_dir / f"{conv_id}.pt" | |
| with open(output_file, "wb") as f: | |
| torch.save( | |
| { | |
| "input_ids": token_ids.cpu(), | |
| "hidden_states": output_hidden_states, | |
| "aux_hidden_states": aux_hidden_states, | |
| "loss_mask": loss_mask[: output_hidden_states.shape[0]].cpu(), | |
| "conversation_id": conv_id, | |
| }, | |
| f, | |
| ) | |
| output_file = output_dir / f"{conv_id}.pt" | |
| temp_output_file = output_file.with_suffix(f"{output_file.suffix}.tmp") | |
| with open(temp_output_file, "wb") as f: | |
| torch.save( | |
| { | |
| "input_ids": token_ids.cpu(), | |
| "hidden_states": output_hidden_states, | |
| "aux_hidden_states": aux_hidden_states, | |
| "loss_mask": loss_mask[: output_hidden_states.shape[0]].cpu(), | |
| "conversation_id": conv_id, | |
| }, | |
| f, | |
| ) | |
| os.replace(temp_output_file, output_file) |
🧰 Tools
🪛 ast-grep (0.45.0)
[warning] 371-371: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(output_file, "wb")
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').
(open-filename-from-request)
🤖 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
`@examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py`
around lines 371 - 382, Update the save path in compute_hidden_states_vllm so
each conversation dump is written atomically: write the torch.save payload to a
temporary file in output_dir first, then replace the final output_file only
after serialization succeeds. Keep the existing output_file naming and payload
structure in the same save block, and anchor the change around the
open(output_file, "wb") / torch.save call site so the resume filter never sees a
partially written .pt file.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2080 +/- ##
==========================================
- Coverage 67.17% 67.02% -0.15%
==========================================
Files 521 521
Lines 59857 59857
==========================================
- Hits 40206 40117 -89
- Misses 19651 19740 +89
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: Bug fix
Fixes two issues in the vLLM offline hidden-state dump
(
examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py).Both are invisible on small dumps and only bite at scale, which is why they survived until
now — they were found while dumping ~194k conversations for a MiniMax-M3 draft.
1. Resume silently re-processed already-finished work.
keep_conversationskips conversations whose.ptalready exists, but that predicate readson-disk state, which is not part of the fingerprint
datasetscomputes forfilter()(it hashes the function and the dataset). With a persistent HF cache reused across a resumed
or requeued run, the cached "keep everything" result from an earlier run — computed when
few or no
.ptfiles existed — is replayed. The run then re-generates and overwritesconversations it had already completed, and reports
Removed 0 conversations due to existing output fileswhile doing so.Observed on a 194k-conversation dump: ~62k
.ptrewritten over a two-hour window with thetotal output count completely flat.
Fix: pass
load_from_cache_file=Falseso the filter re-checks the disk on every run.2. Staging exhausted
/dev/shmpartway through large dumps.The script generated the entire dataset before saving anything. The KV connector stages
each conversation's hidden states under its
shared_storage_path(/dev/shm, i.e. RAM, bydefault) and they are only freed by
cleanup_hidden_states()in the save loop — so everyconversation stayed staged simultaneously. On a large dump this exhausts the space and the
connector starts failing writes:
Fix: generate and save in chunks of
--save-chunk-size(default 256), so at most one chunkis staged at a time. As a side benefit the dump becomes incrementally durable — an
interrupted run (walltime limit, node failure) keeps its finished conversations and the
resume path above continues from them, instead of losing the whole run's work.
Testing
confirmed both fixes on the same workload: after the change the output count advanced
monotonically across requeues (123k → 194k) with no rewrites, and
/dev/shmstayed boundedthrough completion.
pre-commit run --files ...passes (ruff check/format, mypy, bandit, license, rst checks).the default
--save-chunk-size 256is the only new knob.Additional Information
Extracted from #1749, which is otherwise superseded by the streaming DFlash/DSpark path — these
two fixes are model-agnostic and apply to any offline dump, so they are worth landing on their
own.
Before your PR is "Ready for review"
Summary by CodeRabbit
New Features
Bug Fixes