IT之家 3 月 17 日消息,摩爾線程官方今日發(fā)文宣布,已正式開源 MT-MegatronLM 與 MT-TransformerEngine 兩大 AI 框架。這兩大框架通過深度融合 FP8 混合訓(xùn)練策略和高性能算子庫,在國產(chǎn)全功能 GPU 上實(shí)現(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 高密度計(jì)算的潛力和 memory bound 算子效率。兩大框架的技術(shù)突破體現(xiàn)在硬件適配與算法創(chuàng)新的深度協(xié)同:
混合并行訓(xùn)練:支持 Dense、多模態(tài)及 MoE 模型的混合并行訓(xùn)練,可靈活應(yīng)對(duì)不同模型架構(gòu)的復(fù)雜運(yùn)算場(chǎng)景;
FP8 混合訓(xùn)練策略:結(jié)合摩爾線程 GPU 原生支持的 FP8 混合精度訓(xùn)練策略,能夠有效提升訓(xùn)練效率;
高性能算子庫:通過高性能算子庫 muDNN 與通信庫 MCCL 的深度集成,系統(tǒng)性優(yōu)化了計(jì)算密集型任務(wù)與多卡協(xié)同的通信開銷;同時(shí)結(jié)合摩爾線程開源 Simumax 庫,可自動(dòng)進(jìn)行并行策略搜索,并針對(duì)不同模型和加速環(huán)境 spec 最大化并行訓(xùn)練性能;
異常訓(xùn)練處理:框架內(nèi)置的 rewind 異常恢復(fù)機(jī)制,可自動(dòng)回滾至最近穩(wěn)定節(jié)點(diǎn)繼續(xù)訓(xùn)練,大幅提升大規(guī)模訓(xùn)練的穩(wěn)定性;
完整的兼容性:兩個(gè)框架兼容 GPU 主流生態(tài),既保障了現(xiàn)有生態(tài)的平滑遷移,也為開發(fā)者構(gòu)建自有的 AI 技術(shù)棧提供了底層支撐。
實(shí)際應(yīng)用效果如下:
高效訓(xùn)練:在全功能 GPU 集群上,Llama3 8B 模型的訓(xùn)練任務(wù),可以利用 FP8 在 loss 幾乎無損的情況下 MFU 達(dá)到 90% 以上;(如下圖所示)
▲ 利用摩爾線程 FP8 混合精度加速技術(shù)在 loss 無損的情況下得到 28% 的加速
復(fù)現(xiàn) DeepSeek 滿血版訓(xùn)練:摩爾線程已深度集成并開源對(duì) DeepSeek 并行算法 DualPipe 的高效支持,MT-DualPipe 可以完整接入 MT-Megatron 框架和 MT-TransformerEngine 框架,成功實(shí)現(xiàn) DeepSeek V3 訓(xùn)練流程的完整復(fù)現(xiàn),支持 MLA、MTP 及多種專家平衡策略;
性能大幅優(yōu)化:通過多種 Transformer 算子融合技術(shù),顯著提升了內(nèi)存帶寬利用率,有效緩解 memory bound 瓶頸,進(jìn)一步釋放國產(chǎn) GPU 的硬件潛力。
摩爾線程官方表示將持續(xù)優(yōu)化 MT-MegatronLM 與 MT-TransformerEngine 框架,并引入系列功能,具體如下:
Dual Pipe / ZeroBubble 并行策略:進(jìn)一步降低氣泡率,提升并行訓(xùn)練效率;
多種 FP8 優(yōu)化策略:獨(dú)創(chuàng)的 FP8 優(yōu)化策略,提高訓(xùn)練的性能和穩(wěn)定性;
異步 checkpoint 策略:提高訓(xùn)練過程中的容錯(cuò)能力和效率;
優(yōu)化后的重計(jì)算策略:減少計(jì)算和顯存開銷,提高訓(xùn)練速度;
容錯(cuò)訓(xùn)練策略:獨(dú)創(chuàng)的容錯(cuò)訓(xùn)練算法,增強(qiáng)訓(xùn)練過程中的容錯(cuò)能力;
集成摩爾線程 FlashMLA 和 DeepGemm 庫:進(jìn)一步釋放摩爾線程 GPU 的算力和 FP8 計(jì)算能力,提升計(jì)算性能和效率。
IT之家附開源地址如下:
MT-MegatronLM 開源地址:https://github.com/MooreThreads/MT-MegatronLM
MT-TransformerEngine 開源地址:https://github.com/MooreThreads/MT-TransformerEngine
摩爾線程 Simumax 開源地址:https://github.com/MooreThreads/SimuMax
廣告聲明:文內(nèi)含有的對(duì)外跳轉(zhuǎn)鏈接(包括不限于超鏈接、二維碼、口令等形式),用于傳遞更多信息,節(jié)省甄選時(shí)間,結(jié)果僅供參考,IT之家所有文章均包含本聲明。