fix(specdec): correct resume and bound staging in the vLLM hidden-state dump

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. 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)
 ^^^^^^^^^^^^^^^^^
 

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 (``<id>.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."""
@@ -110,6 +137,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=_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 "
+        "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."
     )
@@ -154,12 +190,16 @@ 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)
-    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:
@@ -198,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")
@@ -276,56 +316,74 @@ 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 = 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
+
+            # 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.")