[1/8] Let the hidden-state dump skip aux_hidden_states

examples/speculative_decoding/collect_hidden_states/common.py

@@ -78,6 +78,19 @@ def resolve_aux_layers(args: argparse.Namespace, num_hidden_layers: int) -> list
     return sorted(set(indices))
 
 
+def add_no_aux_hidden_states_args(parser: argparse.ArgumentParser) -> None:
+    """Register the ``--no-aux-hidden-states`` flag on ``parser``."""
+    parser.add_argument(
+        "--no-aux-hidden-states",
+        action="store_true",
+        help=(
+            "Skip dumping aux_hidden_states. External (standalone) draft training never "
+            "reads them, and they dominate the dump: one plane per selected layer against "
+            "a single plane for the base hidden state."
+        ),
+    )
+
+
 def add_answer_only_loss_args(parser: argparse.ArgumentParser) -> None:
     """Register ``--answer-only-loss`` and ``--chat-template`` flags on ``parser``."""
     parser.add_argument(

examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py

@@ -24,6 +24,7 @@
 from common import (
     add_answer_only_loss_args,
     add_aux_layers_args,
+    add_no_aux_hidden_states_args,
     load_chat_template,
     resolve_aux_layers,
     tokenize_with_loss_mask,
@@ -101,6 +102,7 @@ def parse_args() -> argparse.Namespace:
     )
     add_aux_layers_args(parser)
     add_answer_only_loss_args(parser)
+    add_no_aux_hidden_states_args(parser)
 
     return parser.parse_args()
 
@@ -148,9 +150,16 @@ def keep_conversation(entry):
     model = AutoModel.from_pretrained(
         args.model, dtype="auto", device_map="auto", trust_remote_code=args.trust_remote_code
     )
-    num_hidden_layers = getattr(model.config, "num_hidden_layers", None)
+    # Multimodal checkpoints (e.g. *ForConditionalGeneration) keep the language
+    # model's depth under config.text_config.
+    text_config = (
+        getattr(model.config, "text_config", None)
+        or getattr(model.config, "llm_config", None)
+        or model.config
+    )
+    num_hidden_layers = getattr(text_config, "num_hidden_layers", None)
     if num_hidden_layers is None:
-        raise ValueError(f"model.config has no 'num_hidden_layers' attribute: {model.config}")
+        raise ValueError(f"config has no 'num_hidden_layers' attribute: {text_config}")
     selected_layer_ids = resolve_aux_layers(args, num_hidden_layers)
 
     tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=args.trust_remote_code)
@@ -184,8 +193,16 @@ async def dump_hidden_states(
             outputs = model(input_ids=input_ids.to(model.device), output_hidden_states=True)
             # outputs.hidden_states[0] is the embedding output; layer k output is at index k+1.
             hidden_states = outputs.hidden_states
-            aux_hidden_states = torch.cat(
-                [hidden_states[lid + 1].squeeze(0).cpu() for lid in selected_layer_ids], dim=-1
+            # Only build the aux planes when they will be stored: assembling them costs a
+            # concatenation and a device-to-host copy per conversation, which is pure waste
+            # on a dump that is about to drop them.
+            aux_hidden_states = (
+                None
+                if args.no_aux_hidden_states
+                else torch.cat(
+                    [hidden_states[lid + 1].squeeze(0).cpu() for lid in selected_layer_ids],
+                    dim=-1,
+                )
             )
             output_hidden_states = hidden_states[-1].squeeze(0).cpu()
         # Write-then-rename: keep_conversation() skips any existing <id>.pt as finished work,
@@ -199,7 +216,11 @@ async def dump_hidden_states(
                     {
                         "input_ids": input_ids.squeeze(0).cpu(),
                         "hidden_states": output_hidden_states,
-                        "aux_hidden_states": aux_hidden_states,
+                        **(
+                            {}
+                            if args.no_aux_hidden_states
+                            else {"aux_hidden_states": aux_hidden_states}
+                        ),
                         "loss_mask": loss_mask,
                         "conversation_id": conversation_id,
                     },

examples/speculative_decoding/collect_hidden_states/sample_hidden_states.py

@@ -61,18 +61,22 @@ def main(args: argparse.Namespace) -> None:
 
     for i, file in enumerate(sampled_files):
         data = torch.load(file)
-        expected_keys = [
-            "input_ids",
-            "hidden_states",
-            "aux_hidden_states",
-            "conversation_id",
-        ]
-        if set(expected_keys) != set(data.keys()):
-            print(f"File {file} does not contain all expected keys: {expected_keys}")
+        # Require only the keys every dump carries, and report what is missing rather
+        # than comparing the whole set: the optional keys are legitimately absent on
+        # some dumps -- ``aux_hidden_states`` under ``--no-aux-hidden-states``, and
+        # ``loss_mask`` unless ``--answer-only-loss`` was used -- and an exact-set
+        # check rejects those as malformed.
+        required_keys = ["input_ids", "hidden_states", "conversation_id"]
+        missing = [k for k in required_keys if k not in data]
+        if missing:
+            print(f"File {file} is missing required keys: {missing}")
             print(f"  Found keys: {list(data.keys())}")
             continue
         print(f"Sample {i + 1}: {file.name}")
         for key in ["input_ids", "hidden_states", "aux_hidden_states"]:
+            if key not in data:
+                print(f"{key}: absent")
+                continue
             print(f"{key}: {data[key].shape} {data[key].dtype} {data[key].device}")
         print(f"conversation_id: {data['conversation_id']}")
         input_ids_list = data["input_ids"].tolist()

examples/speculative_decoding/eagle_utils.py

@@ -61,6 +61,7 @@ def make_speculative_data_module(
     answer_only_loss=False,
     shift_labels=True,
     final_aux_is_base_hidden=False,
+    aux_hidden_states_optional=False,
 ) -> dict:
     """Create data module for speculative decoding training.
 
@@ -70,6 +71,8 @@ def make_speculative_data_module(
         final_aux_is_base_hidden: Streaming only. True when the draft's top aux layer is the
             base's final layer, so the last captured plane is both the final aux feature and
             the base (KD-target) hidden instead of an extra dedicated plane.
+        aux_hidden_states_optional: accept dumps written with ``--no-aux-hidden-states``.
+            Only modes that never read the aux planes may set this.
     """
     # Load chat template from file if provided
     chat_template = None
@@ -167,7 +170,10 @@ def make_speculative_data_module(
         if data_args.sample_size > 0:
             dumped_files = dumped_files[: data_args.sample_size]
         train_dataset = OfflineSupervisedDataset(
-            dumped_files, answer_only_loss=answer_only_loss, tokenizer=tokenizer
+            dumped_files,
+            answer_only_loss=answer_only_loss,
+            tokenizer=tokenizer,
+            aux_hidden_states_optional=aux_hidden_states_optional,
         )
         data_collator = EagleOfflineDataCollator(train_len=train_len)
 

modelopt/torch/speculative/eagle/utils.py

@@ -90,7 +90,9 @@ class OfflineSupervisedDataset(Dataset):
 
     - ``input_ids``: token IDs of shape ``(seq_len,)``
     - ``hidden_states``: base model last hidden states of shape ``(seq_len, hidden_size)``
-    - ``aux_hidden_states``: auxiliary hidden states of shape ``(seq_len, hidden_size)``
+    - ``aux_hidden_states``: auxiliary hidden states of shape ``(seq_len, hidden_size)``.
+      Absent from dumps written with ``--no-aux-hidden-states``; reading one of those
+      requires ``aux_hidden_states_optional=True``.
     - ``loss_mask`` (optional): per-token assistant mask of shape ``(seq_len,)``,
       present when the dump was produced with ``--answer-only-loss``.
 
@@ -106,23 +108,40 @@ class OfflineSupervisedDataset(Dataset):
             is stored in the file (backward compatible).
         tokenizer: Optional tokenizer used to recover the assistant mask for dumps
             that lack a stored ``loss_mask``.
+        aux_hidden_states_optional (bool): Tolerate dumps written with
+            ``--no-aux-hidden-states``, yielding an empty ``(seq_len, 0)`` plane in
+            their place. Only modes that never read the aux planes may set this; for
+            the others a missing key must still raise.
     """
 
     def __init__(
         self,
         dumped_files,
         answer_only_loss: bool = False,
         tokenizer=None,
+        aux_hidden_states_optional: bool = False,
     ):
         """Initialize with a list of .pt file paths."""
         super().__init__()
         self.dumped_files = dumped_files
         self.answer_only_loss = answer_only_loss
         self.tokenizer = tokenizer
+        self.aux_hidden_states_optional = aux_hidden_states_optional
 
     def __len__(self):
         return len(self.dumped_files)
 
+    def _aux_hidden_states(self, offline_data):
+        """Aux planes from the dump, or an empty plane when the consumer never reads them."""
+        if "aux_hidden_states" in offline_data:
+            return offline_data["aux_hidden_states"]
+        if not self.aux_hidden_states_optional:
+            raise KeyError(
+                "aux_hidden_states missing from the offline dump. Re-dump without "
+                "--no-aux-hidden-states, or use a mode that does not read them."
+            )
+        return offline_data["hidden_states"].new_zeros((offline_data["hidden_states"].shape[0], 0))
+
     def __getitem__(self, i) -> dict[str, torch.Tensor]:
         offline_data = torch.load(self.dumped_files[i], weights_only=True)
 
@@ -152,7 +171,7 @@ def __getitem__(self, i) -> dict[str, torch.Tensor]:
         ret = {
             "input_ids": offline_data["input_ids"],
             "base_model_hidden_states": offline_data["hidden_states"],
-            "aux_hidden_states": offline_data["aux_hidden_states"],
+            "aux_hidden_states": self._aux_hidden_states(offline_data),
             "attention_mask": torch.ones_like(offline_data["input_ids"]),
             "loss_mask": loss_mask,
             "labels": labels,

tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py

@@ -684,6 +684,30 @@ def test_offline_dataset_len_and_getitem(tmp_path):
     assert item["labels"].shape == (SEQ_LEN,)
 
 
+@pytest.mark.parametrize(
+    ("optional", "expect"),
+    [
+        # Modes that read the aux planes must fail loudly on a dump that omits them.
+        (False, "raise"),
+        # The external draft never reads them, so it may opt in to an empty plane.
+        (True, "empty"),
+    ],
+)
+def test_offline_dataset_handles_a_dump_without_aux_hidden_states(tmp_path, optional, expect):
+    """--no-aux-hidden-states dumps are readable only by modes that opt in."""
+    p = tmp_path / "sample.pt"
+    data = _make_offline_pt(p)
+    del data["aux_hidden_states"]
+    torch.save(data, p)
+
+    ds = OfflineSupervisedDataset([str(p)], aux_hidden_states_optional=optional)
+    if expect == "raise":
+        with pytest.raises(KeyError, match="aux_hidden_states missing"):
+            ds[0]
+    else:
+        assert ds[0]["aux_hidden_states"].shape == (SEQ_LEN, 0)
+
+
 def test_offline_dataset_labels_shift(tmp_path):
     """Labels should be input_ids shifted left by 1."""
     p = tmp_path / "sample.pt"