Add first pass at GraniteMoeHybridForCausalLM support

Author: cg123

mergekit/_data/architectures/granite4moehybrid.json

@@ -0,0 +1,99 @@
+{
+    "model_type": "granitemoehybrid",
+    "architectures": [
+        "GraniteMoeHybridForCausalLM"
+    ],
+    "pre_weights": [
+        {
+            "name": "model.embed_tokens.weight",
+            "is_embed": true
+        }
+    ],
+    "post_weights": [
+        {
+            "name": "model.norm.weight"
+        },
+        {
+            "name": "lm_head.weight",
+            "is_embed": true,
+            "optional": true,
+            "tied_names": [
+                "model.embed_tokens.weight"
+            ]
+        }
+    ],
+    "num_layers_config_key": "num_hidden_layers",
+    "layer_templates": {
+        "weights": [
+            {
+                "name": "model.layers.${layer_index}.input_layernorm.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.shared_mlp.input_linear.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.shared_mlp.output_linear.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.block_sparse_moe.input_linear.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.block_sparse_moe.router.layer.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.block_sparse_moe.output_linear.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.post_attention_layernorm.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.self_attn.k_proj.weight",
+                "optional": true
+            },
+            {
+                "name": "model.layers.${layer_index}.self_attn.o_proj.weight",
+                "optional": true
+            },
+            {
+                "name": "model.layers.${layer_index}.self_attn.q_proj.weight",
+                "optional": true
+            },
+            {
+                "name": "model.layers.${layer_index}.self_attn.v_proj.weight",
+                "optional": true
+            },
+            {
+                "name": "model.layers.${layer_index}.mamba.A_log",
+                "optional": true
+            },
+            {
+                "name": "model.layers.${layer_index}.mamba.conv1d.bias",
+                "optional": true
+            },
+            {
+                "name": "model.layers.${layer_index}.mamba.conv1d.weight",
+                "optional": true
+            },
+            {
+                "name": "model.layers.${layer_index}.mamba.D",
+                "optional": true
+            },
+            {
+                "name": "model.layers.${layer_index}.mamba.dt_bias",
+                "optional": true
+            },
+            {
+                "name": "model.layers.${layer_index}.mamba.in_proj.weight",
+                "optional": true
+            },
+            {
+                "name": "model.layers.${layer_index}.mamba.norm.weight",
+                "optional": true
+            },
+            {
+                "name": "model.layers.${layer_index}.mamba.out_proj.weight",
+                "optional": true
+            }
+        ]
+    }
+}

pyproject.toml

@@ -60,6 +60,7 @@ packages = [
     "mergekit.scripts",
     "mergekit.evo",
     "mergekit.tokenizer",
+    "mergekit.tokensurgeon",
     "mergekit.architecture",
     "mergekit._data",
     "mergekit._data.architectures",