From c4a56fa5bab1b36e17c9903077d3de8ff63f561c Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Wed, 5 Aug 2026 05:49:09 -0700 Subject: [PATCH 1/2] fix(specdec): correct resume and bound staging in the vLLM hidden-state 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 --- CHANGELOG.rst | 2 + .../compute_hidden_states_vllm.py | 124 +++++++++++------- 2 files changed, 78 insertions(+), 48 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7a188669cb8..68ec4e6e6d5 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -21,6 +21,8 @@ Changelog **Bug Fixes** +- Fix two issues in the vLLM offline hidden-state dump (``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py``) that only surface on large runs. **Resume:** the filter that skips conversations whose ``.pt`` already exists now runs with ``load_from_cache_file=False``. It depends on on-disk state, which is not part of the fingerprint ``datasets`` computes from the function and the dataset, so with a persistent HF cache reused across a resumed or requeued run the cached "keep everything" result from an earlier run was replayed and the dump re-generated and overwrote conversations it had already finished (observed: tens of thousands of ``.pt`` rewritten while the output count stayed flat). **Staging:** generation is now chunked (``--save-chunk-size``, default 256), so each chunk is saved and its staged hidden states freed before the next chunk is generated. Previously the whole dataset was generated before anything was saved, which kept every conversation staged in the connector's ``shared_storage_path`` (``/dev/shm``, i.e. RAM, by default) at once and exhausted it partway through large dumps. Chunking also makes the dump incrementally durable, so an interrupted run keeps its finished conversations and resumes from them. + 0.46 (2026-08-xx) ^^^^^^^^^^^^^^^^^ diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py index 77441f8f858..e2240c7df71 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py @@ -110,6 +110,15 @@ def parse_args() -> argparse.Namespace: "--trust_remote_code", action="store_true", help="Trust remote code for HF models." ) parser.add_argument("--tp", type=int, default=None, help="Tensor parallel size.") + parser.add_argument( + "--save-chunk-size", + type=int, + default=256, + help="Number of conversations to generate before saving and freeing their staged " + "hidden states. Bounds how much the KV connector stages at once (its shared_storage_path " + "defaults to /dev/shm, i.e. RAM); larger values amortize generate() calls, smaller ones " + "use less staging space.", + ) parser.add_argument( "--debug-max-num-conversations", type=int, default=None, help="Limit conversations." ) @@ -159,7 +168,13 @@ def keep_conversation(entry): return not (output_dir / f"{conversation_id}.pt").exists() original_num = len(dataset) - dataset = dataset.filter(keep_conversation) + # load_from_cache_file=False is required for correctness here: keep_conversation depends on + # on-disk state (which .pt files exist), which is NOT part of the cache fingerprint that + # datasets computes from the function and the dataset. With a persistent HF cache (e.g. + # HF_HOME on shared storage, reused across a resumed/requeued run) a cached result from an + # earlier run -- when fewer or no .pt files existed -- is reused, so the filter keeps + # everything and the run re-dumps and overwrites conversations it already finished. + dataset = dataset.filter(keep_conversation, load_from_cache_file=False) print(f"Removed {original_num - len(dataset)} conversations due to existing output files") if args.debug_max_num_conversations is not None: @@ -276,56 +291,69 @@ def keep_conversation(entry): ) # max_tokens=1: we only need a single forward pass over the prompt tokens. - outputs = llm.generate(prompts, SamplingParams(max_tokens=1)) + # + # Generate and save in chunks of --save-chunk-size rather than generating everything up + # front: the connector stages each conversation's hidden states under its + # shared_storage_path (default /dev/shm, i.e. RAM) and they are only freed by + # cleanup_hidden_states() once saved. Generating the whole dataset first keeps every + # conversation staged simultaneously, which exhausts that space on large runs. Chunking + # bounds staging to one chunk, and makes the dump incrementally durable so an interrupted + # run keeps its finished conversations and can resume. + sampling_params = SamplingParams(max_tokens=1) + chunk_size = max(1, args.save_chunk_size) # Save in the same format as compute_hidden_states_hf.py, including loss_mask. num_success = 0 - for conv_id, loss_mask, output in tqdm( - zip(conversation_ids, loss_masks, outputs), total=len(outputs), desc="Saving" - ): - hidden_states_path = output.kv_transfer_params.get("hidden_states_path") - if hidden_states_path is None: - print(f"WARNING: no hidden_states_path for conversation {conv_id}; skipping") - continue - - obj = example_hidden_states_connector.load_hidden_states(hidden_states_path) - token_ids = obj["token_ids"] - # hidden_states: [num_tokens, num_extracted_layers, hidden_size], ordered to match - # extract_layer_ids. Last layer = final output; the rest = aux layers. - hidden_states = obj["hidden_states"] - - output_hidden_states = hidden_states[:, -1, :].cpu() - if hidden_states.shape[1] > 1: - # Concatenate aux layers along the hidden dim, matching the HF dump format. - aux = hidden_states[:, :-1, :].cpu() - aux_hidden_states = aux.reshape(aux.shape[0], -1) - else: - aux_hidden_states = torch.empty(0) - - # loss_mask is sliced to the dumped length below; a shorter loss_mask would slice - # to itself and silently misalign with the hidden states, so guard explicitly. - n_hs = output_hidden_states.shape[0] - if loss_mask.shape[0] < n_hs: - print( - f"WARNING: {conv_id}: loss_mask ({loss_mask.shape[0]}) shorter than hidden " - f"states ({n_hs}); skipping to avoid misalignment" - ) - continue - - 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, - ) - example_hidden_states_connector.cleanup_hidden_states(hidden_states_path) - num_success += 1 + progress = tqdm(total=len(prompts), desc="Dumping") + for chunk_start in range(0, len(prompts), chunk_size): + chunk = slice(chunk_start, chunk_start + chunk_size) + outputs = llm.generate(prompts[chunk], sampling_params) + for conv_id, loss_mask, output in zip(conversation_ids[chunk], loss_masks[chunk], outputs): + progress.update(1) + hidden_states_path = output.kv_transfer_params.get("hidden_states_path") + if hidden_states_path is None: + print(f"WARNING: no hidden_states_path for conversation {conv_id}; skipping") + continue + + obj = example_hidden_states_connector.load_hidden_states(hidden_states_path) + token_ids = obj["token_ids"] + # hidden_states: [num_tokens, num_extracted_layers, hidden_size], ordered to match + # extract_layer_ids. Last layer = final output; the rest = aux layers. + hidden_states = obj["hidden_states"] + + output_hidden_states = hidden_states[:, -1, :].cpu() + if hidden_states.shape[1] > 1: + # Concatenate aux layers along the hidden dim, matching the HF dump format. + aux = hidden_states[:, :-1, :].cpu() + aux_hidden_states = aux.reshape(aux.shape[0], -1) + else: + aux_hidden_states = torch.empty(0) + + # loss_mask is sliced to the dumped length below; a shorter loss_mask would slice + # to itself and silently misalign with the hidden states, so guard explicitly. + n_hs = output_hidden_states.shape[0] + if loss_mask.shape[0] < n_hs: + print( + f"WARNING: {conv_id}: loss_mask ({loss_mask.shape[0]}) shorter than hidden " + f"states ({n_hs}); skipping to avoid misalignment" + ) + continue + + 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, + ) + example_hidden_states_connector.cleanup_hidden_states(hidden_states_path) + num_success += 1 + progress.close() print(f"Successfully processed {num_success} out of {len(prompts)} conversations.") From abbfb95b4cf6daa316cfe20951a13b98d132cc96 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Wed, 5 Aug 2026 06:04:29 -0700 Subject: [PATCH 2/2] =?UTF-8?q?fix(specdec):=20address=20review=20?= =?UTF-8?q?=E2=80=94=20always=20free=20staging,=20validate=20ids=20and=20c?= =?UTF-8?q?hunk=20size?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- CHANGELOG.rst | 2 +- .../compute_hidden_states_vllm.py | 118 +++++++++++------- 2 files changed, 75 insertions(+), 45 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 68ec4e6e6d5..d7c42c7b9f5 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -21,7 +21,7 @@ Changelog **Bug Fixes** -- Fix two issues in the vLLM offline hidden-state dump (``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py``) that only surface on large runs. **Resume:** the filter that skips conversations whose ``.pt`` already exists now runs with ``load_from_cache_file=False``. It depends on on-disk state, which is not part of the fingerprint ``datasets`` computes from the function and the dataset, so with a persistent HF cache reused across a resumed or requeued run the cached "keep everything" result from an earlier run was replayed and the dump re-generated and overwrote conversations it had already finished (observed: tens of thousands of ``.pt`` rewritten while the output count stayed flat). **Staging:** generation is now chunked (``--save-chunk-size``, default 256), so each chunk is saved and its staged hidden states freed before the next chunk is generated. Previously the whole dataset was generated before anything was saved, which kept every conversation staged in the connector's ``shared_storage_path`` (``/dev/shm``, i.e. RAM, by default) at once and exhausted it partway through large dumps. Chunking also makes the dump incrementally durable, so an interrupted run keeps its finished conversations and resumes from them. +- Fix two issues in the vLLM offline hidden-state dump (``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py``) that only surface on large runs. **Resume:** the filter that skips conversations whose ``.pt`` already exists now runs with ``load_from_cache_file=False``. It depends on on-disk state, which is not part of the fingerprint ``datasets`` computes from the function and the dataset, so with a persistent HF cache reused across a resumed or requeued run the cached "keep everything" result from an earlier run was replayed and the dump re-generated and overwrote conversations it had already finished (observed: tens of thousands of ``.pt`` rewritten while the output count stayed flat). **Staging:** generation is now chunked (``--save-chunk-size``, default 256), so each chunk is saved and its staged hidden states freed before the next chunk is generated. Previously the whole dataset was generated before anything was saved, which kept every conversation staged in the connector's ``shared_storage_path`` (``/dev/shm``, i.e. RAM, by default) at once and exhausted it partway through large dumps. Chunking also makes the dump incrementally durable, so an interrupted run keeps its finished conversations and resumes from them. The save path now also frees each conversation's staged hidden states in a ``finally``, so a conversation skipped mid-loop (e.g. a short ``loss_mask``) can no longer leak its staging file, and conversation ids are validated as plain filenames before being used to build output paths. 0.46 (2026-08-xx) ^^^^^^^^^^^^^^^^^ diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py index e2240c7df71..68cd7874c3c 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py @@ -89,6 +89,33 @@ def _resolve_aux_layers_standalone( return ids +def _conversation_id(entry) -> str: + """Return an entry's conversation id, rejecting values that are not a plain filename. + + The id is used directly as the output filename (``.pt``), so an absolute path or one + containing a separator or ``..`` would resolve outside ``--output-dir``. The dataset is + external input, so validate once here, where it is read, rather than at each use. + """ + conversation_id = entry.get("conversation_id", entry.get("uuid", None)) + if conversation_id is None: + raise ValueError("conversation_id is required") + conversation_id = str(conversation_id) + if conversation_id in {"", ".", ".."} or conversation_id != Path(conversation_id).name: + raise ValueError( + f"conversation_id {conversation_id!r} is not a usable filename: it must not be " + "empty, '.', '..', or contain a path separator." + ) + return conversation_id + + +def _positive_int(value: str) -> int: + """argparse type for options that must be >= 1.""" + parsed = int(value) + if parsed < 1: + raise argparse.ArgumentTypeError(f"must be a positive integer, got {value}") + return parsed + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="""Collect hidden states from conversations using vLLM's native extractor.""" @@ -112,7 +139,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--tp", type=int, default=None, help="Tensor parallel size.") parser.add_argument( "--save-chunk-size", - type=int, + type=_positive_int, default=256, help="Number of conversations to generate before saving and freeing their staged " "hidden states. Bounds how much the KV connector stages at once (its shared_storage_path " @@ -163,9 +190,7 @@ def main(args: argparse.Namespace) -> None: output_dir.mkdir(parents=True, exist_ok=True) def keep_conversation(entry): - conversation_id = entry.get("conversation_id", entry.get("uuid", None)) - assert conversation_id is not None, "conversation_id is required" - return not (output_dir / f"{conversation_id}.pt").exists() + return not (output_dir / f"{_conversation_id(entry)}.pt").exists() original_num = len(dataset) # load_from_cache_file=False is required for correctness here: keep_conversation depends on @@ -213,7 +238,7 @@ def keep_conversation(entry): num_invalid = 0 for entry in dataset: - conversation_id = entry.get("conversation_id", entry.get("uuid")) + conversation_id = _conversation_id(entry) # Accept either the "conversations" or OpenAI-style "messages" key (the # MiniMax synthetic data uses "messages"). conversations = entry.get("conversations") or entry.get("messages") @@ -300,7 +325,7 @@ def keep_conversation(entry): # bounds staging to one chunk, and makes the dump incrementally durable so an interrupted # run keeps its finished conversations and can resume. sampling_params = SamplingParams(max_tokens=1) - chunk_size = max(1, args.save_chunk_size) + chunk_size = args.save_chunk_size # Save in the same format as compute_hidden_states_hf.py, including loss_mask. num_success = 0 @@ -315,44 +340,49 @@ def keep_conversation(entry): print(f"WARNING: no hidden_states_path for conversation {conv_id}; skipping") continue - obj = example_hidden_states_connector.load_hidden_states(hidden_states_path) - token_ids = obj["token_ids"] - # hidden_states: [num_tokens, num_extracted_layers, hidden_size], ordered to match - # extract_layer_ids. Last layer = final output; the rest = aux layers. - hidden_states = obj["hidden_states"] - - output_hidden_states = hidden_states[:, -1, :].cpu() - if hidden_states.shape[1] > 1: - # Concatenate aux layers along the hidden dim, matching the HF dump format. - aux = hidden_states[:, :-1, :].cpu() - aux_hidden_states = aux.reshape(aux.shape[0], -1) - else: - aux_hidden_states = torch.empty(0) - - # loss_mask is sliced to the dumped length below; a shorter loss_mask would slice - # to itself and silently misalign with the hidden states, so guard explicitly. - n_hs = output_hidden_states.shape[0] - if loss_mask.shape[0] < n_hs: - print( - f"WARNING: {conv_id}: loss_mask ({loss_mask.shape[0]}) shorter than hidden " - f"states ({n_hs}); skipping to avoid misalignment" - ) - continue - - 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, - ) - example_hidden_states_connector.cleanup_hidden_states(hidden_states_path) - num_success += 1 + # Staged hidden states must be freed on every path, not just the success one: + # an early `continue` (e.g. a short loss_mask) would otherwise leak this + # conversation's staging file, and enough of them exhaust shared_storage_path. + try: + obj = example_hidden_states_connector.load_hidden_states(hidden_states_path) + token_ids = obj["token_ids"] + # hidden_states: [num_tokens, num_extracted_layers, hidden_size], ordered to match + # extract_layer_ids. Last layer = final output; the rest = aux layers. + hidden_states = obj["hidden_states"] + + output_hidden_states = hidden_states[:, -1, :].cpu() + if hidden_states.shape[1] > 1: + # Concatenate aux layers along the hidden dim, matching the HF dump format. + aux = hidden_states[:, :-1, :].cpu() + aux_hidden_states = aux.reshape(aux.shape[0], -1) + else: + aux_hidden_states = torch.empty(0) + + # loss_mask is sliced to the dumped length below; a shorter loss_mask would slice + # to itself and silently misalign with the hidden states, so guard explicitly. + n_hs = output_hidden_states.shape[0] + if loss_mask.shape[0] < n_hs: + print( + f"WARNING: {conv_id}: loss_mask ({loss_mask.shape[0]}) shorter than hidden " + f"states ({n_hs}); skipping to avoid misalignment" + ) + continue + + 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, + ) + num_success += 1 + finally: + example_hidden_states_connector.cleanup_hidden_states(hidden_states_path) progress.close() print(f"Successfully processed {num_success} out of {len(prompts)} conversations.")