摩尔线程开源 MT-MegatronLM 和 MT-TransformerEngine 两大 AI 框架

xxn 阅读:94601 2025-03-17 20:03:44 评论:0

IT之家 3 月 17 日消息,摩尔线程官方今日发文宣布,已正式推出 MT-MegatronLM 与 MT-TransformerEngine 两款 AI 框架的开源版本。这两款框架结合了 FP8 混合训练战略和高性能算子库,能在国产全功能 GPU 上进行混合并行训练和推理,显著提升了训练效率与稳定性。

据介绍,MT-MegatronLM 是专为全功能 GPU 设计的混合并行训练框架,支持 dense 模型、多模态模型和 MoE(混合专家)模型的高效训练;MT-TransformerEngine 则主要用于 Transformer 模型的训练与推理优化,通过算子融合、并行加速等技术,充分释放了摩尔线程全能 GPU 的计算潜力和 memory bound 算子效率。这两个框架在硬件适配和算法创新上取得了重要突破:

  • 混合并行训练:支持 Dense、多模态及 MoE 模型的混合并行训练,可灵活应对不同模型架构的复杂运算场景;

  • FP8 混合训练战略:结合摩尔线程 GPU 原生支持的 FP8 混合精度训练战略,能有效提升训练效率;

  • 高性能算子库:通过 muDNN 和 MCCL 高性能算子库的深度集成,系统性地优化了计算密集型任务和多卡协同通信的开销;同时,结合摩尔线程开发的开源 Simumax 库,可自动进行并行策略搜索,最大化各种模型和加速环境下的并行训练性能;

  • 异常训练处理:框架内置的 rewind 异常恢复机制,可自动退回到最近的稳定点继续训练,极大地提升了大规模训练的稳定性;

  • 完整的兼容性:两个框架兼容 GPU 主流生态系统,不仅确保了现有生态平滑迁移,也为开发者构建自有的 AI 技术栈提供了底层支持。

实际应用展示如下:

  • 高效训练:在全功能 GPU 集群上,Llama3 8B 模型的训练任务利用 FP8,loss 几乎没有损失,MFU 达到了90% 以上;(如下图所示)

▲ 利用摩尔线程 FP8 混合精度加速技术,在 loss 无损情况下,获得 28% 的加速效果
  • DeepSeek 全新训练:摩尔线程深度整合并开源了 DeepSeek 并行算法 DualPipe 的高效支持,MT-DualPipe 完全兼容 MT-Megatron 和 MT-TransformerEngine 框架,成功实现了 DeepSeek V3 训练流程的完整还原,支持 MLA、MTP 和多种专家平衡策略;

  • 性能大幅优化:通过多种 Transformer 算子融合技术,明显提高了内存带宽利用率,有效缓解了内存瓶颈问题,进一步释放了国产 GPU 的硬件潜力。

摩尔线程官方表示将持续优化 MT-MegatronLM 和 MT-TransformerEngine 框架,并引入系列功能,具体如下:

  • Dual Pipe / ZeroBubble 并行策略:进一步降低气泡率,提升并行训练效率;

  • 多种 FP8 优化策略:独特的 FP8 优化策略,提高训练性能和稳定性;

  • 异步 checkpoint 策略:增强了训练过程中的容错能力和效率;

  • 改进后的重计算策略:减少了计算和显存开销,加快了训练速度;

  • 容错训练策略:独创的容错训练算法,增强了训练过程中的容错能力;

  • 集成摩尔线程 FlashMLA 和 DeepGemm 库:深入释放摩尔线程 GPU 的算力和 FP8 计算能力,提升了计算性能和效率。

IT之家提供的开源地址如下:

广告声明:文内包含的外部链接(包括但不限于超链接、二维码、口令等),旨在传递更多信息,节省您的筛选时间,仅供参考。

声明

1.本站遵循行业规范,任何转载的稿件都会明确标注作者和来源;2.本站的原创文章,请转载时务必注明文章作者和来源,不尊重原创的行为我们将追究责任;3.作者投稿可能会经我们编辑修改或补充。

搜索
排行榜
关注我们

扫一扫关注我们,了解最新精彩内容