Prevent unnecessary Verifier Model load

Signed-off-by: Fynn Schmitt-Ulms <[email protected]>

Author: fynnsu

src/speculators/model.py

@@ -375,7 +375,7 @@ def attach_verifier(
         self,
         verifier: str | os.PathLike | PreTrainedModel,
         mode: Literal["full", "train_only"] | None = None,
-    ) -> PreTrainedModel:
+    ):
         """
         Attach a verifier model for the speculator that is used to attach to
         for running inference/training with the speculator and validates the
@@ -417,14 +417,13 @@ def attach_verifier(
                 "Must be one of 'full', 'train_only', or None."
             )
 
-        verifier = self.resolve_verifier(verifier)
         self.verifier_attachment_mode = mode or "full"
         self.verifier = (
-            verifier if self.verifier_attachment_mode == "full" else None
+            self.resolve_verifier(verifier)
+            if self.verifier_attachment_mode == "full"
+            else None
         )  # Expect subclasses to handle references if train_only
 
-        return verifier
-
     def detach_verifier(self):
         """
         Removes the reference to the attached verifier model and frees up the

src/speculators/models/eagle.py

@@ -15,7 +15,7 @@
 import os
 import re
 import warnings
-from typing import Any, ClassVar, Literal
+from typing import Any, ClassVar, Literal, cast
 
 import torch
 from pydantic import Field, field_serializer, field_validator, model_validator
@@ -308,7 +308,7 @@ def attach_verifier(
         self,
         verifier: str | os.PathLike | PreTrainedModel,
         mode: Literal["full", "train_only"] | None = None,
-    ) -> PreTrainedModel:
+    ):
         """
         Attach a verifier model to the EagleSpeculator for speculative decoding.
         Utilizes the verifier's embed_tokens, rotary_emb, and lm_head layers
@@ -349,25 +349,25 @@ def attach_verifier(
             perform generation until a full verifier is attached.
         :return: The PreTrainedModel instance for the verifier that was attached.
         """
-        verifier = super().attach_verifier(
-            verifier=verifier,
-            mode=mode,
-        )
+        super().attach_verifier(verifier=verifier, mode=mode)
 
-        # Extract layers from the verifier model
+        if self.verifier_attachment_mode == "train_only":
+            verifier_model = self.resolve_verifier(verifier)
+        elif self.verifier_attachment_mode == "full":
+            verifier_model = cast("PreTrainedModel", self.verifier)
+        else:
+            return
 
-        if hasattr(verifier, "model"):
-            self.embed_tokens = verifier.model.embed_tokens  # type: ignore[assignment,union-attr]
-            self.rotary_emb = verifier.model.rotary_emb  # type: ignore[assignment,union-attr]
+        if hasattr(verifier_model, "model"):
+            self.embed_tokens = verifier_model.model.embed_tokens  # type: ignore[assignment,union-attr]
+            self.rotary_emb = verifier_model.model.rotary_emb  # type: ignore[assignment,union-attr]
         else:
             # Bare model structure
-            self.embed_tokens = verifier.embed_tokens  # type: ignore[assignment,attr-defined]
-            self.rotary_emb = verifier.rotary_emb  # type: ignore[assignment,attr-defined]
+            self.embed_tokens = verifier_model.embed_tokens  # type: ignore[assignment,attr-defined]
+            self.rotary_emb = verifier_model.rotary_emb  # type: ignore[assignment,attr-defined]
 
         # lm_head is always at the top level of the verifier
-        self.lm_head = verifier.lm_head  # type: ignore[assignment,attr-defined]
-
-        return verifier
+        self.lm_head = verifier_model.lm_head  # type: ignore[assignment,attr-defined]
 
     def detach_verifier(self):
         """

src/speculators/train/checkpointer.py

@@ -84,7 +84,8 @@ def load_model_state_dict(self, model: PreTrainedModel):
         full_state_dict = load_safetensors_state_dict(
             self.model_path(self.previous_epoch), "cuda:0"
         )
-        model.load_state_dict(full_state_dict)
+        # Note: `strict=False` because we don't load the verifier weights
+        model.load_state_dict(full_state_dict, strict=False)
 
     def load_optimizer_state_dict(
         self,
@@ -110,10 +111,13 @@ def load_model_state_dict(self, model: PreTrainedModel):
         full_state_dict = load_safetensors_state_dict(
             self.model_path(self.previous_epoch), "cpu"
         )
+        # Note: `strict=False` because we don't load the verifier weights
         set_model_state_dict(
             model,
             full_state_dict,  # type: ignore[arg-type]
-            options=StateDictOptions(full_state_dict=True, broadcast_from_rank0=True),
+            options=StateDictOptions(
+                full_state_dict=True, broadcast_from_rank0=True, strict=False
+            ),
         )
         dist.barrier()