[Bugfix][NPU] fix MTP training weight sync for MindSpeed GroupedGemm

vime/backends/megatron_utils/megatron_to_hf/glm4moe.py

@@ -2,6 +2,8 @@
 
 import torch
 
+from vime.utils.common import is_npu
+
 
 def convert_glm4moe_to_hf(args, name, param):
     if name == "module.module.embedding.word_embeddings.weight":
@@ -22,25 +24,30 @@ def convert_glm4moe_to_hf(args, name, param):
     if match:
         layer_idx, rest = match.groups()
 
-        # experts
-        expert_pattern = r"mlp.experts\.(.+)\.weight(\d+)"
-        match = re.match(expert_pattern, rest)
-        if match:
-            rest, expert_idx = match.groups()
-            if rest == "linear_fc1":
-                gate_weight, up_weight = param.chunk(2, dim=0)
-                outputs = [
-                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.gate_proj.weight", gate_weight),
-                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.up_proj.weight", up_weight),
-                ]
-                return outputs
-            elif rest == "linear_fc2":
-                outputs = [
-                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.weight", param),
-                ]
-                return outputs
-            else:
-                raise ValueError(f"Unknown expert parameter name: {name}")
+        if is_npu():
+            npu_outputs = _convert_npu_experts_and_mla(args, name, param, layer_idx, rest)
+            if npu_outputs is not None:
+                return npu_outputs
+        else:
+            # Standard Megatron: one set of weights per expert
+            expert_pattern = r"mlp.experts\.(.+)\.weight(\d+)"
+            match = re.match(expert_pattern, rest)
+            if match:
+                rest, expert_idx = match.groups()
+                if rest == "linear_fc1":
+                    gate_weight, up_weight = param.chunk(2, dim=0)
+                    outputs = [
+                        (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.gate_proj.weight", gate_weight),
+                        (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.up_proj.weight", up_weight),
+                    ]
+                    return outputs
+                elif rest == "linear_fc2":
+                    outputs = [
+                        (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.weight", param),
+                    ]
+                    return outputs
+                else:
+                    raise ValueError(f"Unknown expert parameter name: {name}")
 
         # shared expert
         shared_expert_pattern = r"mlp.shared_experts\.(.+)"
@@ -97,7 +104,7 @@ def convert_glm4moe_to_hf(args, name, param):
             ]
         elif rest == "mlp.linear_fc2.weight":
             return [(f"model.layers.{layer_idx}.mlp.down_proj.weight", param)]
-        elif rest == "self_attention.linear_qkv.layer_norm_weight":
+        elif rest == "self_attention.linear_qkv.layer_norm_weight" or (is_npu() and rest == "input_layernorm.weight"):
             return [(f"model.layers.{layer_idx}.input_layernorm.weight", param)]
         elif rest == "mlp.linear_fc1.layer_norm_weight":
             return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)]
@@ -114,9 +121,14 @@ def convert_glm4moe_to_hf(args, name, param):
 
         # qk norm
         elif rest == "self_attention.q_layernorm.weight":
+            if is_npu():
+                # MindSpeed MLA
+                return [(f"model.layers.{layer_idx}.self_attn.q_a_layernorm.weight", param)]
             return [(f"model.layers.{layer_idx}.self_attn.q_norm.weight", param)]
         elif rest == "self_attention.k_layernorm.weight":
             return [(f"model.layers.{layer_idx}.self_attn.k_norm.weight", param)]
+        elif is_npu() and rest == "self_attention.kv_layernorm.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.kv_a_layernorm.weight", param)]
 
     mtp_layer_pattern = r"module\.module\.mtp\.layers\.(\d+)\.(.+)"
     match = re.match(mtp_layer_pattern, name)
@@ -137,3 +149,115 @@ def convert_glm4moe_to_hf(args, name, param):
             return convert_glm4moe_to_hf(args, name, param)
 
     raise ValueError(f"Unknown parameter name: {name}")
+
+
+def _convert_npu_experts_and_mla(args, name, param, layer_idx, rest):
+    """MindSpeed GroupedGemm / MLA mappings used only on NPU."""
+    # MindSpeed GmmExpertsImpl: "mlp.experts.experts.linear_fc{1,2}.weight"
+    # Standard Megatron: "mlp.experts.linear_fc{1,2}.weight{N}"
+    expert_pattern = r"mlp.experts\.(.+)\.weight(\d*)$"
+    match = re.match(expert_pattern, rest)
+    if match:
+        fc_name, expert_idx = match.groups()
+        # Handle double "experts" in MindSpeed naming
+        if fc_name.startswith("experts."):
+            fc_name = fc_name[len("experts.") :]
+        if fc_name == "linear_fc1":
+            if expert_idx:
+                # Standard Megatron: one expert per param
+                gate_weight, up_weight = param.chunk(2, dim=0)
+                return [
+                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.gate_proj.weight", gate_weight),
+                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.up_proj.weight", up_weight),
+                ]
+            # MindSpeed GroupedGemm: all experts packed in one 3D param
+            # Shape: [num_experts, fc1_output, hidden_size]
+            # fc1_output = 2 * moe_ffn_hidden_size (gate+up packed)
+            # vLLM expects gate_proj/up_proj: [intermediate, hidden_size]
+            num_experts = args.num_experts
+            if param.dim() == 3:
+                # 3D: [num_experts, fc1_output, hidden_size] - slice on dim=1
+                gate_weight, up_weight = param.chunk(2, dim=1)
+                outputs = []
+                for i in range(num_experts):
+                    outputs.append((f"model.layers.{layer_idx}.mlp.experts.{i}.gate_proj.weight", gate_weight[i]))
+                    outputs.append((f"model.layers.{layer_idx}.mlp.experts.{i}.up_proj.weight", up_weight[i]))
+            else:
+                # 2D: [hidden_size, fc1_output * num_experts] - old format
+                gate_up = param.view(num_experts, args.hidden_size, -1)
+                gate_weight, up_weight = gate_up.chunk(2, dim=2)
+                outputs = []
+                for i in range(num_experts):
+                    outputs.append((f"model.layers.{layer_idx}.mlp.experts.{i}.gate_proj.weight", gate_weight[i].t()))
+                    outputs.append((f"model.layers.{layer_idx}.mlp.experts.{i}.up_proj.weight", up_weight[i].t()))
+            return outputs
+        elif fc_name == "linear_fc2":
+            if expert_idx:
+                # Standard Megatron: one expert per param
+                return [
+                    (f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.weight", param),
+                ]
+            # MindSpeed GroupedGemm: all experts packed in one 3D param
+            # Shape: [num_experts, hidden_size, fc2_input]
+            # vLLM expects down_proj: [hidden_size, intermediate]
+            num_experts = args.num_experts
+            if param.dim() == 3:
+                # 3D: [num_experts, hidden_size, fc2_input] - use directly
+                outputs = []
+                for i in range(num_experts):
+                    outputs.append((f"model.layers.{layer_idx}.mlp.experts.{i}.down_proj.weight", param[i]))
+            else:
+                # 2D: [fc2_input * num_experts, hidden_size] - old format
+                down = param.view(num_experts, -1, args.hidden_size)
+                outputs = []
+                for i in range(num_experts):
+                    outputs.append((f"model.layers.{layer_idx}.mlp.experts.{i}.down_proj.weight", down[i].t()))
+            return outputs
+        else:
+            raise ValueError(f"Unknown expert parameter name: {name}")
+
+    # GroupedGemm format: weight1/weight2 (all experts packed together)
+    if rest == "mlp.experts.weight1":
+        # 3D: [num_experts, fc1_output, hidden_size] (MindSpeed GmmExpertsImpl)
+        # 2D: [hidden_size, fc1_output * num_experts] (after EP all-gather + concat)
+        num_experts = args.num_experts
+        if param.dim() == 3:
+            gate_weight, up_weight = param.chunk(2, dim=1)
+            outputs = []
+            for i in range(num_experts):
+                outputs.append((f"model.layers.{layer_idx}.mlp.experts.{i}.gate_proj.weight", gate_weight[i]))
+                outputs.append((f"model.layers.{layer_idx}.mlp.experts.{i}.up_proj.weight", up_weight[i]))
+        else:
+            gate_up = param.view(num_experts, args.hidden_size, -1)
+            gate_weight, up_weight = gate_up.chunk(2, dim=2)
+            outputs = []
+            for i in range(num_experts):
+                outputs.append((f"model.layers.{layer_idx}.mlp.experts.{i}.gate_proj.weight", gate_weight[i].t()))
+                outputs.append((f"model.layers.{layer_idx}.mlp.experts.{i}.up_proj.weight", up_weight[i].t()))
+        return outputs
+    elif rest == "mlp.experts.weight2":
+        # 3D: [num_experts, hidden_size, fc2_input] (MindSpeed GmmExpertsImpl)
+        # 2D: [fc2_input * num_experts, hidden_size] (after EP all-gather + concat)
+        num_experts = args.num_experts
+        if param.dim() == 3:
+            outputs = []
+            for i in range(num_experts):
+                outputs.append((f"model.layers.{layer_idx}.mlp.experts.{i}.down_proj.weight", param[i]))
+        else:
+            down = param.view(num_experts, -1, args.hidden_size)
+            outputs = []
+            for i in range(num_experts):
+                outputs.append((f"model.layers.{layer_idx}.mlp.experts.{i}.down_proj.weight", down[i].t()))
+        return outputs
+
+    # MindSpeed MLA attention: separate q/kv projections
+    if rest == "self_attention.linear_q_down_proj.weight":
+        return [(f"model.layers.{layer_idx}.self_attn.q_a_proj.weight", param)]
+    elif rest == "self_attention.linear_q_up_proj.weight":
+        return [(f"model.layers.{layer_idx}.self_attn.q_b_proj.weight", param)]
+    elif rest == "self_attention.linear_kv_down_proj.weight":
+        return [(f"model.layers.{layer_idx}.self_attn.kv_a_proj_with_mqa.weight", param)]
+    elif rest == "self_attention.linear_kv_up_proj.weight":
+        return [(f"model.layers.{layer_idx}.self_attn.kv_b_proj.weight", param)]
+
+    return None

vime/backends/megatron_utils/update_weight/update_weight_from_distributed.py

@@ -31,7 +31,9 @@
 def _begin_vllm_weight_update_session(rollout_engines: Sequence[ActorHandle]) -> None:
     if dist.get_rank() == 0:
         logger.info("vLLM weight update: start_weight_update")
-        ray.get([engine.start_weight_update.remote(is_checkpoint_format=True) for engine in rollout_engines])
+        # NPU layerwise_reload corrupts weights; use direct mode (same as colocate).
+        _ckpt_fmt = not is_npu()
+        ray.get([engine.start_weight_update.remote(is_checkpoint_format=_ckpt_fmt) for engine in rollout_engines])
     dist.barrier(group=get_gloo_group())
 
 
@@ -274,31 +276,62 @@ def _ep_gather_and_convert(self, named_tensors: list[tuple[str, torch.Tensor]])
         EP all-gather a buffered batch + HF convert on PP source. Returns HF tensors on
         PP source, [] elsewhere. Clears ``named_tensors``.
         """
+        ep_world_size = mpu.get_expert_model_parallel_world_size()
         names = [name for name, _ in named_tensors]
-        all_names = [None] * mpu.get_expert_model_parallel_world_size()
+        all_names = [None] * ep_world_size
         dist.all_gather_object(all_names, names, group=mpu.get_expert_model_parallel_group())
 
-        for names in all_names:
-            assert len(named_tensors) == len(names), f"mismatch names length: {len(named_tensors)} != {len(names)}"
+        for names_list in all_names:
+            assert len(named_tensors) == len(
+                names_list
+            ), f"mismatch names length: {len(named_tensors)} != {len(names_list)}"
 
-        all_gathered_params = [[] for _ in range(mpu.get_expert_model_parallel_world_size())]
+        # NPU MindSpeed GroupedGemm: same param name across all EP ranks
+        # (e.g., mlp.experts.weight1/weight2 or MindSpeed linear_fc1/fc2).
+        is_npu_grouped_gemm = is_npu() and ep_world_size > 1 and all(names_list == names for names_list in all_names)
+
+        device = torch.npu.current_device() if is_npu() else torch.cuda.current_device()
+        all_gathered_params = [[] for _ in range(ep_world_size)]
         handles = []
         for i, (_name, param) in enumerate(named_tensors):
-            params = [
-                torch.empty_like(param.data, device=torch.cuda.current_device())
-                for _ in range(mpu.get_expert_model_parallel_world_size())
-            ]
+            params = [torch.empty_like(param.data, device=device) for _ in range(ep_world_size)]
             handle = dist.all_gather(params, param.data, group=mpu.get_expert_model_parallel_group(), async_op=True)
             handles.append(handle)
-            for ep_rank, names in enumerate(all_names):
-                all_gathered_params[ep_rank].append((names[i], params[ep_rank]))
+            for ep_rank, names_list in enumerate(all_names):
+                all_gathered_params[ep_rank].append((names_list[i], params[ep_rank]))
         for handle in handles:
             handle.wait()
 
         named_tensors.clear()
         if not self._is_pp_src_rank:
             return []
 
+        if is_npu_grouped_gemm:
+            # GroupedGemm: concatenate params from all EP ranks along expert dimension,
+            # then convert the full param (with all experts) to HF format.
+            # 3D MindSpeed: [num_local_experts, ...] -> concat dim=0
+            # 2D weight1/linear_fc1: [hidden, fc * num_local] -> concat dim=1
+            # 2D weight2/linear_fc2: [fc * num_local, hidden] -> concat dim=0
+            saved_names = list(names)
+            converted_hf_tensors = []
+            for i, name in enumerate(saved_names):
+                sample = all_gathered_params[0][i][1]
+                if sample.dim() == 3:
+                    concat_dim = 0
+                elif "weight1" in name or "linear_fc1" in name:
+                    concat_dim = 1
+                else:
+                    concat_dim = 0
+                full_param = torch.cat(
+                    [all_gathered_params[ep_rank][i][1] for ep_rank in range(ep_world_size)],
+                    dim=concat_dim,
+                )
+                converted_hf_tensors += convert_to_hf(
+                    self.args, self.model_name, name, full_param, self.quantization_config
+                )
+            return converted_hf_tensors
+
+        # Original ungrouped format: each EP rank has different expert indices
         all_gathered_params = sum(all_gathered_params, [])
         converted_hf_tensors = []
         for name, param in all_gathered_params: