3 月 17 日消息，摩爾線程官方今日發(fā)文宣布，已正式開源 MT-MegatronLM 與 MT-TransformerEngine 兩大 AI 框架。這兩大框架通過深度融合 FP8 混合訓(xùn)練策略和高性能算子庫，在國產(chǎn)全功能 GPU 上實現(xiàn)混合并行訓(xùn)練和推理，提升了訓(xùn)練效率與穩(wěn)定性。

據(jù)介紹，MT-MegatronLM 是面向全功能 GPU 的開源混合并行訓(xùn)練框架，支持 dense 模型、多模態(tài)模型及 MoE（混合專家）模型的高效訓(xùn)練；MT-TransformerEngine 主要用于 Transformer 模型的訓(xùn)練與推理優(yōu)化，通過算子融合、并行加速策略等技術(shù)，釋放摩爾線程全功能 GPU 高密度計算的潛力和 memory bound 算子效率。兩大框架的技術(shù)突破體現(xiàn)在硬件適配與算法創(chuàng)新的深度協(xié)同：

• 混合并行訓(xùn)練：支持 Dense、多模態(tài)及 MoE 模型的混合并行訓(xùn)練，可靈活應(yīng)對不同模型架構(gòu)的復(fù)雜運算場景；
• FP8 混合訓(xùn)練策略：結(jié)合摩爾線程 GPU 原生支持的 FP8 混合精度訓(xùn)練策略，能夠有效提升訓(xùn)練效率；
• 高性能算子庫：通過高性能算子庫 muDNN 與通信庫 MCCL 的深度集成，系統(tǒng)性優(yōu)化了計算密集型任務(wù)與多卡協(xié)同的通信開銷；同時結(jié)合摩爾線程開源 Simumax 庫，可自動進行并行策略搜索，并針對不同模型和加速環(huán)境 spec 最大化并行訓(xùn)練性能；
• 異常訓(xùn)練處理：框架內(nèi)置的 rewind 異常恢復(fù)機制，可自動回滾至最近穩(wěn)定節(jié)點繼續(xù)訓(xùn)練，大幅提升大規(guī)模訓(xùn)練的穩(wěn)定性；
• 完整的兼容性：兩個框架兼容 GPU 主流生態(tài)，既保障了現(xiàn)有生態(tài)的平滑遷移，也為開發(fā)者構(gòu)建自有的 AI 技術(shù)棧提供了底層支撐。

實際應(yīng)用效果如下：

• 高效訓(xùn)練：在全功能 GPU 集群上，Llama3 8B 模型的訓(xùn)練任務(wù)，可以利用 FP8 在 loss 幾乎無損的情況下 MFU 達到 90% 以上；（如下圖所示）

▲ 利用摩爾線程 FP8 混合精度加速技術(shù)在 loss 無損的情況下得到 28% 的加速。

• 復(fù)現(xiàn) DeepSeek 滿血版訓(xùn)練：摩爾線程已深度集成并開源對 DeepSeek 并行算法 DualPipe 的高效支持，MT-DualPipe 可以完整接入 MT-Megatron 框架和 MT-TransformerEngine 框架，成功實現(xiàn) DeepSeek V3 訓(xùn)練流程的完整復(fù)現(xiàn)，支持 MLA、MTP 及多種專家平衡策略；
• 性能大幅優(yōu)化：通過多種 Transformer 算子融合技術(shù)，顯著提升了內(nèi)存帶寬利用率，有效緩解 memory bound 瓶頸，進一步釋放國產(chǎn) GPU 的硬件潛力。

摩爾線程官方表示將持續(xù)優(yōu)化 MT-MegatronLM 與 MT-TransformerEngine 框架，并引入系列功能，具體如下：

• Dual Pipe / ZeroBubble 并行策略：進一步降低氣泡率，提升并行訓(xùn)練效率；
• 多種 FP8 優(yōu)化策略：獨創(chuàng)的 FP8 優(yōu)化策略，提高訓(xùn)練的性能和穩(wěn)定性；
• 異步 checkpoint 策略：提高訓(xùn)練過程中的容錯能力和效率；
• 優(yōu)化后的重計算策略：減少計算和顯存開銷，提高訓(xùn)練速度；
• 容錯訓(xùn)練策略：獨創(chuàng)的容錯訓(xùn)練算法，增強訓(xùn)練過程中的容錯能力；
• 集成摩爾線程 FlashMLA 和 DeepGemm 庫：進一步釋放摩爾線程 GPU 的算力和 FP8 計算能力，提升計算性能和效率。

IT之家附開源地址如下：

• MT-MegatronLM 開源地址：https://github.com/MooreThreads/MT-MegatronLM
• MT-TransformerEngine 開源地址：https://github.com/MooreThreads/MT-TransformerEngine
• 摩爾線程 Simumax 開源地址：https://github.com/MooreThreads/SimuMax