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

变更摘要

本次变更主要为 NPU MoE（混合专家）模块中的 token unpermute 相关算子新增 DTensor 分片策略支持。在 torch_npu/distributed/tensor/_moe_ops.py 中新增了两个分片策略注册函数，分别为 _npu_moe_token_unpermute 和 npu_moe_token_unpermute_grad_v2 算子提供全复制（Replicate）和 hidden_size 维度分片（Shard(1)）两种分布式策略，使其能够在 DTensor 分布式场景下正确执行。

主要改动

• 新增 _npu_moe_token_unpermute_strategy 分片策略：通过 @register_sharding 装饰器为 npu._npu_moe_token_unpermute 算子注册 DTensor 策略，支持全复制和 hidden_size（第 1 维）分片两种策略，输出和输入张量均按对应策略标记 Replicate 或 Shard(1)。
• 新增 npu_moe_token_unpermute_grad_v2_strategy 分片策略：为 npu.npu_moe_token_unpermute_grad_v2 算子注册 DTensor 反向传播策略，同样支持全复制和 hidden_size 维度分片，其中分片策略下 probs 梯度标记为 Partial() 以适配分布式梯度聚合。

代码审查

✅ 未发现问题

CLA Signature Pass

ascend-ds-bot, thanks for your pull request. All authors of the commits have signed the CLA. 👍

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

💡 Tip:

• Committer can comment /approve or /lgtm
• Commenting /approve implies both code review (lgtm) and intent to merge (approve)

ascend docs pipeline is running...

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

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

/approve

/approve