Add support for Internlm2 (#362)

arcee-ai · Jul 3, 2024 · 3eb146f · 3eb146f
1 parent 21937cd
commit 3eb146f
Showing 1 changed file with 50 additions and 0 deletions.
diff --git a/mergekit/_data/architectures/internlm2.json b/mergekit/_data/architectures/internlm2.json
@@ -0,0 +1,50 @@
+{
+    "model_type": "internlm2",
+    "architectures": [
+        "InternLM2ForCausalLM"
+    ],
+    "pre_weights": [
+        {
+            "name": "model.tok_embeddings.weight",
+            "is_embed": true
+        }
+    ],
+    "post_weights": [
+        {
+            "name": "model.norm.weight"
+        },
+        {
+            "name": "output.weight",
+            "is_embed": true,
+            "aliases": [
+                "model.tok_embeddings.weight"
+            ]
+        }
+    ],
+    "num_layers_config_key": "num_hidden_layers",
+    "layer_templates": {
+        "weights": [
+            {
+                "name": "model.layers.${layer_index}.attention_norm.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.ffn_norm.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.attention.wqkv.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.attention.wo.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.feed_forward.w1.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.feed_forward.w2.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.feed_forward.w3.weight"
+            }
+        ]
+    }
+}