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

modelopt/torch/opt/dynamic.py

@@ -1057,6 +1057,18 @@ def decorator(dm_class: type[_DMRegistryCls.T]) -> type[_DMRegistryCls.T]:
 
         return decorator
 
+    def register_subclass(self, nn_cls: type[nn.Module], key: str) -> None:
+        """Register ``nn_cls`` like its registered base class, if it has one and is not covered yet.
+
+        Needed for subclasses that override ``forward``, which the automatic lookup skips.
+        """
+        if nn_cls in self:
+            return
+        for base_cls in self._registry:
+            if issubclass(nn_cls, base_cls):
+                self.register({nn_cls: key})(self[base_cls])
+                return
+
     def unregister(self, nn_cls: type[nn.Module] | type[T]) -> None:
         """Unregister a previously registered dynamic base module and all its inherited modules.
 

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,50 @@ 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_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/dflash/conversion.py

@@ -77,12 +77,7 @@ def convert_to_dflash_model(model: nn.Module, config: DFlashConfig) -> ConvertRe
             "Expected 'dflash' (default), 'domino', 'dspark', 'dflash2' or 'lilicorr'."
         )
 
-    original_cls = type(model)
-    if original_cls not in registry:
-        for cls in registry._registry:
-            if issubclass(original_cls, cls):
-                registry.register({original_cls: "base_model_class"})(registry[cls])
-                break
+    registry.register_subclass(type(model), "base_model_class")
 
     dflash_model = registry.convert(model)
     dflash_model.modify(config)

modelopt/torch/speculative/eagle/conversion.py

@@ -32,12 +32,7 @@ def convert_to_eagle_model(model: nn.Module, config: EagleConfig) -> ConvertRetu
     # 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 EagleDMRegistry:
-        for cls in EagleDMRegistry._registry:
-            if issubclass(original_cls, cls):
-                EagleDMRegistry.register({original_cls: "base_model_class"})(EagleDMRegistry[cls])
-                break
+    EagleDMRegistry.register_subclass(type(model), "base_model_class")
 
     # merge custom config with default config
     default_arch_config = {

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,65 @@
+# 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
+
+    ExternalDraftDMRegistry.register_subclass(type(model), "draft_model_class")
+
+    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,43 @@
+# 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_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_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/medusa/conversion.py

@@ -31,12 +31,7 @@ def convert_to_medusa_model(model: nn.Module, config: MedusaConfig) -> ConvertRe
     # 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 MedusaDMRegistry:
-        for cls in MedusaDMRegistry._registry:
-            if issubclass(original_cls, cls):
-                MedusaDMRegistry.register({original_cls: "base_model_class"})(MedusaDMRegistry[cls])
-                break
+    MedusaDMRegistry.register_subclass(type(model), "base_model_class")
 
     medusa_model = MedusaDMRegistry.convert(model)
     medusa_model.modify(

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 *