SIGMOD是計算機(jī)科學(xué)領(lǐng)域的頂級國際學(xué)術(shù)會議之一，專注于數(shù)據(jù)庫系統(tǒng)和數(shù)據(jù)管理技術(shù)的研究。作為中國計算機(jī)學(xué)會（CCF）推薦的A類會議，SIGMOD與VLDB、ICDE并稱為數(shù)據(jù)庫領(lǐng)域的三大頂會，近年來論文接收率維持在17%-20%左右。PKU-DAIR實(shí)驗(yàn)室的論文《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模型訓(xùn)練中，輸入序列往往存在顯著的長度差異。傳統(tǒng)做法是通過padding將較短的序列補(bǔ)齊到批次中的最大長度，以便統(tǒng)一并行計算，但這種方式會引入對填充token的冗余計算，造成算力浪費(fèi)。為提升效率，現(xiàn)代訓(xùn)練系統(tǒng)廣泛采用packing技術(shù)，即將多個不同長度的序列拼接成一個較長序列，再通過block-diagonal attention等機(jī)制保證語義隔離，從而減少無效計算。

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

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

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

圖1：迭代內(nèi)與迭代間采樣不均衡

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

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

圖2：流水線內(nèi)與流水線間packing不均衡

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

方法

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

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

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

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

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

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

實(shí)驗(yàn)

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

圖4：端到端實(shí)驗(yàn)對比

總  結(jié)

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

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

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