Weight access with getattr in layer norm folding

Author: degenfabian

transformer_lens/model_bridge/bridge.py

@@ -531,14 +531,20 @@ def fold_layer_norm(self, fold_biases=True, center_weights=True):
             # Fold ln2 into MLP
             if not self.cfg.attn_only:
                 if fold_biases:
-                    self.blocks[l].mlp.input.bias.data = self.blocks[l].mlp.input.bias.data + (
-                        self.blocks[l].mlp.input.weight.data * self.blocks[l].ln2.bias.data[:, None]
-                    ).sum(-2)
+                    getattr(self.blocks[l].mlp, "in").bias.data = getattr(
+                        self.blocks[l].mlp, "in"
+                    ).bias.data + (
+                        getattr(self.blocks[l].mlp, "in").weight.data
+                        * self.blocks[l].ln2.bias.data[:, None]
+                    ).sum(
+                        -2
+                    )
 
                     self.blocks[l].ln2.bias.data = torch.zeros_like(self.blocks[l].ln2.bias)
 
-                self.blocks[l].mlp.input.weight.data = (
-                    self.blocks[l].mlp.input.weight.data * self.blocks[l].ln2.weight.data[:, None]
+                getattr(self.blocks[l].mlp, "in").weight.data = (
+                    getattr(self.blocks[l].mlp, "in").weight.data
+                    * self.blocks[l].ln2.weight.data[:, None]
                 )
 
                 if self.cfg.gated_mlp:
@@ -550,10 +556,10 @@ def fold_layer_norm(self, fold_biases=True, center_weights=True):
                 self.blocks[l].ln2.weight.data = torch.zeros_like(self.blocks[l].ln2.weight)
 
                 if center_weights:
-                    self.blocks[l].mlp.input.weight.data = self.blocks[
-                        l
-                    ].mlp.input.weight.data - einops.reduce(
-                        self.blocks[l].mlp.input.weight.data,
+                    getattr(self.blocks[l].mlp, "in").weight.data = getattr(
+                        self.blocks[l].mlp, "in"
+                    ).weight.data - einops.reduce(
+                        getattr(self.blocks[l].mlp, "in").weight.data,
                         "d_model d_mlp -> 1 d_mlp",
                         "mean",
                     )