Bug new loading (#891)

* tested gpt2 loading

* fixed gemma 1 bugs

* ran format

Author: bryce13950

transformer_lens/weight_conversion/conversion_utils/conversion_steps/rearrange_weight_conversion.py

@@ -19,7 +19,7 @@ def __init__(
         self.pattern = pattern
         self.axes_lengths = axes_lengths
 
-    def handle_conversion(self, input_value: torch.Tensor) -> torch.Tensor:
+    def handle_conversion(self, input_value: torch.Tensor, *full_context) -> torch.Tensor:
         return einops.rearrange(input_value, self.pattern, **self.axes_lengths)
 
     def __repr__(self):

transformer_lens/weight_conversion/conversion_utils/conversion_steps/repeat_weight_conversion.py

@@ -19,7 +19,7 @@ def __init__(
         self.pattern = pattern
         self.axes_lengths = axes_lengths
 
-    def handle_conversion(self, input_value: torch.Tensor) -> torch.Tensor:
+    def handle_conversion(self, input_value: torch.Tensor, *full_context) -> torch.Tensor:
         return einops.repeat(input_value, self.pattern, **self.axes_lengths)
 
     def __repr__(self):

transformer_lens/weight_conversion/conversion_utils/conversion_steps/weight_conversion_set.py

@@ -45,7 +45,7 @@ def process_conversion(
         self, input_value, remote_field: str, conversion: BaseWeightConversion, *full_context
     ):
         field = find_property(remote_field, input_value)
-        if isinstance(field, WeightConversionSet):
+        if isinstance(conversion, WeightConversionSet):
             result = []
             for layer in field:
                 result.append(conversion.convert(layer, input_value, *full_context))

transformer_lens/weight_conversion/gemma.py

@@ -13,7 +13,7 @@
 
 
 class GemmaWeightNormalizationConversion(BaseWeightConversion):
-    def convert(self, input_value):
+    def convert(self, input_value, *full_context):
         return input_value.float() + torch.ones_like(input_value, dtype=torch.float32)
 
     def __repr__(self):
@@ -24,7 +24,7 @@ class GemmaWeightConversion(ArchitectureConversion):
     def __init__(self, cfg: HookedTransformerConfig) -> None:
         super().__init__(
             {
-                "unembed.W_U": "model.lm_head.weight.T",
+                "unembed.W_U": "lm_head.weight.T",
                 "unembed.b_U": torch.zeros(cfg.d_vocab),
                 "ln_final.w": (
                     "model.norm.weight",
@@ -100,4 +100,4 @@ def normalization_before_and_after_conversions(self) -> FIELD_SET:
         }
 
     def standard_normalization_conversions(self) -> FIELD_SET:
-        return {"ln2.w": ("pre_feedforward_layernorm.weight", GemmaWeightNormalizationConversion())}
+        return {"ln2.w": ("post_attention_layernorm.weight", GemmaWeightNormalizationConversion())}