“2025人工智能基礎設施峰會”會場
AI技術飛速演進的當下，神經網絡模型的規(guī)模和復雜度不斷攀升，對訓練過程中的效率和容錯能力提出了更高要求。應對這一挑戰(zhàn)，上?？萍即髮W研究員、博導殷樹教授團隊開展了相關研究工作，在面向大規(guī)模神經網絡的檢查點（Checkpointing）方面取得進展。
在2025人工智能基礎設施峰會-智能算力前沿技術論壇，殷樹教授以”面向神經網絡的探索與優(yōu)化”為題，分享其在面向大規(guī)模神經網絡檢查點方面的最新研究成果，介紹名為Portus的優(yōu)化方法，如何通過優(yōu)化的數據傳輸路徑和索引結構，提升DNN檢查點的效率，并為大規(guī)模模型訓練提供高效的容錯解決方案，深入探討如何通過技術創(chuàng)新提升神經網絡訓練的效率和可靠性。
智能算力前沿技術論壇現場
上?？萍即髮W研究員、博導殷樹教授
隨著深度學習模型規(guī)模的持續(xù)增長（例如PaLM模型已達到5400億參數），訓練過程中的容錯需求變得愈發(fā)重要。Checkpointing技術雖然能夠通過定期保存模型狀態(tài)來實現故障恢復，但在當前3D并行訓練框架下面臨著數據量指數增長、存儲效率低下和傳輸開銷過大等核心挑戰(zhàn)?，F有解決方案如CheckFreq需要消耗高達33%的訓練時間用于檢查點操作，主要瓶頸來自多次數據拷貝、內核層交互以及數據序列化帶來的額外開銷。
針對這些問題，研究團隊提出了名為Portus的新型優(yōu)化方案。該系統(tǒng)的核心創(chuàng)新在于設計了反向RDMA通道，通過改變傳統(tǒng)數據傳輸方向，使計算節(jié)點只需注冊GPU內存地址，而PMEM設備可以主動通過RDMA讀取或推送數據，從而徹底消除了GPU到PMEM的數據拷貝開銷。此外，Portus采用雙窗口PMEM管理機制，通過輪換寫入策略優(yōu)化了持久內存的訪問效率。在架構設計上，系統(tǒng)采用輕量級元數據管理，服務器端負責集中維護檢查點元數據，而客戶端則以PyTorch插件形式部署，實現了零拷貝的張量收集功能。
實驗驗證環(huán)節(jié)采用了由雙路Xeon服務器（配備6塊256GB PMEM設備）構成的存儲端，以及搭載V100/A40 GPU的計算節(jié)點，網絡環(huán)境為100Gbps Infiniband。測試結果顯示，Portus在224億參數GPT模型上的表現顯著優(yōu)于傳統(tǒng)方案：完成一次Checkpointing僅需15秒，相比BeeGFS的127秒和torch.save的130秒實現了8倍以上的性能提升；在數據恢復（Restoration）方面更是達到9.23倍的加速效果。進一步分析表明，Portus成功將GPU利用率提升至76.4%，同時將主要開銷集中在RDMA通信環(huán)節(jié)，基本消除了序列化和內存拷貝帶來的性能損耗。據估算，在典型的24小時訓練周期中，Portus可以累計節(jié)省1.5小時的檢查點操作時間。
這項研究的重要意義在于首次實現了PMEM與RDMA的高效協(xié)同，為千億參數規(guī)模的大模型訓練提供了切實可用的高性價比容錯解決方案。展望未來，研究團隊計劃進一步探索將CXL內存等新型硬件技術納入系統(tǒng)優(yōu)化范疇，同時深入研究檢查點機制與訓練流水線的深度協(xié)同優(yōu)化策略。這些工作將持續(xù)推動人工智能基礎設施的技術創(chuàng)新，為更大規(guī)模、更復雜的神經網絡訓練提供可靠支持。
