SIGMOD是計算機科學領域的頂級國際學術會議之一，專注于數據庫系統(tǒng)和數據管理技術的研究。作為中國計算機學會（CCF）推薦的A類會議，SIGMOD與VLDB、ICDE并稱為數據庫領域的三大頂會，近年來論文接收率維持在17%-20%左右。PKU-DAIR實驗室的論文《Hydraulis: Balancing Large Transformer Model Training via Co-designing Parallel Strategies and Data Assignment》在SIGMOD第二輪評審中成功被錄用。

Hydraulis: Balancing Large Transformer Model Training via Co-designing Parallel Strategies and Data Assignment

作者：Haoyang Li, Fangcheng Fu, Sheng Lin, Hao Ge, Xuanyu Wang, Jiawen Niu, Jinbao Xue, Yangyu Tao, Di Wang, Jie Jiang, Bin Cui

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

背景

在大規(guī)模Transformer模型訓練中，輸入序列往往存在顯著的長度差異。傳統(tǒng)做法是通過padding將較短的序列補齊到批次中的最大長度，以便統(tǒng)一并行計算，但這種方式會引入對填充token的冗余計算，造成算力浪費。為提升效率，現代訓練系統(tǒng)廣泛采用packing技術，即將多個不同長度的序列拼接成一個較長序列，再通過block-diagonal attention等機制保證語義隔離，從而減少無效計算。

目前，主流的訓練系統(tǒng)通常采用兩類方法：一方面，并行策略上采用靜態(tài)、同構的配置，即所有訓練流水線（replicas）在整個訓練過程中都保持相同的并行度組合（如固定的tensor parallel與pipeline parallel度數）；另一方面，數據管理上使用max-length packing，即先設定一個最大序列長度（context length），再將不同長度的原始序列packing成不超過該閾值的長序列，并平均分配到各設備上。這種“靜態(tài)同構并行+最大長度packing”的組合在實現上簡潔，但卻存在四種負載不均：

1. 迭代內采樣不均衡（intra-iteration sampling imbalance）：同一批次中，短序列占主導地位，而長序列稀少。這種長度差異使得同構并行策略為了給長序列預留顯存，不得不采取低效的資源分配方案，從而拖累了短序列的訓練效率。

2. 迭代間采樣不均衡（inter-iteration sampling imbalance）：不同批次的序列長度分布差異顯著，最大序列長度波動劇烈。靜態(tài)的并行策略通常按最長序列配置，結果在大多數短序列批次中浪費資源。

圖1：迭代內與迭代間采樣不均衡

     3. 流水線內 packing 不均衡（intra-pipeline packing imbalance）：packing雖然平衡了顯存占用，但注意力機制的計算復雜度隨序列長度平方增長，導致不同packed序列的計算時間差異極大，增加流水線氣泡。

    4. 流水線間 packing 不均衡（inter-pipeline packing imbalance）：當不同流水線處理的packed序列復雜度差異較大時，最慢的流水線決定了整體迭代速度，從而拉低集群效率。

圖2：流水線內與流水線間packing不均衡

這四類不均衡使得現有依賴靜態(tài)同構并行與最大長度packing的方法在實際訓練中效率低下。因此，需要在并行策略與數據管理層面協同優(yōu)化，以適應真實世界中變長序列的特性。

方法

我們提出Hydraulis系統(tǒng)，從并行策略與數據分配兩個維度協同設計，核心方法包括：

   1. 動態(tài)異構并行策略：針對采樣不均衡，Hydraulis支持在不同迭代動態(tài)選擇合適的并行策略；在同一迭代內，不同流水線可采用異構并行配置（如tensor/pipeline/context parallel度數不同），使長短序列各自匹配最優(yōu)方案，從而同時緩解intra-iteration與inter-iteration的不均衡。

2.優(yōu)化—傳播解耦與子圖機制：Hydraulis 通過將傳播階段（前向/反向計算）與優(yōu)化階段（梯度同步與更新）解耦，保持優(yōu)化階段的參數分片固定不變，同時允許傳播階段靈活切換異構并行配置。為進一步解決兩者之間的異構通信問題，我們引入了 pull（拉取參數）和 push（同步梯度）操作?；谧訄D抽象，我們進一步統(tǒng)一表達了不同并行策略下的通信模式，實現了異構并行策略之間的無縫切換。

3. 兩階段序列分配（two-stage sequence assignment）：針對packing不均衡，Hydraulis先在流水線間分配原始序列以保證負載均衡（解決inter-pipeline imbalance），再在流水線內部通過基于整數線性規(guī)劃的優(yōu)化方法進行packing，使不同微批次的執(zhí)行時間更加接近（解決intra-pipeline imbalance）。

4.數據分布感知的策略候選生成：結合數據集的全局序列長度分布，Hydraulis使用動態(tài)規(guī)劃生成一組高質量的候選策略，訓練時快速評估與選擇，避免了迭代中窮舉搜索的開銷。

圖3：Hydraulis系統(tǒng)框架

實驗

在LLaMA2 7B、13B和32B模型上，以CommonCrawl和GitHub兩個大規(guī)模數據集為基準，我們在64張Nvidia A800 GPU的集群中開展實驗。結果顯示，Hydraulis相較于Megatron-LM、DeepSpeed和HotSPa，在端到端性能上提升1.32–2.66倍。進一步實驗表明，在 GPU 數量擴展、最大序列長度增大和batch size調整等不同條件下，Hydraulis均展現出良好的擴展性和穩(wěn)定性。

圖4：端到端實驗對比

總  結

本文首次系統(tǒng)化地提出通過并行策略與數據管理的協同優(yōu)化來應對大模型訓練中的四類不均衡問題。Hydraulis通過動態(tài)異構并行以及合理的序列分配，有效提升了對變長序列訓練的適應性，大幅改善了負載均衡。實驗驗證了其在真實大規(guī)模集群上的顯著優(yōu)勢，為未來大模型訓練在變長數據場景下的系統(tǒng)設計提供了新的方向。

實驗室簡介

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