“2025人工智能基礎(chǔ)設(shè)施峰會”會場

AI技術(shù)飛速演進的當(dāng)下，神經(jīng)網(wǎng)絡(luò)模型的規(guī)模和復(fù)雜度不斷攀升，對訓(xùn)練過程中的效率和容錯能力提出了更高要求。應(yīng)對這一挑戰(zhàn)，上?？萍即髮W(xué)研究員、博導(dǎo)殷樹教授團隊開展了相關(guān)研究工作，在面向大規(guī)模神經(jīng)網(wǎng)絡(luò)的檢查點（Checkpointing）方面取得進展。

在2025人工智能基礎(chǔ)設(shè)施峰會-智能算力前沿技術(shù)論壇，殷樹教授以”面向神經(jīng)網(wǎng)絡(luò)的探索與優(yōu)化”為題，分享其在面向大規(guī)模神經(jīng)網(wǎng)絡(luò)檢查點方面的最新研究成果，介紹名為Portus的優(yōu)化方法，如何通過優(yōu)化的數(shù)據(jù)傳輸路徑和索引結(jié)構(gòu)，提升DNN檢查點的效率，并為大規(guī)模模型訓(xùn)練提供高效的容錯解決方案，深入探討如何通過技術(shù)創(chuàng)新提升神經(jīng)網(wǎng)絡(luò)訓(xùn)練的效率和可靠性。

智能算力前沿技術(shù)論壇現(xiàn)場

上?？萍即髮W(xué)研究員、博導(dǎo)殷樹教授

隨著深度學(xué)習(xí)模型規(guī)模的持續(xù)增長（例如PaLM模型已達到5400億參數(shù)），訓(xùn)練過程中的容錯需求變得愈發(fā)重要。Checkpointing技術(shù)雖然能夠通過定期保存模型狀態(tài)來實現(xiàn)故障恢復(fù)，但在當(dāng)前3D并行訓(xùn)練框架下面臨著數(shù)據(jù)量指數(shù)增長、存儲效率低下和傳輸開銷過大等核心挑戰(zhàn)?，F(xiàn)有解決方案如CheckFreq需要消耗高達33%的訓(xùn)練時間用于檢查點操作，主要瓶頸來自多次數(shù)據(jù)拷貝、內(nèi)核層交互以及數(shù)據(jù)序列化帶來的額外開銷。

針對這些問題，研究團隊提出了名為Portus的新型優(yōu)化方案。該系統(tǒng)的核心創(chuàng)新在于設(shè)計了反向RDMA通道，通過改變傳統(tǒng)數(shù)據(jù)傳輸方向，使計算節(jié)點只需注冊GPU內(nèi)存地址，而PMEM設(shè)備可以主動通過RDMA讀取或推送數(shù)據(jù)，從而徹底消除了GPU到PMEM的數(shù)據(jù)拷貝開銷。此外，Portus采用雙窗口PMEM管理機制，通過輪換寫入策略優(yōu)化了持久內(nèi)存的訪問效率。在架構(gòu)設(shè)計上，系統(tǒng)采用輕量級元數(shù)據(jù)管理，服務(wù)器端負責(zé)集中維護檢查點元數(shù)據(jù)，而客戶端則以PyTorch插件形式部署，實現(xiàn)了零拷貝的張量收集功能。

實驗驗證環(huán)節(jié)采用了由雙路Xeon服務(wù)器（配備6塊256GB PMEM設(shè)備）構(gòu)成的存儲端，以及搭載V100/A40 GPU的計算節(jié)點，網(wǎng)絡(luò)環(huán)境為100Gbps Infiniband。測試結(jié)果顯示，Portus在224億參數(shù)GPT模型上的表現(xiàn)顯著優(yōu)于傳統(tǒng)方案：完成一次Checkpointing僅需15秒，相比BeeGFS的127秒和torch.save的130秒實現(xiàn)了8倍以上的性能提升；在數(shù)據(jù)恢復(fù)（Restoration）方面更是達到9.23倍的加速效果。進一步分析表明，Portus成功將GPU利用率提升至76.4%，同時將主要開銷集中在RDMA通信環(huán)節(jié)，基本消除了序列化和內(nèi)存拷貝帶來的性能損耗。據(jù)估算，在典型的24小時訓(xùn)練周期中，Portus可以累計節(jié)省1.5小時的檢查點操作時間。

這項研究的重要意義在于首次實現(xiàn)了PMEM與RDMA的高效協(xié)同，為千億參數(shù)規(guī)模的大模型訓(xùn)練提供了切實可用的高性價比容錯解決方案。展望未來，研究團隊計劃進一步探索將CXL內(nèi)存等新型硬件技術(shù)納入系統(tǒng)優(yōu)化范疇，同時深入研究檢查點機制與訓(xùn)練流水線的深度協(xié)同優(yōu)化策略。這些工作將持續(xù)推動人工智能基礎(chǔ)設(shè)施的技術(shù)創(chuàng)新，為更大規(guī)模、更復(fù)雜的神經(jīng)網(wǎng)絡(luò)訓(xùn)練提供可靠支持。

xiesc