feat(moe): NPU平台 EP experts forward计算改用 NPU grouped matmul#252
Open
werwrewe wants to merge 4 commits into
Open
feat(moe): NPU平台 EP experts forward计算改用 NPU grouped matmul#252werwrewe wants to merge 4 commits into
werwrewe wants to merge 4 commits into
Conversation
两个预存 bug 导致 EP 训练在 torch_npu 2.10 + CANN 9.0.0 上无法运行： 1. preprocess()：HCCL all_gather 在独立 stream 上执行，紧随其后的 host 读（.tolist()）与通信完成无序；且 non_blocking=True 的 D2H 拷贝未落地就被 host 消费，产生垃圾 split size（实测报错 "Trying to create tensor with negative dimension"）。 修复：集合通信后增加 torch.npu.synchronize()（仅 NPU）， 两处 D2H 拷贝改为同步。 2. permute()：index_select 的反向走 aclnnIndexAdd，在本环境对任意 dtype/shape 均失败（error 161001）。改为数学等价的高级索引 x[idx]（反向为 index_put，可正常累加重复下标的梯度）。
kernelize 的 NPU patch 无法生效，导致NPU forward被ep forward覆盖。考虑在ep forward中单独给NPU提供forward。ep forward原实现逐 expert Python 循环，在 EP 闭包内部复用现有 NPU 组件（GmmFunction + npu_swiglu + _get_cached_expert_weights）替换循环实现并行化。每次前向做资格检查，不满足条件自动回退原循环并打印一次 WARNING 说明原因。 端到端验证： loss 曲线与原实现一致，2*910B qwen3.5-35B-A3B 每 12 step 约 200s -> 60s。
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
kevssim
reviewed
Jul 23, 2026
| # observed in practice as garbage split sizes. Force a device sync before | ||
| # reading gathered results back on the host. | ||
| if num_local_tokens_per_expert.device.type == 'npu': | ||
| torch.npu.synchronize() |
Collaborator
There was a problem hiding this comment.
最好使用平台无关的封装：
from twinkle import torch_util
torch_util.synchronize()
kevssim
reviewed
Jul 23, 2026
| # all-to-all, causing EP collective order divergence. | ||
| return permuted_tokens | ||
|
|
||
| if _ep_experts_can_use_npu_gmm(self): |
Collaborator
There was a problem hiding this comment.
更好的做法是把gmm定义成一个接口，有不同的impl，然后在expert_parallel.py中直接调用接口，隐藏impl细节。
替换写死的 torch.npu.synchronize()，GPU/NPU 自动分发。
新增 kernel/ops/：EpExpertsGmm 基类 + 注册表 + 分发器，NPU 实现迁移为 NpuEpExpertsGmm；expert_parallel.py 只调接口，GMM 命中日志改为只打印一次。
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
kernelize 的 NPU patch 无法生效，导致NPU forward被ep forward覆盖。考虑在ep forward中单独给NPU提供forward。ep forward原实现逐 expert Python 循环，在 EP 闭包内部复用现有 NPU 组件（GmmFunction + npu_swiglu + _get_cached_expert_weights）替换循环实现并行化。每次前向做资格检查，不满足条件自动回退原循环并打印一次 WARNING 说明原因。
端到端验证： loss 曲线与原实现一致，使用ep_fsdp2_lora_qwen3_5_moe.sh 设备2*910B 模型qwen3.5-35B-A3B 每 12 step 约 200s -> 60s。