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

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

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

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

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

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

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

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

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

xiesc