Create bridge for every module in Phi 1 (#1055)

Author: degenfabian

transformer_lens/model_bridge/sources/transformers.py

@@ -230,6 +230,7 @@ def boot(
     # Load the tokenizer
     tokenizer = tokenizer
     default_padding_side = getattr(adapter.cfg, "default_padding_side", None)
+    use_fast = getattr(adapter.cfg, "use_fast", True)
 
     if tokenizer is not None:
         tokenizer = setup_tokenizer(tokenizer, default_padding_side=default_padding_side)
@@ -239,6 +240,7 @@ def boot(
             AutoTokenizer.from_pretrained(
                 model_name,
                 add_bos_token=True,
+                use_fast=use_fast,
                 token=huggingface_token if len(huggingface_token) > 0 else None,
             ),
             default_padding_side=default_padding_side,

transformer_lens/model_bridge/supported_architectures/phi.py

@@ -11,6 +11,7 @@
     AttentionBridge,
     BlockBridge,
     EmbeddingBridge,
+    LinearBridge,
     MLPBridge,
     NormalizationBridge,
     UnembeddingBridge,
@@ -28,6 +29,8 @@ def __init__(self, cfg: Any) -> None:
         """
         super().__init__(cfg)
 
+        self.default_cfg = {"use_fast": False}
+
         self.conversion_rules = HookConversionSet(
             {
                 "embed.e": "transformer.wte.weight",
@@ -78,13 +81,30 @@ def __init__(self, cfg: Any) -> None:
         # Set up component mapping
         self.component_mapping = {
             "embed": EmbeddingBridge(name="model.embed_tokens"),
+            "rotary_emb": EmbeddingBridge(name="model.rotary_emb"),
             "blocks": BlockBridge(
                 name="model.layers",
                 submodules={
                     "ln1": NormalizationBridge(name="input_layernorm", config=self.cfg),
+                    "attn": AttentionBridge(
+                        name="self_attn",
+                        config=self.cfg,
+                        submodules={
+                            "q": LinearBridge(name="q_proj"),
+                            "k": LinearBridge(name="k_proj"),
+                            "v": LinearBridge(name="v_proj"),
+                            "o": LinearBridge(name="dense"),
+                        },
+                    ),
+                    # Layer norm 1 and 2 are tied.
                     "ln2": NormalizationBridge(name="input_layernorm", config=self.cfg),
-                    "attn": AttentionBridge(name="self_attn", config=self.cfg),
-                    "mlp": MLPBridge(name="mlp"),
+                    "mlp": MLPBridge(
+                        name="mlp",
+                        submodules={
+                            "in": LinearBridge(name="fc1"),
+                            "out": LinearBridge(name="fc2"),
+                        },
+                    ),
                 },
             ),
             "ln_final": NormalizationBridge(name="model.final_layernorm", config=self.cfg),