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

80 0 0

GoodNav 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 算子融合技术，显著提高了内存带宽利用率，缓解了 memory bound 瓶颈，进一步释放国产 GPU 的硬件潜力。

摩尔线程官方表示将持续对 MT-MegatronLM 和 MT-TransformerEngine 框架进行优化，并引入一系列功能，具体如下：

Dual Pipe / ZeroBubble 并行策略：进一步降低气泡率，提升并行训练效率；

多种 FP8 优化策略：独创 FP8 优化策略，以提高训练性能和稳定性；

异步 checkpoint 策略：增强训练过程中的容错能力和效率；

优化后的重计算策略：减少计算和显存开销，提升训练速度；

容错训练策略：独特的容错训练算法，增强训练过程中的容错能力；

集成摩尔线程 FlashMLA 和 DeepGemm 库：进一步释放摩尔线程 GPU 的算力和 FP8 计算能力，提高计算性能和效率。