[Bugfix][NPU] fix MTP training weight sync for MindSpeed GroupedGemm - #375
[Bugfix][NPU] fix MTP training weight sync for MindSpeed GroupedGemm#375CalvinXKY wants to merge 1 commit into
Conversation
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 <kyxiezju@163.com>
Documentation build overview
48 files changed ·
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds NPU-specific support for GLM4MoE weight conversion, including handling MindSpeed GroupedGemm and MLA attention mappings, and disables checkpoint format during NPU weight updates to prevent corruption. The review feedback highlights critical issues where directly calling .view() on 2D expert weights would scramble them due to memory layout, suggesting a .reshape().permute() sequence instead. Additionally, the reviewer recommends using .reshape() instead of .view() elsewhere to safely handle potentially non-contiguous tensors.
| gate_up = param.view(num_experts, args.hidden_size, -1) | ||
| gate_weight, up_weight = gate_up.chunk(2, dim=2) |
There was a problem hiding this comment.
In the 2D case, param has the shape [hidden_size, fc1_output * num_experts]. Directly calling .view(num_experts, args.hidden_size, -1) will scramble the weights because the expert dimension is interleaved in memory. To correctly reshape and split the expert dimension, you should first reshape to [hidden_size, num_experts, fc1_output] and then permute to [num_experts, hidden_size, fc1_output]. Additionally, using .reshape() is safer than .view() to avoid potential runtime errors on non-contiguous tensors.
| gate_up = param.view(num_experts, args.hidden_size, -1) | |
| gate_weight, up_weight = gate_up.chunk(2, dim=2) | |
| gate_up = param.reshape(args.hidden_size, num_experts, -1).permute(1, 0, 2) | |
| gate_weight, up_weight = gate_up.chunk(2, dim=2) |
| gate_up = param.view(num_experts, args.hidden_size, -1) | ||
| gate_weight, up_weight = gate_up.chunk(2, dim=2) |
There was a problem hiding this comment.
Similar to the 2D case in linear_fc1, param here has the shape [hidden_size, fc1_output * num_experts]. Directly calling .view(num_experts, args.hidden_size, -1) will scramble the weights. You should reshape to [hidden_size, num_experts, fc1_output] and then permute to [num_experts, hidden_size, fc1_output]. Using .reshape() is also safer than .view() here.
| gate_up = param.view(num_experts, args.hidden_size, -1) | |
| gate_weight, up_weight = gate_up.chunk(2, dim=2) | |
| gate_up = param.reshape(args.hidden_size, num_experts, -1).permute(1, 0, 2) | |
| gate_weight, up_weight = gate_up.chunk(2, dim=2) |
| 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) |
There was a problem hiding this comment.
Using .reshape() is safer than .view() here to prevent potential runtime errors if the tensor is non-contiguous.
| down = param.view(num_experts, -1, args.hidden_size) | |
| down = param.reshape(num_experts, -1, args.hidden_size) |
| 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) |
There was a problem hiding this comment.
Using .reshape() is safer than .view() here to prevent potential runtime errors if the tensor is non-contiguous.
| down = param.view(num_experts, -1, args.hidden_size) | |
| down = param.reshape(num_experts, -1, args.hidden_size) |
|
1 participant
Details
Fix NPU MTP acceptance collapsing to 0% in disaggregated RL training.
NPU-only (
is_npu()gated):glm4moe.py(wrongview()corrupted experts)layerwise_reloadin disaggregated weight update (is_checkpoint_format=not is_npu())GPU path unchanged.
Companion vllm-ascend fix: vllm-project/vllm-ascend#13209
Testing
GLM-4.7-Flash (30B-A3B MoE + MTP), 8x Ascend 910B1, disaggregated 4 train + 4 rollout, TP=4, EP=4 (together with vllm-ascend#13209):
Pure-inference MTP baseline (same machine): +82.5% throughput (49.6 → 90.6 tok/s).