From fd6c0538336b97cbade56d9c879f988fe9c88ddc Mon Sep 17 00:00:00 2001 From: sohv Date: Tue, 4 Aug 2026 15:22:47 +0000 Subject: [PATCH] fix(bridge): remove orphaned convert_weights override from nanogpt adapter NanogptArchitectureAdapter.convert_weights ended in super().convert_weights(remote_module), but ArchitectureAdapter has had no such method since 3efbd6e ("Cleanup (#1129)"), so every call raised AttributeError. The failure was invisible to CI because of the # type: ignore[misc] on that line; without it mypy reports '"convert_weights" undefined in superclass'. The override had no callers anywhere in the tree, and its _orig_mod. prefix strip was a no-op regardless: nn.Module.state_dict() returns a fresh dict, so the loop mutated a throwaway copy before passing the original module to super(). Removing it cannot regress behaviour, since every path through it raised. Dropping the ignore leaves the CI type-check job as the guard against reintroduction. Not re-homed into preprocess_weights: that hook runs on self.state_dict() inside process_weights, i.e. after load with TL-renamed keys, so it never observes _orig_mod.-prefixed checkpoint keys. --- .../supported_architectures/nanogpt.py | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/transformer_lens/model_bridge/supported_architectures/nanogpt.py b/transformer_lens/model_bridge/supported_architectures/nanogpt.py index 13d381d2c..5ba52e4bf 100644 --- a/transformer_lens/model_bridge/supported_architectures/nanogpt.py +++ b/transformer_lens/model_bridge/supported_architectures/nanogpt.py @@ -1,7 +1,5 @@ from typing import Any -import torch - from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion from transformer_lens.conversion_utils.param_processing_conversion import ( ParamProcessingConversion, @@ -88,16 +86,3 @@ def __init__(self, cfg: Any) -> None: ), # Final layer norm "unembed": UnembeddingBridge(name="lm_head"), } - - def convert_weights(self, remote_module: Any) -> dict[str, torch.Tensor]: - # Nanogpt models saved after torch.compile() have this unwanted prefix - # This is a simple way to remove it - unwanted_prefix = "_orig_mod." - state_dict: dict[str, torch.Tensor] = ( - remote_module.state_dict() if hasattr(remote_module, "state_dict") else remote_module - ) - for k, v in list(state_dict.items()): - if k.startswith(unwanted_prefix): - state_dict[k[len(unwanted_prefix) :]] = state_dict.pop(k) - - return super().convert_weights(remote_module) # type: ignore[misc]