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..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 @@ -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 .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, }, 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/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..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__( @@ -113,16 +119,29 @@ 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.""" 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"