[sync] PR-40312: Update moe token unpermute DTensor support-pytorch-AtomGit

Star843

Star843

已合并

[sync] PR-40312: Update moe token unpermute DTensor support #42584

shawnylee233创建于 27 天前从shawnylee233/pytorch： feature/moe_token_permute_memory_opt_v2.12.0 合入到 Ascend/pytorch： v2.12.0

[sync] PR-40312: Update moe token unpermute DTensor support #42584

已合并

shawnylee233创建于 27 天前

从shawnylee233/pytorch feature/moe_token_permute_memory_opt_v2.12.0合入到Ascend/pytorch v2.12.0

shawnylee233成员

27 天前

27 天前

### 1. Origin pull request: https://gitcode.com/Ascend/pytorch/merge_requests/40312 ### 2. Original pull request related issue(s): https://gitcode.com/Ascend/pytorch/issues/3108 ### 3. Original pull request related commit(s): | Sha | Datetime | Message | |---|---|---| |[fc5ecd2f](https://gitcode.com/Ascend/pytorch/commit/fc5ecd2f0f2456827110f8f0a635d9cb51668594)|2026-07-16 09:43:21 +0800 CST|Update moe token unpermute DTensor support
|

Pull Request已成功合入, 合并人@ascend-robot

（感谢 shawnylee233 的贡献）

S shawnylee233成员

27 天前创建了 pull request，commit 01fef155

27 天前创建了 pull request，commit 01fef155

ascend-robot成员

27 天前评论：

27 天前评论：

openLiBing可接管PR, PR详情链接:
https://www.openlibing.com/apps/transitionPrDetail?owner=Ascend&repo=pytorch&number=42584&platform=gitcode&codeHostingPlatformFlag=gitcode

27 天前评论：

27 天前评论：

变更摘要

此PR为MoE（混合专家）模块中的token unpermute操作及其反向传播V2版本新增了DTensor分布式张量分片策略支持。通过在 torch_npu/distributed/tensor/_moe_ops.py 中注册两个新的策略函数，使得NPU上的MoE token unpermute操作能够在分布式场景下选择全复制（Replicate）或hidden_size维度分片（Shard）两种并行策略。

主要改动

新增 _npu_moe_token_unpermute_strategy 策略函数：为 npu._npu_moe_token_unpermute 算子注册DTensor分片策略，支持全复制策略（输入输出均为 Replicate()）和hidden_size维度分片策略（permuted_tokens 与输出在维度1使用 Shard(1)，sorted_indices 保持 Replicate()）。
新增 npu_moe_token_unpermute_grad_v2_strategy 策略函数：为 npu.npu_moe_token_unpermute_grad_v2 反向算子注册DTensor分片策略，支持全复制策略和hidden_size维度分片策略，其中分片策略下 permuted_tokens 梯度在维度1使用 Shard(1)，probs 梯度使用 Partial() 归约。

27 天前评论：

27 天前评论：

代码审查

✅ 未发现问题

ascend-robot成员

27 天前添加了label：ascend-cla/yes

27 天前添加了label：ascend-cla/yes

ascend-robot成员

27 天前评论：

27 天前评论：

Thanks for your pull-request.
The full list of commands accepted by me can be found at here。
You can get sig-info at here

PR Approval Progress

✅ Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

module	lgtm status	approve status
torch_npu/distributed	✅ liangsongwei, chengpeng25 (2/2)	✅ liangsongwei (1/1)

💡 Tip:

Committer can comment /approve or /lgtm

Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

shawnylee233, thanks for your pull request. All authors of the commits have signed the CLA. 👍

ascend-robot成员

27 天前添加了label：needs-issue

27 天前添加了label：needs-issue

ascend-robot成员

27 天前评论：

27 天前评论：

Linking Issue Notice

@shawnylee233 , the pull request must be linked to at least one issue.
If an issue has already been linked, but the needs-issue label remains, you can remove the label by commenting /check-issue .

ascend-robot成员

27 天前评论：

27 天前评论：

当前仓库存在以下 保护分支 ：

Protected Branch	Version	Release
master
v2.7.1-26.1.0
v2.7.1
v2.9.0
v2.10.0
v2.11.0
v2.12.0
v2.10.0-26.1.0
v2.11.0-26.1.0
v2.12.0-26.1.0
v2.9.0-26.1.0
ci-test

评论 /sync <branch1> <branch2> ... 可将当前 PR 修改同步到其它分支（创建同步 PR）：
a) 如果当前 PR 是 Open 状态，同步操作将延迟到 PR 被合并时执行
b) 如果当前 PR 已经 Merged，将立即执行同步操作

注意：

/sync 命令可以指定同步到多个分支，仅最后一个 /sync 命令生效

如果创建的同步 PR 不正确，可通过向同步 PR 的源分支提交轻量级 PR 完善，或使用 /close 命令关闭

S shawnylee233成员

27 天前修改标题为 “[sync] PR-40312: Update moe token unpermute DTensor support”，原标题为“Update moe token unpermute DTensor support”

27 天前修改标题为 “[sync] PR-40312: Update moe token unpermute DTensor support”，原标题为“Update moe token unpermute DTensor support”

S shawnylee233成员

27 天前修改了pull request 的描述

27 天前修改了pull request 的描述

ascend-robot成员

27 天前添加了label：ci-pipeline-running

27 天前添加了label：ci-pipeline-running

ascend-robot成员

27 天前删除了label：needs-issue

27 天前删除了label：needs-issue

ascend-robot成员

27 天前评论：

27 天前评论：

ascend docs pipeline is running...

ascend-robot成员

27 天前添加了label：docs-ci-pipeline-running

27 天前添加了label：docs-ci-pipeline-running

ascend-robot成员

27 天前评论：

27 天前评论：

✅ 跳过 docs ci 检查，没有需要检查的文档文件

ascend-robot成员

27 天前删除了label：docs-ci-pipeline-running

27 天前删除了label：docs-ci-pipeline-running

ascend-robot成员

27 天前添加了label：docs-ci-pipeline-success

27 天前添加了label：docs-ci-pipeline-success

ascend-robot成员

26 天前删除了label：ci-pipeline-running

26 天前删除了label：ci-pipeline-running

ascend-robot成员

26 天前添加了label：ci-pipeline-passed

26 天前添加了label：ci-pipeline-passed

ascend-robot成员

26 天前评论：

26 天前评论：

流水线 PR-pipeline_pytorch#48957 [ commitID：9a8011a4 ] 已完成

>>>代码风格自动修复执行成功（无修复内容）

阶段	任务名	状态	详情
编译构建	Build_X86	✅	>>>
	Build_ARM	✅	>>>
	Build_LibTorch_x86	✅	>>>
	Build_LibTorch_ARM	✅	>>>
	Build_X86_torchair	🛑	>>>
	Build_ARM_torchair	🛑	>>>
	patch_test	🛑	>>>
恶意代码检查	Antipoison	✅	>>>
编码安全与规范检查	CodeCheck	✅	>>>
	check_error	✅	>>>
	CodeCheck_lintrunner	✅	>>>
开源片段检查	SCA	✅	>>>
开发者测试	UT_ARM_A3_Part_01	🛑	>>>
	UT_ARM_A3_Part_02	🛑	>>>
	UT_ARM_A2_Part_01	✅	>>>
	UT_ARM_A2_Part_02	✅	>>>
	UT_ARM_A2_Part_03	✅	>>>
	UT_inductor_Part_01	🛑	>>>
	UT_inductor_Part_02	🛑	>>>
	UT_inductor_Part_03	🛑	>>>
	UT_inductor_Part_04	🛑	>>>
	UT_DIST_ARM_Part_01	✅	>>>
	UT_DIST_ARM_Part_02	✅	>>>
	UT_DIST_ARM_Part_03	✅	>>>
	UT_DIST_ARM_Part_04	✅	>>>
	UT_ARM_A2_Select_Part_01	✅	>>>
	UT_ARM_A2_Select_Part_02	✅	>>>
流水线	PR-pipeline_pytorch	✅	>>>

此流水线已支持下列评论快捷指令，仅PR创建者和白名单成员[wujinyuan1, huangjingwei, liangsongwei, yashi999, culechan, Dring, wuyouqi1, L1919_snow, qq_52711437, WhiteNight12, nomiz, xiu_21, ffmh, wanglijun55, hss-shuai, husichao, smallsilly, lanshaozuishuai, jimmyisme1, lzy0920232, alpha-junh, Sunshine_Youngster, wei_zhuoyi, zhangyihuiben, zyw-hw, zzzkeke, rmch, yangch0324, LucciC, AACAES, renyujin, wjlflyer, senzhen-town, pengjingyou, qsc97, limuan, yule100, xiaoqi-zhou, kuhn7, chenxingying, hanye02, zichun_ye, anyrenwei, kkjocker, wangzili121, Lu_G, yvjc, puddingfjz, HandsoemLemon, bigprestigee1, huawuyi, zhenyu10, dairenjie, du-jin-hang, zou-jieyu, adelaideliu, TrHan, wanlinan, Windwindzzz, pengqihw, kisnwang, yuheng_wang, honghao_wang, jizewei, zhangguoguang, sunyu-xuan, chenrayray, hbhu_bin, liujunzhu, c_34, LiNuoh, maoyuanpeng1, zzhongmin, zhaoyu65, bellatan, jiabaolin, zhuofanshen, wencaiwen, lu_zhuge, caoshuyang, molly12, lyx324521, LQ1206, gitcode-chenjiao, cai-weiwei1989, CHDong, ogqin, yuanlipingGit, xuqinglin1, lqz2, zouwei1, chaoluoa, paradox325, jackzhang1116, yaoyao, akh, yujiacheng, dengjie0116, Hubert11111, Shine_Ws, wslhj555, longqiand, OYtao666, JiaqingQiang, luyyyy, Kingbelial, zhanghaiyu0101, wenxp1018, yanliu-luoluo, ksun_sekiro, liyong328, wgzheng, tangky, vivi_is_coding, aoiaoisola, weixin_44494597, wangmengmengwang65667, hid57809721, qq_35468730, comeonup, C547032, gcw_m5OQChA4, yao_yao_ling_xian, cnnbwcy, szqfes_12, cora_19, cann_lilin, can, shawnylee233, fanglanyue0916, hhz0, LiNuohang, taohuoquan]评论有效

compile、compile_inductor、compile_torchair : 运行流水线
retry : 重试流水线所有失败子任务
retry <任务名> : 仅重试指定失败子任务
stop : 停止流水线

梁松伟成员

26 天前评论：

26 天前评论：

/approve

ascend-robot成员

26 天前添加了label：approved

26 天前添加了label：approved

chengpeng25成员

26 天前评论：

26 天前评论：

/lgtm

ascend-robot成员

26 天前添加了label：lgtm

26 天前添加了label：lgtm

ascend-robot成员

26 天前合入了pull request

26 天前合入了pull request

ascend-robot成员

26 天前评论：

26 天前评论：

流水线 pytorch_gitcode_PR_multiVersion#13101 [ commitID：9a8011a4 ] 已完成