[Feature]: 为 Torch MoE 提供 MegaMoe 专家模块与训练支持

🚀 功能描述

[Feature]: 为 Torch MoE 提供 MegaMoe 专家模块与训练支持

🚀 功能描述

现有 MegaMoe 提供低层前反向算子，接入 Torch MoE 模型时，调用方需要手动组织路由、执行计划、
通信 buffer 和资源生命周期。希望提供可直接接入模型的专家模块，降低训练集成和维护成本。

• 接收模型 Router 的 Top-K 结果，支持输入、路由权重及专家权重的梯度，并复用 Router 已计算的 counts。
• 默认容量保证 lossless；显式容量不足时明确报错。
• 支持串行层共享 workspace，并统一普通 SHMEM 与模块资源的初始化、共享和释放。
• 提供 common MoE 对照、Qwen 训练示例和必要的使用指南。

初版面向 Torch NPU BF16、固定 shape 和覆盖完整 world 的 EP。

现有替代方案

• 直接调用 mega_moe / mega_moe_grad，由模型侧维护路由和执行资源。
• 使用现有 common MoE 与 EP 组合完成训练，作为精度和性能对照。

与DFX相关性DF

通过模块化接口提升易用性；通过容量检查、共享资源排序和统一释放提升可靠性；
以同配置对照验证训练精度、性能和显存成本。

提议的新API

新增 MegaMoeExperts 和资源管理基类 MulticoreModule。
MegaMoeExperts.forward 接收 hidden_states、topk_ids、topk_weights 和可选的
tokens_per_expert；通过 share_execution_resources() 共享资源，通过 close() 释放资源。

是否影响现有API

保留现有 mega_moe / mega_moe_grad 低层入口。Torch SHMEM 的初始化和释放统一到进程级 owner 管理，
需验证普通 SHMEM 调用与 managed 模块共存时的生命周期兼容性。

补充信息

关联 PR：#1315。

验收要求：

• 与 common MoE 的输出、全部梯度和 optimizer 更新一致，覆盖单层及 Qwen 整网训练。
• 多层共享、多个 outstanding forward 后 backward、交替 stream 和重复 step 正确；
  不同 owner 关闭顺序均安全，末次关闭释放资源，稳定迭代无显存增长。
• 默认容量、显式容量及省略 counts 的路径分别验收；性能包含 Router→专家完整路径，记录峰值和稳定显存。

具体验收数据随关联 PR 提供。

现有替代方案

无

与DFX相关性DF

提议的新API

是否影响现有API

无

补充信息

Thanks for contributing 🎉!

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 370
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/370

Start by reading the existing mega_moe and mega_moe_grad entry points, then review linked PR #1315 for the proposed MegaMoeExperts and MulticoreModule interfaces. Validate the requested single-layer and Qwen training comparisons, gradient and optimizer consistency, shared-resource lifetimes, capacity paths, and memory stability described in the acceptance requirements.