Added gemma-2 2b (#687)

TransformerLensOrg · Jul 31, 2024 · 332d8d7 · 332d8d7
1 parent bf64ede
commit 332d8d7
Showing 1 changed file with 32 additions and 0 deletions.
diff --git a/transformer_lens/loading_from_pretrained.py b/transformer_lens/loading_from_pretrained.py
@@ -206,6 +206,8 @@
     "google/gemma-7b",
     "google/gemma-2b-it",
     "google/gemma-7b-it",
+    "google/gemma-2-2b",
+    "google/gemma-2-2b-it",
     "google/gemma-2-9b",
     "google/gemma-2-9b-it",
     "google/gemma-2-27b",
@@ -628,8 +630,10 @@
     "google/gemma-7b": ["gemma-7b"],
     "google/gemma-2b-it": ["gemma-2b-it"],
     "google/gemma-7b-it": ["gemma-7b-it"],
+    "google/gemma-2-2b": ["gemma-2-2b"],
     "google/gemma-2-9b": ["gemma-2-9b"],
     "google/gemma-2-27b": ["gemma-2-27b"],
+    "google/gemma-2-2b-it": ["gemma-2-2b-it"],
     "google/gemma-2-9b-it": ["gemma-2-9b-it"],
     "google/gemma-2-27b-it": ["gemma-2-27b-it"],
     "01-ai/Yi-6B": ["yi-6b", "Yi-6B"],
@@ -1218,6 +1222,34 @@ def convert_hf_model_config(model_name: str, **kwargs):
             "gated_mlp": True,
             "final_rms": True,
         }
+    elif official_model_name.startswith("google/gemma-2-2b"):
+        # Architecture for Gemma-2 2b and Gemma-2 2b Instruct models
+        cfg_dict = {
+            "d_model": 2304,
+            "d_head": 256,
+            "n_heads": 8,
+            "d_mlp": 9216,
+            "n_layers": 26,
+            "n_ctx": 8192,
+            "eps": 1e-06,
+            "d_vocab": 256000,
+            "act_fn": "gelu_pytorch_tanh",
+            "initializer_range": 0.02,
+            "normalization_type": "RMS",
+            "rotary_base": 10000.0,
+            "positional_embedding_type": "rotary",
+            "use_attn_scale": True,
+            "attn_scale": math.sqrt(224),
+            "n_key_value_heads": 4,
+            "window_size": 4096,
+            "use_local_attn": True,
+            "attn_types": ["global", "local"] * 21,  # Alternate global and local attn
+            "attn_scores_soft_cap": 50.0,
+            "output_logits_soft_cap": 30.0,
+            "gated_mlp": True,
+            "final_rms": True,
+            "use_normalization_before_and_after": True,
+        }
     elif official_model_name.startswith("google/gemma-2-9b"):
         # Architecture for Gemma-2 9b and Gemma-2 9b Instruct models
         cfg_dict = {