[2/8] Add the external draft speculative-decoding mode

examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py

@@ -152,10 +152,14 @@ 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}")
+        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)
@@ -189,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,

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()

modelopt/torch/speculative/__init__.py

@@ -15,6 +15,6 @@
 
 """Speculative Decoding Optimizations."""
 
-from . import eagle, medusa, mode, plugins
+from . import eagle, external, medusa, mode, plugins
 from .config import *
 from .speculative_decoding import *

modelopt/torch/speculative/config.py

@@ -30,6 +30,7 @@
     "EAGLE_MTP_DEFAULT_CFG",
     "DFlashConfig",
     "EagleConfig",
+    "ExternalDraftConfig",
     "MedusaConfig",
     "eagle3_default_config",
     "eagle_mtp_default_config",
@@ -468,6 +469,58 @@ class MedusaConfig(ModeloptBaseConfig):
     )
 
 
+class ExternalDraftConfig(ModeloptBaseConfig):
+    """External (standalone) draft model config.
+
+    The converted module is a pretrained causal LM used directly as the draft. It
+    keeps its own embeddings and lm_head; the base model contributes only cached
+    hidden states, projected by the base lm_head to obtain teacher logits.
+    """
+
+    external_offline: bool = ModeloptField(
+        default=True,
+        description=(
+            "Whether the draft consumes pre-computed base hidden states instead of running "
+            "the base model in-process. Only the offline path is implemented."
+        ),
+    )
+
+    external_loss: Literal["soft_ce", "tvd", "tvd_deploy"] = ModeloptField(
+        default="tvd",
+        description=(
+            "Training objective. ``soft_ce`` matches the Eagle offline objective. "
+            "``tvd`` is total-variation distance between the two untruncated "
+            "distributions. ``tvd_deploy`` puts both through the serving filter first "
+            "(``external_top_k`` then ``external_top_p``), making it the acceptance the "
+            "deployed pair will see: ``alpha = 1 - TVD``."
+        ),
+    )
+
+    external_top_k: int = ModeloptField(
+        default=20,
+        ge=0,
+        description=(
+            "Top-k of the serving filter, applied to *both* the base and the "
+            "draft by ``external_loss='tvd_deploy'``. 0 disables the top-k stage."
+        ),
+    )
+
+    external_top_p: float = ModeloptField(
+        default=0.95,
+        ge=0.0,
+        le=1.0,
+        description=(
+            "Top-p of the serving filter, applied to *both* the base and the "
+            "draft by ``external_loss='tvd_deploy'``. 0 or 1.0 disables the nucleus stage."
+        ),
+    )
+
+    external_report_acc: bool = ModeloptField(
+        default=True,
+        description=("Whether to report top-1 agreement with the base model during training."),
+    )
+
+
 class EagleConfig(ModeloptBaseConfig):
     """Eagle config."""
 

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

modelopt/torch/speculative/external/__init__.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""External Draft Optimization Method."""
+
+from .conversion import *
+from .external_model import *

modelopt/torch/speculative/external/conversion.py

@@ -0,0 +1,72 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""External draft conversion/restore utilities."""
+
+from torch import nn
+
+from modelopt.torch.opt.conversion import ModelLikeModule
+from modelopt.torch.opt.dynamic import _DMRegistryCls
+from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict
+
+from ..config import ExternalDraftConfig
+
+ExternalDraftDMRegistry = _DMRegistryCls(prefix="ExternalDraft")  # global instance for the registry
+
+__all__ = [
+    "ExternalDraftDMRegistry",
+    "convert_to_external_draft_model",
+    "restore_external_draft_model",
+]
+
+
+def convert_to_external_draft_model(
+    model: nn.Module, config: ExternalDraftConfig
+) -> ConvertReturnType:
+    """Convert a pretrained causal LM into a trainable external draft as per `config`.
+
+    Note that, unlike the other speculative modes, ``model`` here is the *draft*
+    rather than the base model. Nothing is grafted onto a target; the draft keeps
+    its own embeddings and lm_head.
+    """
+    # initialize the true module if necessary
+    model = model.init_modellike() if isinstance(model, ModelLikeModule) else model
+
+    original_cls = type(model)
+    if original_cls not in ExternalDraftDMRegistry:
+        for cls in ExternalDraftDMRegistry._registry:
+            if issubclass(original_cls, cls):
+                ExternalDraftDMRegistry.register({original_cls: "draft_model_class"})(
+                    ExternalDraftDMRegistry[cls]
+                )
+                break
+
+    external_draft_model = ExternalDraftDMRegistry.convert(model)
+    external_draft_model.modify(config)
+
+    # no metadata, all specified via config.
+    metadata = {}
+
+    return external_draft_model, metadata
+
+
+def restore_external_draft_model(
+    model: nn.Module, config: ExternalDraftConfig, metadata: MetadataDict
+) -> nn.Module:
+    """Function for restoring a previously converted model to an external draft model."""
+    # the metadata should be empty
+    assert not metadata, "No metadata expected!"
+
+    return convert_to_external_draft_model(model, config)[0]

modelopt/torch/speculative/external/external_model.py

@@ -0,0 +1,45 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""External draft model to support speculative decoding with a standalone draft."""
+
+from modelopt.torch.opt.dynamic import DynamicModule
+
+__all__ = ["ExternalDraftModel"]
+
+
+class ExternalDraftModel(DynamicModule):
+    """Base external draft model.
+
+    Unlike Eagle/Medusa/DFlash, the converted module *is* the draft: a pretrained
+    causal LM that owns its embeddings and lm_head. The base (target) model is
+    never grafted onto it; it contributes only cached hidden states, which are
+    projected by the target lm_head to obtain teacher logits.
+    """
+
+    def _setup(self):
+        self._register_temp_attribute("external_offline", True)
+        self._register_temp_attribute("external_loss", "tvd")
+        self._register_temp_attribute("external_report_acc", True)
+        self._register_temp_attribute("external_top_k", 20)
+        self._register_temp_attribute("external_top_p", 0.95)
+
+    def modify(self, config):
+        """Base external draft modify function. Child class should implement the details."""
+        self.external_offline = config.external_offline
+        self.external_loss = config.external_loss
+        self.external_report_acc = config.external_report_acc
+        self.external_top_k = config.external_top_k
+        self.external_top_p = config.external_top_p

modelopt/torch/speculative/mode.py

@@ -23,9 +23,10 @@
     _ModeRegistryCls,
 )
 
-from .config import DFlashConfig, EagleConfig, MedusaConfig
+from .config import DFlashConfig, EagleConfig, ExternalDraftConfig, MedusaConfig
 from .dflash.conversion import convert_to_dflash_model, restore_dflash_model
 from .eagle.conversion import convert_to_eagle_model, restore_eagle_model
+from .external.conversion import convert_to_external_draft_model, restore_external_draft_model
 from .medusa.conversion import convert_to_medusa_model, restore_medusa_model
 
 SpeculativeDecodingModeRegistry = _ModeRegistryCls("speculative")
@@ -113,3 +114,31 @@ def convert(self) -> ConvertEntrypoint:
     def restore(self) -> RestoreEntrypoint:
         """The mode's entrypoint for restoring a model."""
         return restore_eagle_model
+
+
+@SpeculativeDecodingModeRegistry.register_mode
+class ExternalDraftModeDescriptor(ModeDescriptor):
+    """Class to describe the ``"external"`` mode.
+
+    The properties of this mode can be inspected via the source code.
+    """
+
+    @property
+    def name(self) -> str:
+        """Returns the value (str representation) of the mode."""
+        return "external"
+
+    @property
+    def config_class(self) -> type[ModeloptBaseConfig]:
+        """Specifies the config class for the mode."""
+        return ExternalDraftConfig
+
+    @property
+    def convert(self) -> ConvertEntrypoint:
+        """The mode's entrypoint for converting a model."""
+        return convert_to_external_draft_model
+
+    @property
+    def restore(self) -> RestoreEntrypoint:
+        """The mode's entrypoint for restoring a model."""
+        return restore_external_draft_model

modelopt/torch/speculative/plugins/__init__.py

@@ -35,5 +35,6 @@
     from .hf_domino import *
     from .hf_dspark import *
     from .hf_eagle import *
+    from .hf_external import *
     from .hf_lilicorr import *
     from .hf_medusa import *