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

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)練中，輸入序列往往存在顯著的長(zhǎng)度差異。傳統(tǒng)做法是通過padding將較短的序列補(bǔ)齊到批次中的最大長(zhǎng)度，以便統(tǒng)一并行計(jì)算，但這種方式會(huì)引入對(duì)填充token的冗余計(jì)算，造成算力浪費(fèi)。為提升效率，現(xiàn)代訓(xùn)練系統(tǒng)廣泛采用packing技術(shù)，即將多個(gè)不同長(zhǎng)度的序列拼接成一個(gè)較長(zhǎng)序列，再通過block-diagonal attention等機(jī)制保證語(yǔ)義隔離，從而減少無(wú)效計(jì)算。

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

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

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

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

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

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

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

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

方法

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

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

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

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

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

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

實(shí)驗(yàn)

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

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

總  結(jié)

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

實(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è)開源項(xiàng)目。課題組同學(xué)曾數(shù)十次獲得包括CCF優(yōu)博、ACM中國(guó)優(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)化落地。