Fix Theseus checkpoint loading

bert_squeeze/models/lt_theseus_bert.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import Mapping
 from typing import Optional, Union
 
 import lightning.pytorch as pl
@@ -19,6 +20,33 @@
 from .custom_transformers import TheseusBertModel
 
 
+def _should_initialize_successor_layers(
+    model: TheseusBertModel, loading_info: object
+) -> bool:
+    if not isinstance(loading_info, Mapping):
+        raise TypeError("Theseus loading information must be a mapping.")
+
+    missing_keys = loading_info.get("missing_keys")
+    if not isinstance(missing_keys, list) or not all(
+        isinstance(key, str) for key in missing_keys
+    ):
+        raise TypeError("Theseus loading information must contain missing key names.")
+
+    model_prefix = f"{model.base_model_prefix}."
+    normalized_missing_keys = {key.removeprefix(model_prefix) for key in missing_keys}
+    expected_successor_keys = {
+        f"encoder.successor_layers.{key}"
+        for key in model.encoder.successor_layers.state_dict()
+    }
+    missing_successor_keys = expected_successor_keys & normalized_missing_keys
+
+    if not missing_successor_keys:
+        return False
+    if missing_successor_keys != expected_successor_keys:
+        raise ValueError("Theseus checkpoint has incomplete successor layer weights.")
+    return True
+
+
 class LtTheseusBert(BaseSequenceClassificationTransformerModule):
     """
     Lightning module to fine-tune a TheseusBert based model on a sequence classification
@@ -50,12 +78,18 @@ def __init__(
         **kwargs,
     ):
         if model is None:
-            model = TheseusBertModel.from_pretrained(
+            loaded_model, loading_info = TheseusBertModel.from_pretrained(
                 pretrained_model,
                 config=AutoConfig.from_pretrained(
                     pretrained_model, num_labels=num_labels
                 ),
+                output_loading_info=True,
             )
+            if not isinstance(loaded_model, TheseusBertModel):
+                raise TypeError("Expected a TheseusBertModel checkpoint.")
+            model = loaded_model
+            if _should_initialize_successor_layers(model, loading_info):
+                model.encoder.init_successor_layers()
 
         super().__init__(
             training_config, pretrained_model, num_labels, model, scorer, **kwargs
@@ -116,10 +150,8 @@ def forward(
     def _before_training_step(self) -> None:
         self.replacement_scheduler.step()
 
-    def _build_model(self):
-        """"""
+    def _build_model(self) -> None:
         self.encoder = self.model
-        self.encoder.encoder.init_successor_layers()
 
         self.classifier = torch.nn.Sequential(
             torch.nn.Dropout(self.model_config.hidden_dropout_prob),

tests/test_custom_model_initialization.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 from typing import Callable
 
+import pytest
 import torch
 from omegaconf import OmegaConf
 from transformers import (
@@ -227,8 +228,11 @@ def hook(*args: object) -> None:
     assert layer_calls == [1, 0]
 
 
-def test_theseus_loads_and_uses_a_custom_pretrained_encoder(tmp_path):
+def test_theseus_restores_trained_successor_layers(tmp_path: Path) -> None:
     source_encoder = TheseusBertModel(_bert_config(num_hidden_layers=6))
+    with torch.no_grad():
+        source_encoder.encoder.layer[0].attention.self.query.weight.fill_(1.0)
+        source_encoder.encoder.successor_layers[0].attention.self.query.weight.fill_(7.0)
     source_encoder.save_pretrained(tmp_path)
 
     module = LtTheseusBert(
@@ -250,6 +254,52 @@ def test_theseus_loads_and_uses_a_custom_pretrained_encoder(tmp_path):
         module.encoder.embeddings.word_embeddings.weight,
         source_encoder.embeddings.word_embeddings.weight,
     )
+    assert torch.equal(
+        module.encoder.encoder.successor_layers[0].attention.self.query.weight,
+        source_encoder.encoder.successor_layers[0].attention.self.query.weight,
+    )
     assert logits.shape == (2, 2)
     assert torch.isfinite(loss)
     assert module.replacement_scheduler.step_counter == 1
+
+
+def test_theseus_initializes_successors_missing_from_bert_checkpoint(
+    tmp_path: Path,
+) -> None:
+    source_encoder = BertForSequenceClassification(_bert_config(num_hidden_layers=6))
+    with torch.no_grad():
+        source_encoder.bert.encoder.layer[0].attention.self.query.weight.fill_(3.0)
+    source_encoder.save_pretrained(tmp_path)
+
+    module = LtTheseusBert(
+        training_config=_training_config(),
+        pretrained_model=str(tmp_path),
+        num_labels=2,
+        replacement_scheduler=OmegaConf.create(
+            {"type": "constant", "replacing_rate": 1.0}
+        ),
+    )
+
+    assert torch.equal(
+        module.encoder.encoder.successor_layers[0].attention.self.query.weight,
+        source_encoder.bert.encoder.layer[0].attention.self.query.weight,
+    )
+
+
+def test_theseus_rejects_partial_successor_checkpoint(tmp_path: Path) -> None:
+    source_encoder = TheseusBertModel(_bert_config(num_hidden_layers=6))
+    source_encoder.save_pretrained(tmp_path, safe_serialization=False)
+    checkpoint_path = tmp_path / "pytorch_model.bin"
+    state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+    state_dict.pop("encoder.successor_layers.0.attention.self.query.weight")
+    torch.save(state_dict, checkpoint_path)
+
+    with pytest.raises(ValueError, match="incomplete successor layer weights"):
+        LtTheseusBert(
+            training_config=_training_config(),
+            pretrained_model=str(tmp_path),
+            num_labels=2,
+            replacement_scheduler=OmegaConf.create(
+                {"type": "constant", "replacing_rate": 1.0}
+            ),
+        )