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)領域最具影響力的會議之一，由 USENIX 主辦，關注操作系統(tǒng)、分布式系統(tǒng)及系統(tǒng)與機器學習交叉等方向。OSDI長期被CCF評為A類會議，論文接收率通常在15%左右，與SOSP并稱為計算機系統(tǒng)領域的奧斯卡獎。

PKU-DAIR實驗室論文《HSPMD: Hierarchical and Heterogeneous SPMD for Distributed DL Training》被OSDI 2026接收。

一、背景

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

表1. 不同異構負載場景的空間/時間特征，以及代表工作

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

圖1. 不同訓練框架的范式對比

二、方法

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

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

圖2. 層次化、異構的切分注解

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

圖3. 層次化通信解析

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

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

三、實驗

我們在16張H800和32張H20 GPU組成的異構集群上評估HSPMD，采用Llama系列模型（7B至32B），從三個代表性場景展開。

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

圖5. 異構設備實驗

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

圖6. 不穩(wěn)定設備實驗

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

圖7. 變長數(shù)據(jù)實驗

四．總 結

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

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

實驗室簡介

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