Create bridge for every module in Qwen 2 (#1061)

Author: degenfabian

transformer_lens/model_bridge/generalized_components/base.py

@@ -288,4 +288,7 @@ def has_bias(self) -> bool:
             raise RuntimeError(
                 f"Original component not set for {self.name}. Call set_original_component() first."
             )
+
+        if not hasattr(self.original_component, "bias"):
+            return False
         return self.original_component.bias is not None

transformer_lens/model_bridge/generalized_components/normalization.py

@@ -74,11 +74,11 @@ def forward(
         hidden_states = self.hook_normalized(hidden_states / scale)
 
         if not self.config.layer_norm_folding:
-            if self.config.uses_rms_norm:
-                # No bias if using RMSNorm
+            if self.config.uses_rms_norm or not self.has_bias():
+                # No bias if using RMSNorm or if the original component has no bias
                 hidden_states = hidden_states * self.weight
             else:
-                # Add bias if using LayerNorm
+                # Add bias if using LayerNorm and the original component has a bias
                 hidden_states = hidden_states * self.weight + self.bias
 
         output = self.hook_out(hidden_states)

transformer_lens/model_bridge/supported_architectures/qwen2.py

@@ -11,6 +11,7 @@
     AttentionBridge,
     BlockBridge,
     EmbeddingBridge,
+    LinearBridge,
     MLPBridge,
     NormalizationBridge,
     UnembeddingBridge,
@@ -24,6 +25,10 @@ def __init__(self, cfg: Any) -> None:
         """Initialize the Qwen2 architecture adapter."""
         super().__init__(cfg)
 
+        self.cfg.default_prepend_bos = False
+        self.cfg.gated_mlp = True
+        self.cfg.uses_rms_norm = True
+
         self.conversion_rules = HookConversionSet(
             {
                 "embed.e": "model.embed_tokens.weight",
@@ -65,8 +70,24 @@ def __init__(self, cfg: Any) -> None:
                 submodules={
                     "ln1": NormalizationBridge(name="input_layernorm", config=self.cfg),
                     "ln2": NormalizationBridge(name="post_attention_layernorm", config=self.cfg),
-                    "attn": AttentionBridge(name="self_attn", config=self.cfg),
-                    "mlp": MLPBridge(name="mlp"),
+                    "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="o_proj"),
+                        },
+                    ),
+                    "mlp": MLPBridge(
+                        name="mlp",
+                        submodules={
+                            "gate": LinearBridge(name="gate_proj"),
+                            "in": LinearBridge(name="up_proj"),
+                            "out": LinearBridge(name="down_proj"),
+                        },
+                    ),
                 },
             ),
             "ln_final": NormalizationBridge(name="model.norm", config=self.cfg),