add qwen3 moe

mergekit/_data/architectures/kormo.json

@@ -0,0 +1,57 @@
+{
+    "model_type": "kormo",
+    "architectures": [
+        "KORMoForCausalLM"
+    ],
+    "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}.pre_attention_layernorm.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.self_attn.q_proj.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.self_attn.k_proj.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.self_attn.v_proj.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.self_attn.o_proj.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.pre_mlp_layernorm.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.mlp.gate_proj.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.mlp.up_proj.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.mlp.down_proj.weight"
+            }
+        ]
+    }
+}

mergekit/_data/architectures/kormo_moe.json

@@ -0,0 +1,48 @@
+{
+    "model_type": "kormo_moe",
+    "architectures": [
+        "KORMoMoeForCausalLM"
+    ],
+    "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}.pre_attention_layernorm.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.self_attn.q_proj.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.self_attn.k_proj.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.self_attn.v_proj.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.self_attn.o_proj.weight"
+            },
+            {
+                "name": "model.layers.${layer_index}.pre_mlp_layernorm.weight"
+            }
+        ]
+    }
+}

mergekit/architecture/__init__.py

@@ -20,6 +20,7 @@
 from mergekit.architecture.moe_defs import (
     MixtralModuleArchitecture,
     Qwen3MoeModuleArchitecture,
+    KORMoMoeModuleArchitecture
 )
 from mergekit.options import MergeOptions
 
@@ -34,6 +35,7 @@
 def arch_info_for_config(config: PretrainedConfig) -> Optional[ModelArchitecture]:
     if len(config.architectures) != 1:
         raise RuntimeError("More than one architecture in config?")
+
     arch_name = config.architectures[0]
 
     if arch_name == MixtralModuleArchitecture.ARCHITECTURE_NAME:
@@ -50,6 +52,13 @@ def arch_info_for_config(config: PretrainedConfig) -> Optional[ModelArchitecture
             architectures=[arch_name],
             model_type="qwen3_moe",
         )
+    elif arch_name == KORMoMoeModuleArchitecture.ARCHITECTURE_NAME:  # 추가
+        module = KORMoMoeModuleArchitecture.from_config(config)
+        return ModelArchitecture(
+            modules={"default": ModuleDefinition(architecture=module)},
+            architectures=[arch_name],
+            model_type="kormo_moe",
+        )
     elif arch_name in NAME_TO_ARCH:
         candidates = list(NAME_TO_ARCH[arch_name])
         if len(candidates) == 1:

mergekit/architecture/moe_defs.py

@@ -85,19 +85,96 @@ def num_layers_config_key(self) -> str:
     def layer_weights(
         self, index: int, config: PretrainedConfig
     ) -> Optional[List[WeightInfo]]:
+        num_experts = self.num_experts
         prefix = f"model.layers.{index}"
         tensor_names = []
-        for expert_idx in range(self.num_experts):
+
+        # Expert weights 추가
+        for expert_idx in range(num_experts):
             for param in ("up_proj", "gate_proj", "down_proj"):
                 tensor_names.append(
                     prefix + f".mlp.experts.{expert_idx}.{param}.weight"
                 )
+
+        # Shared expert weights 추가 - 이 부분이 중요!
+        for param in ("up_proj", "gate_proj", "down_proj"):
+            tensor_names.append(
+                prefix + f".mlp.shared_expert.{param}.weight"
+            )
+
+        # Gate weights 추가
         tensor_names.append(prefix + ".mlp.gate.weight")
+        tensor_names.append(prefix + ".mlp.shared_expert_gate.weight")
+
         res = []
         for name in tensor_names:
             res.append(WeightInfo(name=name))
+
+        # 기존 Qwen3 weights 중에서 MLP를 제외한 것들 추가
         for weight_info in QWEN3_MODULE_ARCH.layer_weights(index, config):
             if ".mlp." in weight_info.name:
                 continue
             res.append(weight_info)
+
         return res
+
+# 파일 상단 import 부분에 추가
+KORMO_INFO = NAME_TO_ARCH["KORMoForCausalLM"][0]
+KORMO_MODULE_ARCH = KORMO_INFO.modules["default"].architecture
+
+
+class KORMoMoeModuleArchitecture(ModuleArchitecture, BaseModel):
+    ARCHITECTURE_NAME: ClassVar[str] = "KORMoMoeForCausalLM"
+    num_experts: int
+
+    def name(self) -> str:
+        return "kormo_moe"
+
+    @classmethod
+    def from_config(cls, config: PretrainedConfig):
+        return KORMoMoeModuleArchitecture(num_experts=config.num_experts)
+
+    def pre_weights(self, config: PretrainedConfig) -> List[WeightInfo]:
+        return KORMO_MODULE_ARCH.pre_weights(config)
+
+    def post_weights(self, config: PretrainedConfig) -> List[WeightInfo]:
+        return KORMO_MODULE_ARCH.post_weights(config)
+
+    def num_layers_config_key(self) -> str:
+        return KORMO_MODULE_ARCH.num_layers_config_key()
+
+    def layer_weights(
+        self, index: int, config: PretrainedConfig
+    ) -> Optional[List[WeightInfo]]:
+        num_experts = self.num_experts
+        prefix = f"model.layers.{index}"
+        tensor_names = []
+
+        # Expert weights 추가
+        for expert_idx in range(num_experts):
+            for param in ("gate_proj", "up_proj", "down_proj"):
+                tensor_names.append(
+                    prefix + f".mlp.experts.{expert_idx}.{param}.weight"
+                )
+
+        # Shared expert weights 추가
+        for param in ("gate_proj", "up_proj", "down_proj"):
+            tensor_names.append(
+                prefix + f".mlp.shared_expert.{param}.weight"
+            )
+
+        # Gate weights 추가
+        tensor_names.append(prefix + ".mlp.gate.weight")
+        tensor_names.append(prefix + ".mlp.shared_expert_gate.weight")
+
+        res = []
+        for name in tensor_names:
+            res.append(WeightInfo(name=name))
+
+        # 기존 KORMo weights 중에서 MLP를 제외한 것들 추가
+        for weight_info in KORMO_MODULE_ARCH.layer_weights(index, config):
+            if ".mlp." in weight_info.name:
+                continue
+            res.append(weight_info)
+
+        return res

mergekit/moe/__init__.py

@@ -6,14 +6,31 @@
 
 ALL_OUTPUT_ARCHITECTURES: List[MoEOutputArchitecture] = [MixtralMoE(), DeepseekMoE()]
 
+# Qwen3MoE를 먼저 추가
+try:
+    from mergekit.moe.qwen3 import Qwen3MoE
+except ImportError:
+    pass
+else:
+    ALL_OUTPUT_ARCHITECTURES.append(Qwen3MoE())
+
+# QwenMoE를 나중에 추가 (fallback용)
 try:
     from mergekit.moe.qwen import QwenMoE
 except ImportError:
     pass
 else:
     ALL_OUTPUT_ARCHITECTURES.append(QwenMoE())
 
+# KORMo MoE 추가
+try:
+    from mergekit.moe.kormo import KORMoMoE
+except ImportError:
+    pass
+else:
+    ALL_OUTPUT_ARCHITECTURES.append(KORMoMoE())
+
 __all__ = [
     "ALL_OUTPUT_ARCHITECTURES",
     "MoEOutputArchitecture",
-]
+]