From bef2d93bcbbf8e817cb0360c40060435a0b0091d Mon Sep 17 00:00:00 2001 From: kaiyuan Date: Thu, 30 Jul 2026 19:50:02 +0800 Subject: [PATCH] [NPU] fix MTP training weight sync for MindSpeed GroupedGemm Correct 3D expert weight conversion and disable NPU layerwise_reload in disaggregated update so MTP acceptance recovers during train/infer split. NPU-only paths are gated by is_npu(); GPU behavior is unchanged. Signed-off-by: kaiyuan --- .../megatron_utils/megatron_to_hf/glm4moe.py | 164 +++++++++++++++--- .../update_weight_from_distributed.py | 55 ++++-- 2 files changed, 188 insertions(+), 31 deletions(-) diff --git a/vime/backends/megatron_utils/megatron_to_hf/glm4moe.py b/vime/backends/megatron_utils/megatron_to_hf/glm4moe.py index 33a64e6e8..9af9c644d 100644 --- a/vime/backends/megatron_utils/megatron_to_hf/glm4moe.py +++ b/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 diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/vime/backends/megatron_utils/update_weight/update_weight_from_distributed.py index dd7ff1bb1..adf1ce857 100644 --- a/vime/backends/megatron_utils/update_weight/update_weight_from_distributed.py +++ b/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,24 +276,29 @@ 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() @@ -299,6 +306,32 @@ def _ep_gather_and_convert(self, named_tensors: list[tuple[str, torch.Tensor]]) 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: