已开启
[Feature]: 为 Torch MoE 提供 MegaMoe 专家模块与训练支持 #370
Feiran Zhang创建于  14 天前
Feiran Zhang
Feiran Zhang
14 天前 创建

🚀 功能描述

[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_statestopk_idstopk_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 🎉!

likedislike