上海科技大學(xué)計(jì)算機(jī)科學(xué)與技術(shù)學(xué)院殷樹教授團(tuán)隊(duì)在峰會的智能算力前沿技術(shù)論壇中，分享了針對大規(guī)模神經(jīng)網(wǎng)絡(luò)訓(xùn)練的Checkpointing優(yōu)化研究成果。該研究針對當(dāng)前3D并行訓(xùn)練框架下數(shù)據(jù)量激增、存儲效率低、傳輸開銷大的核心痛點(diǎn)，提出了名為Portus的創(chuàng)新型優(yōu)化方案，為千億參數(shù)級大模型訓(xùn)練提供了高性價比的容錯解決方案。

文字編輯｜宋雨涵

1

技術(shù)突破

Portus的新型優(yōu)化方案

隨著深度學(xué)習(xí)模型規(guī)模的持續(xù)增長（例如PaLM模型已達(dá)到5400億參數(shù)），訓(xùn)練過程中的容錯需求變得愈發(fā)重要。

01

Checkpointing技術(shù)面臨的問題？

Checkpointing技術(shù)雖然能夠通過定期保存模型狀態(tài)來實(shí)現(xiàn)故障恢復(fù)，但在當(dāng)前3D并行訓(xùn)練框架下面臨著數(shù)據(jù)量指數(shù)增長、存儲效率低下和傳輸開銷過大等核心挑戰(zhàn)。現(xiàn)有解決方案如CheckFreq需要消耗高達(dá)33%的訓(xùn)練時間用于檢查點(diǎn)操作，主要瓶頸來自多次數(shù)據(jù)拷貝、內(nèi)核層交互以及數(shù)據(jù)序列化帶來的額外開銷。

02

創(chuàng)新設(shè)計(jì)

• 反向RDMA通道設(shè)計(jì)

該系統(tǒng)的核心創(chuàng)新在于設(shè)計(jì)了反向RDMA通道，通過改變傳統(tǒng)數(shù)據(jù)傳輸方向，使計(jì)算節(jié)點(diǎn)只需注冊GPU內(nèi)存地址，而PMEM設(shè)備可以主動通過RDMA讀取或推送數(shù)據(jù)，從而徹底消除了GPU到PMEM的數(shù)據(jù)拷貝開銷。

• 雙窗口PMEM管理機(jī)制

此外，Portus采用雙窗口PMEM管理機(jī)制，通過輪換寫入策略優(yōu)化了持久內(nèi)存的訪問效率。在架構(gòu)設(shè)計(jì)上，系統(tǒng)采用輕量級元數(shù)據(jù)管理，服務(wù)器端負(fù)責(zé)集中維護(hù)檢查點(diǎn)元數(shù)據(jù)，而客戶端則以PyTorch插件形式部署，實(shí)現(xiàn)了零拷貝的張量收集功能。

2

實(shí)驗(yàn)驗(yàn)證環(huán)節(jié)

顯著能效提升

研究團(tuán)隊(duì)在由雙路Xeon服務(wù)器（配備6塊256GB PMEM設(shè)備）和V100/A40 GPU計(jì)算節(jié)點(diǎn)組成的集群中，通過100Gbps Infiniband網(wǎng)絡(luò)完成實(shí)驗(yàn)驗(yàn)證：

效率提升

在224億參數(shù)GPT模型訓(xùn)練中，Portus完成一次Checkpointing僅需15秒，較BeeGFS（127秒）和torch.save（130秒）提升8倍以上。

數(shù)據(jù)恢復(fù)（Restoration）速度達(dá)到9.23倍加速，顯著縮短模型故障恢復(fù)時間。

資源優(yōu)化

Portus成功將GPU利用率提升至76.4%，同時將主要開銷集中在RDMA通信環(huán)節(jié)，基本消除了序列化和內(nèi)存拷貝帶來的性能損耗。據(jù)估算，在典型的24小時訓(xùn)練周期中，Portus可以累計(jì)節(jié)省1.5小時的檢查點(diǎn)操作時間。

技術(shù)意義與未來方向

這項(xiàng)研究的重要意義在于首次實(shí)現(xiàn)了PMEM與RDMA的高效協(xié)同，為千億參數(shù)規(guī)模的大模型訓(xùn)練提供了切實(shí)可用的高性價比容錯解決方案

展望未來，研究團(tuán)隊(duì)計(jì)劃進(jìn)一步探索將CXL內(nèi)存等新型硬件技術(shù)納入系統(tǒng)優(yōu)化范疇，同時深入研究檢查點(diǎn)機(jī)制與訓(xùn)練流水線的深度協(xié)同優(yōu)化策略。這些工作將持續(xù)推動人工智能基礎(chǔ)設(shè)施的技術(shù)創(chuàng)新，為更大規(guī)模、更復(fù)雜的神經(jīng)網(wǎng)絡(luò)訓(xùn)練提供可靠支持。

人工智能基礎(chǔ)設(shè)施加速向“效率與安全并重”演進(jìn)

從算法優(yōu)化到硬件適配，從單點(diǎn)突破到生態(tài)共建，人工智能基礎(chǔ)設(shè)施的創(chuàng)新正加速向“效率與安全并重”演進(jìn)。Portus系統(tǒng)的落地，為我國在AI基礎(chǔ)設(shè)施技術(shù)攻關(guān)中邁進(jìn)一大步，為全球AI大模型訓(xùn)練提供了中國方案。

lixiangjing

算力豹主編