现有 MegaMoe 提供低层前反向算子,接入 Torch MoE 模型时,调用方需要手动组织路由、执行计划、 通信 buffer 和资源生命周期。希望提供可直接接入模型的专家模块,降低训练集成和维护成本。
初版面向 Torch NPU BF16、固定 shape 和覆盖完整 world 的 EP。
mega_moe
mega_moe_grad
通过模块化接口提升易用性;通过容量检查、共享资源排序和统一释放提升可靠性; 以同配置对照验证训练精度、性能和显存成本。
新增 MegaMoeExperts 和资源管理基类 MulticoreModule。 MegaMoeExperts.forward 接收 hidden_states、topk_ids、topk_weights 和可选的 tokens_per_expert;通过 share_execution_resources() 共享资源,通过 close() 释放资源。
MegaMoeExperts
MulticoreModule
MegaMoeExperts.forward
hidden_states
topk_ids
topk_weights
tokens_per_expert
share_execution_resources()
close()
保留现有 mega_moe / mega_moe_grad 低层入口。Torch SHMEM 的初始化和释放统一到进程级 owner 管理, 需验证普通 SHMEM 调用与 managed 模块共存时的生命周期兼容性。
关联 PR:#1315。
验收要求:
具体验收数据随关联 PR 提供。
无
Thanks for contributing 🎉!
🚀 功能描述
[Feature]: 为 Torch MoE 提供 MegaMoe 专家模块与训练支持
🚀 功能描述
现有 MegaMoe 提供低层前反向算子,接入 Torch MoE 模型时,调用方需要手动组织路由、执行计划、
通信 buffer 和资源生命周期。希望提供可直接接入模型的专家模块,降低训练集成和维护成本。
初版面向 Torch NPU BF16、固定 shape 和覆盖完整 world 的 EP。
现有替代方案
mega_moe/mega_moe_grad,由模型侧维护路由和执行资源。与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。
验收要求:
不同 owner 关闭顺序均安全,末次关闭释放资源,稳定迭代无显存增长。
具体验收数据随关联 PR 提供。
现有替代方案
无
与DFX相关性DF
提议的新API
是否影响现有API
无
补充信息
Thanks for contributing 🎉!