From 2572119114266f87c0bd46c66348e7760a96cc04 Mon Sep 17 00:00:00 2001 From: Peter Hicks Date: Thu, 1 Oct 2026 09:26:38 +0000 Subject: [PATCH 1/4] [1/8] Let the hidden-state dump skip aux_hidden_states Add --no-aux-hidden-states to the HF dump script. Only the head-based modes read aux_hidden_states, and they dominate the dump: one plane per selected layer against a single plane for the base hidden state. OfflineSupervisedDataset now falls back to an empty plane when the key is absent, so dumps written with the flag still load. Also resolve num_hidden_layers through config.text_config when present; multimodal checkpoints keep the language model's depth there and the dump previously failed on them. Signed-off-by: Peter Hicks Co-Authored-By: Claude Opus 5 Signed-off-by: Peter Hicks --- .../collect_hidden_states/common.py | 13 ++++++++++ .../compute_hidden_states_hf.py | 13 ++++++++-- examples/speculative_decoding/eagle_utils.py | 8 ++++++- modelopt/torch/speculative/eagle/utils.py | 23 ++++++++++++++++-- .../plugins/test_hf_speculative_offline.py | 24 +++++++++++++++++++ 5 files changed, 76 insertions(+), 5 deletions(-) diff --git a/examples/speculative_decoding/collect_hidden_states/common.py b/examples/speculative_decoding/collect_hidden_states/common.py index 78b317853b6..742c2d54661 100644 --- a/examples/speculative_decoding/collect_hidden_states/common.py +++ b/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( diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py index 6b94e1cd6a6..7d23ee4bcc1 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py +++ b/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,7 +150,10 @@ 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 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}") selected_layer_ids = resolve_aux_layers(args, num_hidden_layers) @@ -199,7 +204,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, }, diff --git a/examples/speculative_decoding/eagle_utils.py b/examples/speculative_decoding/eagle_utils.py index 45df2c7e341..f05c092c0fe 100644 --- a/examples/speculative_decoding/eagle_utils.py +++ b/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) diff --git a/modelopt/torch/speculative/eagle/utils.py b/modelopt/torch/speculative/eagle/utils.py index 24b7dfb8b34..c00b691c083 100644 --- a/modelopt/torch/speculative/eagle/utils.py +++ b/modelopt/torch/speculative/eagle/utils.py @@ -113,16 +113,35 @@ def __init__( dumped_files, answer_only_loss: bool = False, tokenizer=None, + aux_hidden_states_optional: bool = False, ): - """Initialize with a list of .pt file paths.""" + """Initialize with a list of .pt file paths. + + Args: + aux_hidden_states_optional: tolerate dumps written with + ``--no-aux-hidden-states``. Only modes that never read the aux planes + may set this; for the others a missing key must still raise. + """ 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, diff --git a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py b/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py index 16408e2c3e8..d124a2a4de1 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py +++ b/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" From 35247e34b57cece91ebfa0047dadefa03afc81d2 Mon Sep 17 00:00:00 2001 From: Peter Hicks <60696413+phicksy@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:37:48 +0100 Subject: [PATCH 2/4] Update examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Signed-off-by: Peter Hicks <60696413+phicksy@users.noreply.github.com> --- .../collect_hidden_states/compute_hidden_states_hf.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py index 7d23ee4bcc1..047dfb28863 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py @@ -152,7 +152,11 @@ def keep_conversation(entry): ) # Multimodal checkpoints (e.g. *ForConditionalGeneration) keep the language # model's depth under config.text_config. - text_config = getattr(model.config, "text_config", None) or model.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}") From b2faa5274f8cc9eb3a4dd16c7f72af0c4d7a4ebb Mon Sep 17 00:00:00 2001 From: Peter Hicks <60696413+phicksy@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:38:03 +0100 Subject: [PATCH 3/4] Update examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Signed-off-by: Peter Hicks <60696413+phicksy@users.noreply.github.com> --- .../collect_hidden_states/compute_hidden_states_hf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py index 047dfb28863..3b51994b52d 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py @@ -159,7 +159,7 @@ def keep_conversation(entry): ) 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) From c00021c53cbbe54966d8a868bb3acdf5e7cfb087 Mon Sep 17 00:00:00 2001 From: Peter Hicks Date: Tue, 13 Oct 2026 12:40:17 +0000 Subject: [PATCH 4/4] Address review suggestions on the aux-hidden-states dump path - Skip building the aux planes under --no-aux-hidden-states: the concatenation and device-to-host copy ran per conversation on a dump about to drop them. - Document aux_hidden_states_optional beside its sibling arguments, and record in the dump-format key list that the key is absent under --no-aux-hidden-states. That list is the contract a new producer or consumer is written against. - Let sample_hidden_states inspect dumps that omit an optional key. It compared the whole key set, so it reported every --no-aux-hidden-states dump as damaged, and every --answer-only-loss dump along with it. Co-Authored-By: Claude Opus 5 Signed-off-by: Peter Hicks --- .../compute_hidden_states_hf.py | 12 +++++++++-- .../sample_hidden_states.py | 20 +++++++++++-------- modelopt/torch/speculative/eagle/utils.py | 16 +++++++-------- 3 files changed, 30 insertions(+), 18 deletions(-) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py index 3b51994b52d..21e1b652527 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py @@ -193,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 .pt as finished work, diff --git a/examples/speculative_decoding/collect_hidden_states/sample_hidden_states.py b/examples/speculative_decoding/collect_hidden_states/sample_hidden_states.py index 75a88969d38..618ef1f6597 100644 --- a/examples/speculative_decoding/collect_hidden_states/sample_hidden_states.py +++ b/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() diff --git a/modelopt/torch/speculative/eagle/utils.py b/modelopt/torch/speculative/eagle/utils.py index c00b691c083..4528bbfe5b5 100644 --- a/modelopt/torch/speculative/eagle/utils.py +++ b/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,6 +108,10 @@ 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__( @@ -115,13 +121,7 @@ def __init__( tokenizer=None, aux_hidden_states_optional: bool = False, ): - """Initialize with a list of .pt file paths. - - Args: - aux_hidden_states_optional: tolerate dumps written with - ``--no-aux-hidden-states``. Only modes that never read the aux planes - may set this; for the others a missing key must still raise. - """ + """Initialize with a list of .pt file paths.""" super().__init__() self.dumped_files = dumped_files self.answer_only_loss = answer_only_loss