Cleanup old implementations/features

scripts/train.py

@@ -177,8 +177,6 @@ def main(args: argparse.Namespace):
     )
 
     # Get model class from registry and create model using its factory method
-    if SpeculatorModel.registry_auto_discovery:
-        SpeculatorModel.auto_populate_registry()
 
     if args.speculator_type not in SpeculatorModel.registry:
         raise ValueError(

src/speculators/__init__.py

@@ -23,22 +23,23 @@
 from .config import (
     SpeculatorModelConfig,
     SpeculatorsConfig,
-    TokenProposalConfig,
     VerifierConfig,
-    reload_and_populate_configs,
+    reload_schemas,
 )
-from .model import SpeculatorModel, reload_and_populate_models
+from .model import SpeculatorModel
+from .models import Eagle3DraftModel, Eagle3SpeculatorConfig
+from .proposals import TokenProposalConfig
 
 __all__ = [
+    "Eagle3DraftModel",
+    "Eagle3SpeculatorConfig",
     "SpeculatorModel",
     "SpeculatorModelConfig",
     "SpeculatorsConfig",
     "TokenProposalConfig",
     "VerifierConfig",
-    "reload_and_populate_configs",
-    "reload_and_populate_models",
+    "reload_schemas",
 ]
 
 # base imports complete, run auto loading for base classes
-reload_and_populate_configs()
-reload_and_populate_models()
+reload_schemas()

src/speculators/config.py

@@ -25,49 +25,17 @@
 from pydantic import BaseModel, ConfigDict, Field
 from transformers import PretrainedConfig
 
+from speculators.proposals import TokenProposalConfig
 from speculators.utils import PydanticClassRegistryMixin, ReloadableBaseModel
 
 __all__ = [
     "SpeculatorModelConfig",
     "SpeculatorsConfig",
-    "TokenProposalConfig",
     "VerifierConfig",
-    "reload_and_populate_configs",
+    "reload_schemas",
 ]
 
 
-class TokenProposalConfig(PydanticClassRegistryMixin):
-    """
-    The base config for a token proposal method which defines how tokens are generated
-    by the speculator, how they are passed to the verifier, and how they are scored
-    for acceptance or rejection. All implementations of token proposal methods
-    must inherit from this class, set the proposal_type to a unique value, and
-    add any additional parameters needed to instantiate and implement the method.
-
-    It uses pydantic to validate the parameters, provide default values, and
-    enable automatic serialization and deserialization of the correct class
-    types based on the proposal_type field.
-    """
-
-    @classmethod
-    def __pydantic_schema_base_type__(cls) -> type["TokenProposalConfig"]:
-        if cls.__name__ == "TokenProposalConfig":
-            return cls
-
-        return TokenProposalConfig
-
-    auto_package: ClassVar[str] = "speculators.proposals"
-    registry_auto_discovery: ClassVar[bool] = True
-    schema_discriminator: ClassVar[str] = "proposal_type"
-
-    proposal_type: str = Field(
-        description=(
-            "The type of token proposal the config is for. "
-            "Must be a supported proposal type from the Speculators repo."
-        ),
-    )
-
-
 class VerifierConfig(BaseModel):
     """
     The base config for a verifier model which defines the parameters that are required
@@ -330,12 +298,12 @@ def to_diff_dict(self) -> dict[str, Any]:
         return super().to_diff_dict()
 
 
-def reload_and_populate_configs():
+def reload_schemas():
     """
     Automatically populates the registry for all PydanticClassRegistryMixin subclasses
     and reloads schemas for all Config classes to ensure their schemas are up-to-date
     with the current registry state.
     """
-    TokenProposalConfig.auto_populate_registry()
+    TokenProposalConfig.reload_schema()
     SpeculatorsConfig.reload_schema()
-    SpeculatorModelConfig.auto_populate_registry()
+    SpeculatorModelConfig.reload_schema()

src/speculators/convert/eagle/eagle3_legacy_model.py

@@ -242,7 +242,7 @@ class Eagle3Speculator(SpeculatorModel):
     """
 
     config_class: ClassVar[type[Eagle3SpeculatorConfig]] = Eagle3SpeculatorConfig  # type: ignore[misc]
-    _keys_to_ignore_on_load_missing: ClassVar[list[str]] = [  # type: ignore[misc]
+    _keys_to_ignore_on_load_missing: ClassVar[list[str]] = [  # type: ignore[assignment,misc]
         "verifier*",
     ]
     _keys_to_ignore_on_save: ClassVar[list[str]] = []  # type: ignore[misc,assignment]

src/speculators/convert/eagle/eagle_converter.py

@@ -9,13 +9,16 @@
 from transformers import LlamaConfig, PretrainedConfig
 
 from speculators.config import SpeculatorsConfig, VerifierConfig
+from speculators.convert.eagle.eagle_legacy_model import (
+    EagleSpeculator,
+    EagleSpeculatorConfig,
+)
 from speculators.convert.eagle.utils import (
     detect_fusion_bias_and_layernorms,
     ensure_checkpoint_is_local,
     load_checkpoint_config,
     load_checkpoint_weights,
 )
-from speculators.models.eagle import EagleSpeculator, EagleSpeculatorConfig
 from speculators.proposals.greedy import GreedyTokenProposalConfig
 
 

src/speculators/models/eagle.py → src/speculators/convert/eagle/eagle_legacy_model.py

@@ -20,7 +20,12 @@
 import torch
 from pydantic import Field, field_serializer, field_validator, model_validator
 from torch import nn
-from transformers import AutoConfig, PretrainedConfig, PreTrainedModel
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    PretrainedConfig,
+    PreTrainedModel,
+)
 from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING
@@ -234,15 +239,17 @@ class EagleSpeculator(SpeculatorModel):
 
     # PreTrainedModel settings
     config_class: ClassVar[type[EagleSpeculatorConfig]] = EagleSpeculatorConfig  # type: ignore[misc]
-    _keys_to_ignore_on_load_missing: ClassVar[list[str]] = [  # type: ignore[misc]
+    _keys_to_ignore_on_load_missing: ClassVar[list[str]] = [  # type: ignore[assignment,misc]
         "verifier*",
         "embed_tokens*",
         "lm_head*",
     ]
+
     _keys_to_ignore_on_save: ClassVar[list[str]] = [  # type: ignore[assignment,misc]
         "embed_tokens.weight",
         "lm_head.weight",
         "lm_head.bias",
+        "verifier*",
     ]
 
     @classmethod
@@ -340,13 +347,16 @@ def __init__(
         self.rotary_emb: nn.Module | None = None
         self.lm_head: nn.Linear | None = None
 
-        # Delayed initialization to ensure everything needed for attach_verifier is set
-        super().__init__(
-            config=config,
-            verifier=verifier,
-            verifier_attachment_mode=verifier_attachment_mode,
+        super().__init__(config=config)
+        self.verifier: PreTrainedModel | None = None
+        self.verifier_attachment_mode: Literal["detached", "full", "train_only"] = (
+            "detached"
         )
 
+        verifier = verifier or config.speculators_config.verifier.name_or_path
+        if verifier is not None and verifier_attachment_mode != "detached":
+            self.attach_verifier(verifier, mode=verifier_attachment_mode)
+
         self._decoder_class, self._layernorm_class = self._import_model_classes()
         # Initialize layers based on the configuration
         self.embedding_layernorm: nn.Module | None = self._create_layernorm()
@@ -360,6 +370,67 @@ def __init__(
 
         self.post_init()  # type: ignore[attr-defined]
 
+    def resolve_verifier(
+        self, verifier: str | os.PathLike | PreTrainedModel
+    ) -> PreTrainedModel:
+        """
+        Resolves the verifier model from a given path or identifier.
+
+        This method loads the verifier model from a specified path or identifier,
+        ensuring it is compatible with the speculator's configuration. If the
+        verifier is already attached, it returns the existing verifier instance.
+
+        :param verifier: The verifier model to resolve. Can be a path to a local
+            model directory, a Hugging Face model identifier, or an instance of
+            PreTrainedModel.
+        :return: The resolved PreTrainedModel instance for the verifier.
+        """
+        if not verifier:
+            raise ValueError(
+                "Verifier must be provided as a path, identifier, or PreTrainedModel. "
+            )
+
+        if not isinstance(verifier, (str, os.PathLike, PreTrainedModel)):
+            raise TypeError(
+                f"Expected verifier to be a PreTrainedModel, a string path, "
+                f"or an os.PathLike object, got {type(verifier)} {verifier}."
+            )
+
+        if isinstance(verifier, PreTrainedModel):
+            return verifier
+
+        return AutoModelForCausalLM.from_pretrained(verifier)
+
+    def state_dict(
+        self,
+        *,
+        destination: dict[str, Any] = None,  # type: ignore[assignment]
+        prefix: str = "",
+        keep_vars: bool = False,
+    ):
+        """
+        Overrides the state_dict method from PyTorch to ensure that save pathways
+        within Transformers PreTrainedModel do not include the verifier model's
+        parameters. This is important to ensure that the speculator model
+        can be saved and loaded without including the verifier's state, which
+        is expected to be managed separately.
+
+        :param destination: Optional dictionary to store the state.
+        :param prefix: Optional prefix for parameter names.
+        :param keep_vars: Whether to keep Variables in the state_dict.
+        :return: A dictionary containing the state of the speculator model,
+            excluding the verifier model's parameters. This dictionary can be used
+            to save the model's state to disk or for further processing.
+        """
+        tmp_verifier = self.verifier
+        self.verifier = None
+        state = super().state_dict(  # type: ignore[misc]
+            destination=destination, prefix=prefix, keep_vars=keep_vars
+        )
+        self.verifier = tmp_verifier
+
+        return state
+
     def attach_verifier(
         self,
         verifier: str | os.PathLike | PreTrainedModel,
@@ -405,7 +476,24 @@ def attach_verifier(
             perform generation until a full verifier is attached.
         :return: The PreTrainedModel instance for the verifier that was attached.
         """
-        super().attach_verifier(verifier=verifier, mode=mode)
+        if self.verifier_attachment_mode != "detached":
+            raise RuntimeError(
+                "Cannot attach a verifier when the speculator is not in detached mode. "
+                "Detach the current verifier first using `detach_verifier()`."
+            )
+
+        if mode not in {"full", "train_only", None}:
+            raise ValueError(
+                f"Invalid verifier_attachment_mode: {mode}. "
+                "Must be one of 'full', 'train_only', or None."
+            )
+
+        self.verifier_attachment_mode = mode or "full"
+        self.verifier = (
+            self.resolve_verifier(verifier)
+            if self.verifier_attachment_mode == "full"
+            else None
+        )  # Expect subclasses to handle references if train_only
 
         if self.verifier_attachment_mode == "train_only":
             verifier_model = self.resolve_verifier(verifier)
@@ -432,7 +520,17 @@ def detach_verifier(self):
         be able to perform forward passes or generation until a new verifier
         is attached.
         """
-        super().detach_verifier()
+        if self.verifier_attachment_mode == "detached":
+            raise RuntimeError(
+                "Verifier is already detached, cannot be called again until "
+                "a new verifier is attached."
+            )
+
+        if self.verifier is not None:
+            del self.verifier
+
+        self.verifier = None
+        self.verifier_attachment_mode = "detached"
 
         del self.embed_tokens
         self.embed_tokens = None