近日,摩尔线程正式开源MT-MegatronLM与MT-TransformerEngine两大AI框架,这一举措在国产GPU领域掀起波澜,为AI训练和推理带来全新的国产化解决方案。
MT-MegatronLM是面向全功能GPU的开源混合并行训练框架,支持dense模型、多模态模型及MoE(混合专家)模型的高效训练。它利用全功能GPU支持的FP8混合精度策略,搭配高性能算子库muDNN与集合通信库MCCL,显著提升了国产全功能GPU集群的算力利用率。
MT-TransformerEngine则专注于Transformer模型的高效训练与推理优化,通过算子融合、并行加速策略等技术,充分挖掘摩尔线程全功能GPU高密度计算的潜力,提升memory bound算子的效率。
这两大框架的技术突破体现在硬件适配与算法创新的深度协同。在混合并行训练方面,它们支种模型架构的复杂运算场景;FP8混合训练策略结合摩尔线程GPU原生支持的FP8混合精度训练,有效提升了训练效率。高性能算子库与通信库的深度集成,优化了计算密集型任务与多卡协同的通信开销,并通过摩尔线程开源Simumax库,可自动进行并行策略搜索,最大化并行训练性能。框架内置的rewind异常恢复机制,能自动回滚至最近稳定节点继续训练,提升了大规模训练的稳定性。此外,两个框架兼容GPU主流生态,保障了现有生态的平滑迁移,也为开发者构建自有的AI技术栈提供了底层支撑。
在实际应用中,这两个框架成果显著。在全功能GPU集群上,Llama3 8B模型的训练任务,利用FP8在loss几乎无损的情况下MFU达到90%以上。(如下图所示)
图注:利用摩尔线程FP8混合精度加速技术在loss无损的情况下得到28%的加速
摩尔线程还深度集成并开源对DeepSeek并行算法DualPipe的高效支持,成功复现DeepSeek V3训练流程。通过多种Transformer算子融合技术,显著提升了内存带宽利用率,释放了国产GPU的硬件潜力。
为加速国产GPU生态发展与建设,摩尔线程将持续优化这两大框架。未来将引入Dual Pipe/ZeroBubble并行策略降低气泡率,提升并行训练效率;推出多种FP8优化策略:独创的FP8优化策略提高训练的性能和稳定性;采用异步checkpoint策略提高训练过程中的容错能力和效率;优化重计算策略减少计算和显存开销;运用容错训练策略增强训练过程中的容错能力;集成摩尔线程FlashMLA和DeepGemm库,进一步释放摩尔线程GPU的算力和FP8计算能力。
摩尔线程通过开源MT-MegatronLM与MT-TransformerEngine,不仅展示了其在国产GPU技术上的深厚积累,更为国产GPU在AI大模型领域的应用提供了有力支持,加速了国产全功能GPU在AI计算领域的规模化应用进程。
