HSPMD: Hierarchical and Heterogeneous SPMD for Distributed DL Training

作者：Haoyang Li, Fangcheng Fu, Hao Ge, Sheng Lin, Xuanyu Wang, Jiawen Niu, Yuming Zhou, Xupeng Miao, Bin Cui

OSDI（Operating Systems Design and Implementation）是系統(tǒng)領(lǐng)域最具影響力的會(huì)議之一，由 USENIX 主辦，關(guān)注操作系統(tǒng)、分布式系統(tǒng)及系統(tǒng)與機(jī)器學(xué)習(xí)交叉等方向。OSDI長(zhǎng)期被CCF評(píng)為A類(lèi)會(huì)議，論文接收率通常在15%左右，與SOSP并稱(chēng)為計(jì)算機(jī)系統(tǒng)領(lǐng)域的奧斯卡獎(jiǎng)。

PKU-DAIR實(shí)驗(yàn)室論文《HSPMD: Hierarchical and Heterogeneous SPMD for Distributed DL Training》被OSDI 2026接收。

一、背景

隨著GPT、Gemini等大模型規(guī)模不斷擴(kuò)大，分布式訓(xùn)練已成為基礎(chǔ)設(shè)施核心。其中，單程序多數(shù)據(jù)，即SPMD（Single-Program Multiple-Data）的范式為分布式深度學(xué)習(xí)訓(xùn)練中的多種并行維度提供了統(tǒng)一抽象。用戶只需寫(xiě)一份程序，通過(guò)張量/算子標(biāo)注即可自動(dòng)映射到數(shù)據(jù)并行（DP）、張量并行（TP）、流水并行（PP）等多種策略。然而，SPMD依賴一個(gè)關(guān)鍵前提：訓(xùn)練負(fù)載是均勻劃分的，即設(shè)備同構(gòu)且數(shù)據(jù)負(fù)載一致。但在真實(shí)環(huán)境中，這一假設(shè)逐漸失效。一方面，GPU代際混用、性能差異、甚至頻繁故障，使設(shè)備側(cè)呈現(xiàn)顯著異構(gòu)性；另一方面，原始數(shù)據(jù)（如文本、圖像、視頻）具有天然不均勻性，導(dǎo)致計(jì)算負(fù)載動(dòng)態(tài)變化。這些因素共同打破了SPMD的對(duì)稱(chēng)性假設(shè)。

表1. 不同異構(gòu)負(fù)載場(chǎng)景的空間/時(shí)間特征，以及代表工作

      現(xiàn)有方法要么采用MPMD（多程序多數(shù)據(jù)）方案，但復(fù)雜且難擴(kuò)展；要么在SPMD基礎(chǔ)上引入定制調(diào)度器以支持特定場(chǎng)景下的非對(duì)稱(chēng)執(zhí)行，然而，這類(lèi)方法通常強(qiáng)綁定具體場(chǎng)景，缺乏通用性。為此，我們提出了HSPMD，其從更基礎(chǔ)的抽象出發(fā)，在原語(yǔ)層擴(kuò)展SPMD，使其原生支持異構(gòu)執(zhí)行，并提出面向“空間/時(shí)間異構(gòu)”的統(tǒng)一建模與機(jī)制設(shè)計(jì)。

圖1. 不同訓(xùn)練框架的范式對(duì)比

二、方法

HSPMD的核心思路是在SPMD的底層原語(yǔ)層面引入不對(duì)稱(chēng)性，而非像現(xiàn)有工作那樣在調(diào)度器層面打補(bǔ)丁。其設(shè)計(jì)圍繞三個(gè)關(guān)鍵創(chuàng)新展開(kāi)：

1. 層次化、異構(gòu)的切分注解：傳統(tǒng)SPMD注解只能對(duì)張量進(jìn)行均勻切分。HSPMD擴(kuò)展了注解體系，引入DG Union（設(shè)備組聯(lián)合）和DS Union（分布式狀態(tài)聯(lián)合），允許一個(gè)張量在不同設(shè)備子組內(nèi)采用不同的切分方式。同時(shí)引入HDim（異構(gòu)維度）和HSize（異構(gòu)大小），描述跨子組的切分關(guān)系，從而在單設(shè)備編程視角下原生支持不對(duì)稱(chēng)并行。

圖2. 層次化、異構(gòu)的切分注解

2. 層次化通信解析：根據(jù)切分注解的層次結(jié)構(gòu)，HSPMD將通信分為底層通信（子組內(nèi)部）和頂層通信（跨子組）。底層通信盡可能復(fù)用標(biāo)準(zhǔn)SPMD的集體通信原語(yǔ)（如all-reduce、all-gather），利用同質(zhì)設(shè)備間的高帶寬鏈路；頂層通信則處理異構(gòu)鏈路，設(shè)計(jì)了SplitAR、SplitAG等算子，并在復(fù)雜場(chǎng)景下使用Batched-Send-Receive（BSR）機(jī)制。BSR通過(guò)識(shí)別最細(xì)粒度切片、構(gòu)建映射表、基于帶寬優(yōu)先和負(fù)載均衡的啟發(fā)式調(diào)度，實(shí)現(xiàn)高效的非對(duì)稱(chēng)通信。

圖3. 層次化通信解析

3. 圖特化與圖切換：針對(duì)空間異構(gòu)性（負(fù)載靜態(tài)不平衡），HSPMD引入圖特化：從單一定義圖和注解計(jì)劃出發(fā)，為每個(gè)設(shè)備生成專(zhuān)屬的可執(zhí)行圖，設(shè)備間可執(zhí)行不同邏輯。針對(duì)時(shí)間異構(gòu)性（負(fù)載動(dòng)態(tài)變化），HSPMD引入圖切換：當(dāng)訓(xùn)練策略需重配時(shí)（如設(shè)備故障、數(shù)據(jù)分布變化），通過(guò)BSR在線重分片模型權(quán)重，無(wú)縫切換到新策略，無(wú)需重啟或重新加載檢查點(diǎn)。

圖4. HSPMD工作流：圖特化與圖切換作為模塊化構(gòu)件，靈活組合以適配異構(gòu)設(shè)備、不穩(wěn)定設(shè)備、變長(zhǎng)數(shù)據(jù)三類(lèi)場(chǎng)景

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

我們?cè)?6張H800和32張H20 GPU組成的異構(gòu)集群上評(píng)估HSPMD，采用Llama系列模型（7B至32B），從三個(gè)代表性場(chǎng)景展開(kāi)。

1. 異構(gòu)設(shè)備：在混合GPU類(lèi)型（H800與H20混用）的靜態(tài)異構(gòu)環(huán)境下，HSPMD相較于標(biāo)準(zhǔn)SPMD系統(tǒng)（DeepSpeed、Megatron）和異構(gòu)專(zhuān)用系統(tǒng)（HexiScale），訓(xùn)練吞吐提升顯著。標(biāo)準(zhǔn)SPMD因?qū)ΨQ(chēng)切分導(dǎo)致設(shè)備間負(fù)載失衡，而HexiScale的調(diào)度器難以支持靈活流水線（如1F1B），且僅支持粗粒度廣播。HSPMD通過(guò)聲明式注解解耦策略與執(zhí)行，配合層次化通信，實(shí)現(xiàn)更優(yōu)的工作負(fù)載均衡與通信效率。

圖5. 異構(gòu)設(shè)備實(shí)驗(yàn)

2. 不穩(wěn)定設(shè)備：我們模擬GPU故障場(chǎng)景，對(duì)比彈性訓(xùn)練性能。標(biāo)準(zhǔn)SPMD只能丟棄整節(jié)點(diǎn)，且依賴檢查點(diǎn)重啟，恢復(fù)開(kāi)銷(xiāo)大。專(zhuān)用系統(tǒng)Oobleck支持無(wú)重啟重配，但策略空間受限（只能使用固定的pipeline templates）且重配通信粒度粗（只能使用broadcast）。HSPMD可利用剩余所有GPU，通過(guò)張量級(jí)聲明注解探索更優(yōu)策略，且重配時(shí)通過(guò)Fused BSR合并通信、均衡負(fù)載，顯著降低切換開(kāi)銷(xiāo)。

圖6. 不穩(wěn)定設(shè)備實(shí)驗(yàn)

3. 變長(zhǎng)數(shù)據(jù)：我們訓(xùn)練32B模型，處理序列長(zhǎng)度動(dòng)態(tài)變化的混合數(shù)據(jù)。標(biāo)準(zhǔn)SPMD固定策略在短序列為主時(shí)效率低下；HotSPA雖支持策略熱切換，但受限于SPMD對(duì)稱(chēng)性，無(wú)法使用空間異構(gòu)策略。HSPMD預(yù)先生成多種異構(gòu)策略，根據(jù)序列長(zhǎng)度分布動(dòng)態(tài)切換，在性能和靈活性上均優(yōu)于基線。

圖7. 變長(zhǎng)數(shù)據(jù)實(shí)驗(yàn)

四．總 結(jié)

HSPMD把“SPMD范式”往前推了一步：不再局限于對(duì)稱(chēng)的、靜態(tài)的并行假設(shè)，而是正面面對(duì)設(shè)備異構(gòu)、設(shè)備故障、數(shù)據(jù)分布不均等真實(shí)世界挑戰(zhàn)。它通過(guò)在原語(yǔ)層面擴(kuò)展注解和通信，讓SPMD自身具備表達(dá)和執(zhí)行異構(gòu)并行策略的能力，而非依賴調(diào)度器層面打補(bǔ)丁。

對(duì)系統(tǒng)設(shè)計(jì)而言，這篇工作帶來(lái)兩個(gè)啟示：一是處理異構(gòu)性應(yīng)當(dāng)從底層原語(yǔ)入手，而非在高層堆砌特設(shè)機(jī)制；二是空間異構(gòu)性與時(shí)間異構(gòu)性可以抽象為圖特化與圖切換兩個(gè)正交的模塊化構(gòu)件，組合起來(lái)即可覆蓋多種場(chǎng)景。這種“底層擴(kuò)展+模塊化抽象”的思路，為構(gòu)建通用、高效的分布式訓(xùn)練系統(tǒng)提供了新的方向。

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