PPoPP（Principles and Practice of Parallel Programming）是并行與高性能計(jì)算領(lǐng)域的CCF-A類國際會議，關(guān)注并行程序設(shè)計(jì)、系統(tǒng)與運(yùn)行時(shí)等研究方向。第31屆“ACM SIGPLAN并行編程原理與實(shí)踐研討會”（ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, PPoPP '26）將于2026年1月31日至2月4日在澳大利亞悉尼召開，本次會議從280篇投稿論文中接收51篇，接收率為18.2%。PKU-DAIR實(shí)驗(yàn)室論文《Elastic and Efficient Model Partitioning and Checkpointing for Fault-Tolerant Distributed Training》被接收。

Elastor: Elastic and Efficient Model Partitioning and Checkpointing for Fault-Tolerant Distributed Training

作者：Xuanyu Wang，F(xiàn)angcheng Fu，Haoyang Li，Hao Ge，Sheng Lin，Jiawen Niu，Bin Cui

代碼鏈接：https://github.com/PKU-DAIR/Hetu

一、 背景

大模型訓(xùn)練離不開分布式：數(shù)據(jù)并行（DP）負(fù)責(zé)擴(kuò)吞吐，張量并行（TP）/流水并行（PP）負(fù)責(zé)把超大模型拆到多張GPU上。但現(xiàn)實(shí)世界的集群并不“理想”：GPU宕機(jī)、節(jié)點(diǎn)掉線、網(wǎng)絡(luò)故障會讓可用GPU數(shù)量在訓(xùn)練中波動。如果系統(tǒng)只能按“整節(jié)點(diǎn)失敗”去設(shè)計(jì)，一旦出現(xiàn)“部分GPU不可用”，要么浪費(fèi)仍然健康的GPU，要么被迫長時(shí)間停機(jī)等待。

更麻煩的是，訓(xùn)練策略一變（例如從32卡變成28卡、PP stage數(shù)和TP組大小都發(fā)生變化），檢查點(diǎn)也隨之變得難處理。很多框架按“當(dāng)前并行策略切分參數(shù)”來存儲權(quán)重，恢復(fù)時(shí)如果切分方式不同，就會出現(xiàn)冗余讀取與重分片開銷，在共享文件系統(tǒng)（如NAS）上尤其致命——I/O調(diào)用次數(shù)多、單次I/O延遲高，恢復(fù)速度很容易被拖垮。

《Elastor: Elastic and Efficient Model Partitioning and Checkpointing for Fault-Tolerant Distributed Training》聚焦于以上兩個(gè)問題：”當(dāng)GPU/節(jié)點(diǎn)在訓(xùn)練過程中失效、可用GPU數(shù)量發(fā)生變化時(shí)，如何既能快速恢復(fù)訓(xùn)練，又不把時(shí)間浪費(fèi)在反復(fù)的檢查點(diǎn)保存/加載與重分片上”，并提供了創(chuàng)新的協(xié)同設(shè)計(jì)解決方案：一方面讓模型切分足夠彈性，能在任意數(shù)量GPU上恢復(fù)；另一方面讓檢查點(diǎn)足夠“策略無關(guān)”，盡量避免因?yàn)榍蟹肿兓a(chǎn)生重復(fù)I/O，并把周期性檢查點(diǎn)的額外開銷隱藏到訓(xùn)練流水線里。

二、方法

Elastor的核心可以概括為四件事：彈性切分（HMP）、策略搜索、細(xì)粒度檢查點(diǎn)、以及訓(xùn)練-保存的重疊優(yōu)化。

1. 異構(gòu)模型并行（HMP, Heterogeneous Model Parallelism）：當(dāng)某些GPU失效時(shí)，系統(tǒng)仍能用剩余GPU繼續(xù)訓(xùn)練。HMP允許不同DP rank內(nèi)的TP組大小不一致，并在此基礎(chǔ)上組織PP階段與通信組，從而適配“非整除”的GPU數(shù)量。

圖1：異構(gòu)模型并行切分方案

2. 恢復(fù)時(shí)的策略搜索（Strategy Searching）：當(dāng)GPU數(shù)量變化后，Elastor會在候選的{DP, TP_max}組合中搜索合適的并行策略。其流程包含：①把可用GPU劃分成若干TP組并分配給各DP rank；②在每個(gè)DP rank內(nèi)部進(jìn)一步?jīng)Q定層/數(shù)據(jù)如何分配，并通過微批（micro-batch）分配平衡不同rank的計(jì)算。在論文的模擬中，策略搜索在1024張GPU規(guī)模下也能在數(shù)秒內(nèi)完成。

圖2：自適應(yīng)策略搜索示意圖

3. 細(xì)粒度、分片驅(qū)動的檢查點(diǎn)（Fine-grained Checkpointing via Splits）：將參數(shù)張量統(tǒng)一切成全局的global_split份（split），并保證任意HMP策略下每張GPU都持有整數(shù)個(gè)split。這樣恢復(fù)時(shí)每張GPU只需要加載“自己負(fù)責(zé)的split”，避免了因切分變化導(dǎo)致的冗余讀取。同時(shí)，Elastor用JSON元數(shù)據(jù)記錄split與文件位置的映射，做到策略變化下仍能精確定位所需數(shù)據(jù)。

圖3: 細(xì)粒度，自適應(yīng)模型切分

4. 高效保存/加載與重疊（Overlapping Training & Checkpointing）：保存過程被拆成GPU→CPU內(nèi)存與CPU內(nèi)存→文件系統(tǒng)兩段，通過共享內(nèi)存與多進(jìn)程/多線程把參數(shù)搬運(yùn)、序列化（Safetensors）和寫盤解耦，并與訓(xùn)練計(jì)算流并行執(zhí)行，盡量把檢查點(diǎn)成本“藏起來”。加載階段則通過重排與合并I/O，把大量小I/O盡可能合并為更少的順序讀取，降低共享文件系統(tǒng)上的開銷。

圖4: 高效的異步存儲方案和流水線

三、實(shí)驗(yàn)

論文在32張A100-40G的集群上評估Elastor：4臺服務(wù)器每臺8卡，機(jī)內(nèi)NVLink帶寬約600GB/s，機(jī)間InfiniBand帶寬約200GB/s。文件系統(tǒng)使用NAS，單文件寫入帶寬約800MB/s、讀取約1800MB/s，總帶寬超過5TB/s。

工作負(fù)載選擇了3個(gè)LLM：LLaMA2-7B、LLaMA2-13B與Qwen2.5-32B；默認(rèn)上下文長度4096，全局batch size為256。為了貼近真實(shí)環(huán)境，作者根據(jù)集群故障統(tǒng)計(jì)構(gòu)造了5種GPU可用性軌跡（Case A-E），包括單GPU故障、多個(gè)節(jié)點(diǎn)內(nèi)GPU故障、整節(jié)點(diǎn)掉線/斷網(wǎng)、以及混合故障等。

對比基線主要包括：FSDP2+PyTorch Distributed Checkpoint（DCP），以及Megatron配合不同檢查點(diǎn)方案（如MCP與BCP）。實(shí)驗(yàn)從三個(gè)維度評估：訓(xùn)練效率、模型加載效率、以及模型保存效率。

• 訓(xùn)練效率：在無故障（Case A）下，Elastor與強(qiáng)基線訓(xùn)練效率接近；當(dāng)GPU數(shù)量動態(tài)變化（Case B-E）時(shí)，Elastor能更穩(wěn)定地維持MFU，并在端到端訓(xùn)練時(shí)間上取得約1.12×–1.40×的加速。

圖5: 訓(xùn)練效率

• 加載效率：由于檢查點(diǎn)對并行策略更“無關(guān)”，且I/O合并更充分，Elastor在不同故障場景下的加載耗時(shí)顯著降低，整體可達(dá)約1.95×–4.98×的加速。
• 保存效率：通過訓(xùn)練-保存流水線化與線程解耦，模型保存階段也獲得約1.62×–3.94×的提升，降低了周期性檢查點(diǎn)對長期訓(xùn)練的侵蝕。

四、總結(jié)

Elastor把“彈性訓(xùn)練”往前推進(jìn)了一步：不再只假設(shè)整節(jié)點(diǎn)失敗，而是正面面對更常見的部分GPU不可用。它通過HMP讓模型切分能適配任意GPU數(shù)量，又通過細(xì)粒度split把檢查點(diǎn)做成策略無關(guān)，避免了恢復(fù)時(shí)的冗余I/O與重分片；最后再用重疊與I/O合并把檢查點(diǎn)成本壓到更低。

對工程實(shí)踐而言，這篇工作有兩個(gè)啟示：一是故障恢復(fù)能力要與并行策略的動態(tài)變化綁定考慮；二是檢查點(diǎn)格式與加載路徑的設(shè)計(jì)，往往比“寫不寫檢查點(diǎn)”本身更決定系統(tǒng)能否在真實(shí)集群里跑得穩(wěn)、跑得快。

實(shí)驗(yàn)室簡介

北京大學(xué)數(shù)據(jù)與智能實(shí)驗(yàn)室（Data And Intelligence Research Lab at Peking Univeristy，PKU-DAIR實(shí)驗(yàn)室）由北京大學(xué)計(jì)算機(jī)學(xué)院崔斌教授領(lǐng)導(dǎo)，長期從事數(shù)據(jù)庫系統(tǒng)、大數(shù)據(jù)管理與分析、人工智能等領(lǐng)域的前沿研究，在理論和技術(shù)創(chuàng)新以及系統(tǒng)研發(fā)上取得多項(xiàng)成果，已在國際頂級學(xué)術(shù)會議和期刊發(fā)表學(xué)術(shù)論文200余篇，發(fā)布多個(gè)開源項(xiàng)目。課題組同學(xué)曾數(shù)十次獲得包括CCF優(yōu)博、ACM中國優(yōu)博、北大優(yōu)博、微軟學(xué)者、蘋果獎(jiǎng)學(xué)金、谷歌獎(jiǎng)學(xué)金等榮譽(yù)。PKU-DAIR實(shí)驗(yàn)室持續(xù)與工業(yè)界展開卓有成效的合作，與騰訊、阿里巴巴、蘋果、微軟、百度、快手、中興通訊等多家知名企業(yè)開展項(xiàng)目合作和前沿探索，解決實(shí)際問題，進(jìn)行科研成果的轉(zhuǎn)化落地。