1719
2024-11-13 21:07:29（已編輯）
第30屆“ACM操作系統(tǒng)原理大會(huì)”（SOSP: ACM Symposium on Operating Systems Principles）于2024年11月4日至6日在美國的德克薩斯州召開,。SOSP與OSDI并稱為計(jì)算機(jī)系統(tǒng)領(lǐng)域兩個(gè)最高水平的學(xué)術(shù)會(huì)議，擁有50多年的悠久歷史,。本次SOSP大會(huì)共有248篇論文投稿,，43篇被接收，錄用率僅為17.3%,。
PKU-DAIR實(shí)驗(yàn)室論文《Enabling Parallelism Hot Switching for Efficient Training of Large Language Models》被計(jì)算機(jī)系統(tǒng)領(lǐng)域頂級會(huì)議SOSP 2024錄用,，論文介紹了新型的支持并行熱切換的大模型訓(xùn)練系統(tǒng)，系PKU-DAIR實(shí)驗(yàn)室自研分布式深度學(xué)習(xí)系統(tǒng)河圖Hetu（https://github.com/PKU-DAIR/Hetu）圍繞大模型訓(xùn)練的新成果,。
論文介紹
Distributed Training
Large Language Model
Parallelism Strategy
近年來，以ChatGPT為代表的大語言模型（Large Language Model, LLM）引起了廣泛的關(guān)注,，它的性能提升得益于模型參數(shù)量,、上下文和數(shù)據(jù)量的增長，同時(shí)也給系統(tǒng)優(yōu)化帶來了更多的挑戰(zhàn)?，F(xiàn)有系統(tǒng)通常假設(shè)工作負(fù)載是恒定的,，從而采用靜態(tài)的并行策略組合來進(jìn)行大規(guī)模的分布式訓(xùn)練。然而,，真實(shí)數(shù)據(jù)集的序列長度在不同樣本間差異較大,，且大多呈長尾分布，本工作首次揭示了在這類長短文混訓(xùn)的動(dòng)態(tài)場景中,，現(xiàn)有系統(tǒng)的靜態(tài)并行策略會(huì)大大拖慢短序列的訓(xùn)練效率,。
針對該問題，我們基于自研的分布式深度學(xué)習(xí)系統(tǒng)河圖Hetu,，創(chuàng)新性地提出了首個(gè)支持并行熱切換的HotSPa系統(tǒng),，通過對mini-batch內(nèi)的序列進(jìn)行分組并使用不同的并行策略來最大化訓(xùn)練效率。HotSPa利用熱切換技術(shù)完成模型權(quán)重,、梯度在策略間的高效轉(zhuǎn)換和累積,，在保證精度無損的前提下，最大化利用硬件內(nèi)存和計(jì)算資源,。實(shí)驗(yàn)結(jié)果表明,，與Megatron-LM、DeepSpeed等使用靜態(tài)并行策略的系統(tǒng)相比,，HotSPa在不同規(guī)模的LLaMA2模型和不同上下文長度下,，可以獲得2.99x的加速比,。
近年來,，大規(guī)模預(yù)訓(xùn)練模型得到了快速的發(fā)展,，它的性能提升得益于模型參數(shù)量、上下文和數(shù)據(jù)量的增長,，同時(shí)也給系統(tǒng)優(yōu)化帶來了更多的挑戰(zhàn)?，F(xiàn)有的分布式訓(xùn)練系統(tǒng)提出了一系列并行策略，從而能夠在多個(gè)設(shè)備中處理大規(guī)模的模型和數(shù)據(jù),。并行策略的選擇取決于如顯存占用,、計(jì)算開銷、通信代價(jià)等工作負(fù)載,，現(xiàn)有的方法通常假設(shè)工作負(fù)載是恒定的,，因此在訓(xùn)練過程中會(huì)使用靜態(tài)的并行策略組合。然而,，在真實(shí)數(shù)據(jù)集中,，不同樣本間的序列長度差異較大，導(dǎo)致了樣本間工作負(fù)載的不均衡,，即使在一個(gè)mini-batch內(nèi)部,，這樣的現(xiàn)象也是相當(dāng)顯著的，因此現(xiàn)有系統(tǒng)的靜態(tài)并行策略并不是最優(yōu)解,。
為了支持較長序列的訓(xùn)練,，通常需要采用更節(jié)省內(nèi)存的并行策略以避免內(nèi)存溢出。然而,，對于短序列來說這類并行策略會(huì)引入大量不必要的通信開銷,，導(dǎo)致效率低下。結(jié)合CommonCrawl和GitHub這兩個(gè)數(shù)據(jù)集的序列長度分布可以發(fā)現(xiàn),，雖然數(shù)據(jù)集中都包含長序列,，但占比不高，以張量并行為例,，可以看到隨著上下文的增加,，需要增大TP以避免OOM，但與此同時(shí),，對于在數(shù)據(jù)集中占比大多數(shù)的短序列來說,，更高的TP意味著更低的吞吐,，從而拉低了整體的訓(xùn)練性能,。
總的來說，現(xiàn)有系統(tǒng)忽視了大模型訓(xùn)練中樣本間工作負(fù)載不均衡的問題,，只是簡單地使用靜態(tài)并行策略來處理,。因此，我們嘗試從一個(gè)尚未被探索的方向來加速模型訓(xùn)練過程：我們能否針對不同工作負(fù)載/序列長度的序列采用不同的并行策略？
根據(jù)上述發(fā)現(xiàn),，我們提出了HotSPa，一個(gè)支持并行熱切換的訓(xùn)練系統(tǒng),，核心貢獻(xiàn)如下：
（1）首次提出并行熱切換的訓(xùn)練范式：我們提出了一個(gè)基于并行熱切換的全新訓(xùn)練范式,，對每個(gè)mini-batch內(nèi)的數(shù)據(jù)，我們會(huì)根據(jù)其工作負(fù)載的差異進(jìn)行分組,，并對每個(gè)分組使用最合適的并行策略,，在任意兩組策略之間，我們的系統(tǒng)都會(huì)對模型參數(shù)和梯度進(jìn)行自動(dòng),、高效且無感的切換,，并在模型更新前完成不同策略間的梯度累積，以保證訓(xùn)練效果不受影響,。
（2）統(tǒng)一的計(jì)算圖表示和編譯：現(xiàn)有的系統(tǒng),，如Megatron-LM和DeepSpeed，由于其復(fù)雜的系統(tǒng)設(shè)計(jì),，在訓(xùn)練過程中僅支持一種固定的并行策略組合,，從而無法對不同的序列負(fù)載使用不同的并行策略。而我們的工作設(shè)計(jì)了專門的圖編譯器（graph compiler）,，能夠用一張統(tǒng)一的邏輯圖同時(shí)表示多組不同的并行策略組合,，并進(jìn)一步編譯生成對應(yīng)的多組可執(zhí)行計(jì)算圖，共享模型狀態(tài)的存儲(chǔ),，從而才能支持復(fù)雜的并行熱切換語義,。
（3）并行熱切換技術(shù)：給定任意兩組不同的并行策略組合，它們之間的熱切換需要在不同的設(shè)備中交換模型的參數(shù)和梯度,，不可避免地會(huì)引入額外的通信開銷,。為了解決這個(gè)問題，我們設(shè)計(jì)了熱切換規(guī)劃器（hot switch planner）,，提出一種啟發(fā)式算法來尋找任意兩組策略之間的最優(yōu)通信方案,，并引入了一系列通信和顯存拷貝的優(yōu)化技術(shù)來進(jìn)一步降低切換開銷。
（4）HotSPa系統(tǒng)基于Graph Compiler和Hot Switch Planner,，通過支持并行熱切換的訓(xùn)練范式,，與現(xiàn)有系統(tǒng)相比，可以獲得2.99x的加速比,。
圖編譯器（graph compiler）支持用一張統(tǒng)一的邏輯圖同時(shí)表示多組不同的并行策略組合,，并進(jìn)一步編譯生成對應(yīng)的多組可執(zhí)行計(jì)算圖，且共享模型狀態(tài)的存儲(chǔ),。對應(yīng)三個(gè)核心步驟：
（1）邏輯圖（Logic Graph）：用統(tǒng)一的邏輯圖來表示多組不同的并行策略組合,。
基于DTensor（Distributed Tensor, 分布式張量）架構(gòu),，用DStates（Distributed States, 分布式狀態(tài)）來表示一組并行策略組合：用Splits, Partial, Duplicate來表示參數(shù)在不同設(shè)備中的切分方式，用DeviceGroup來表示不同的參數(shù)切片和設(shè)備的映射關(guān)系,。
用DStates表示并行策略
為了同時(shí)表示多組分布式策略,，本工作創(chuàng)新性地將一個(gè)DTensor與多組DStates綁定，并支持在整張計(jì)算圖上同時(shí)進(jìn)行多組分布式狀態(tài)的推導(dǎo),。具體來說,，令每個(gè)參數(shù)和輸入變量都同時(shí)綁定多組DStates，并在構(gòu)建邏輯計(jì)算圖的過程中,，每個(gè)算子都會(huì)同時(shí)對這多組DStates進(jìn)行推導(dǎo),，并自動(dòng)插入中間算子以保證功能的完整性。
由于不同的并行策略組合,，通過推導(dǎo)DStates所得到的計(jì)算圖并不一定相同（如下圖step1）,。為了能夠用同一張計(jì)算圖表示這多組策略，會(huì)通過自動(dòng)插入空算子（dummy op）來處理不同策略在通信等中間算子上的差異（如下圖step2）,。
通過上述方案,，本工作支持了用統(tǒng)一的邏輯圖來表示多組不同的并行策略組合。
step1：兩組不同的并行策略,，對應(yīng)兩組不同的計(jì)算圖
step2：通過引入dummy op,，使得兩張計(jì)算圖的表現(xiàn)形式一致
step3: 將兩張計(jì)算圖用同一張邏輯圖來統(tǒng)一表示
（2）可執(zhí)行圖（Exec Graphs）：基于統(tǒng)一的邏輯圖編譯、生成多組可執(zhí)行計(jì)算圖,，每組可執(zhí)行計(jì)算圖都對應(yīng)一組獨(dú)立的并行策略組合,，它們之間共享同一份模型狀態(tài)存儲(chǔ)。
邏輯圖只是一種抽象表示,，為了能夠編譯,、生成真正可執(zhí)行的分布式計(jì)算圖，編譯器會(huì)進(jìn)行算子的插入,、合并,、剪枝、替換,、重排等操作：
下圖給了將一個(gè)邏輯圖轉(zhuǎn)化為兩組可執(zhí)行圖（兩組并行策略組合）的具體例子：
邏輯圖（左上角）：編譯、生成兩組可執(zhí)行計(jì)算圖（左下角,，右側(cè)）
（3）編排可執(zhí)行圖（Orchestrate Exec Graphs）：基于Hot Switch Planner給出的代價(jià)分析,，編排這些可執(zhí)行圖（并行策略組合）的執(zhí)行順序，以最小化存儲(chǔ)和通信的開銷,。
選取初始化圖（instantiation graph）：選取最小化模型狀態(tài)（model states）存儲(chǔ)的策略對應(yīng)的可執(zhí)行圖作為初始化圖,。（以下圖為例，選取TP2,PP2為初始化圖）
重排可執(zhí)行圖（executable graph）：基于Hot Switch Planner給出的任意兩組策略間熱切換的代價(jià),，確定不同策略執(zhí)行的先后順序,，以最小化整體的熱切換代價(jià)。（以下圖為例,，tp2pp2->dp4->dp2tp2,，只有第一次熱切換需要引入額外的通信代價(jià)，其余只需要本地切分即可）
剪枝&梯度累積：除初始化圖之外,，其余可執(zhí)行圖需要剪枝不必要的類型轉(zhuǎn)換算子（type casting op）,、參數(shù)更新算子（update op）等，從而保證只在初始化圖上做模型的更新,，在其他可執(zhí)行圖上僅作梯度累積,，確保模型的精度不受影響。
選取tp2,pp2為初始化圖,，重排可執(zhí)行圖順序?yàn)椋簍p2,pp2->dp4->dp2,tp2
HotSPa支持將一個(gè)mini-batch內(nèi)的數(shù)據(jù)按照序列長度進(jìn)行分組,，每組采用不同的并行策略（對應(yīng)不同的可執(zhí)行計(jì)算圖），不同策略的梯度會(huì)累積到同一個(gè)梯度緩沖區(qū)（grad buffer）中以保證模型收斂不受影響,。
在兩組不同的并行策略之間,，HotSPa會(huì)自動(dòng)對模型的權(quán)重和梯度進(jìn)行熱切換，而熱切換規(guī)劃器（Hot Switch Planner）的核心作用就是推導(dǎo)出任意兩組不同的并行策略之間切換代價(jià)最小的通信方案,，具體來說,，分為以下2個(gè)核心步驟：
（1）基于啟發(fā)式算法建模熱切換（Model Hot Switching）
通信方案存在多個(gè)可行解：熱切換指分布式狀態(tài)的切換（從初始策略->目標(biāo)策略），需要在整個(gè)集群中對模型的參數(shù)和梯度進(jìn)行重新劃分,，是一個(gè)多對多的復(fù)雜通信,，由于數(shù)據(jù)并行的存在，同一個(gè)模型切片在不同設(shè)備中存在多份重復(fù)的拷貝,，對應(yīng)的發(fā)送方并不唯一,，因此對于通信方案來說存在大量的可行解。
通信的基本單位-模型切片：對于任意一個(gè)模型參數(shù)或梯度來說,，整體可以看作是一個(gè)全局的抽象的ParamBlock,。每個(gè)ParamBlock會(huì)根據(jù)并行策略所對應(yīng)的分布式狀態(tài)被劃分成多個(gè)ParamSlice，由于數(shù)據(jù)并行的存在,，同一個(gè)ParamSlice可能被多個(gè)設(shè)備同時(shí)擁有,。對于任意兩組并行策略（如下圖中的src策略和dst策略）,，它們對ParamSlice劃分的交集，被定義為熱切換通信的基本單位-模型切片,。
通信的基本單位：兩組并行策略對ParamSlice劃分（左側(cè)）的交集（右側(cè)）
熱切換問題定義：假設(shè)從當(dāng)前策略熱切換到目標(biāo)策略,，則對任意一個(gè)模型切片，遍歷目標(biāo)策略中需要該切片的每個(gè)設(shè)備（needed devices）,，并從當(dāng)前策略中擁有它的所有設(shè)備里（owned devices）選擇最合適的發(fā)送方,。
啟發(fā)式算法：基于兩個(gè)基本原則，我們提出了一種啟發(fā)式的搜索算法,，從而能夠?qū)ふ易詈线m的熱切換通信方案,。
基于上述兩個(gè)基本原則,，熱切換通信方案推導(dǎo)的啟發(fā)式算法流程如下：
（2）優(yōu)化熱切換開銷：Message Fusion & Layout Optimization
熱切換不可避免地會(huì)帶來額外開銷，包括通信開銷和顯存拷貝開銷,，這里利用消息合并（Message Fusion）和布局優(yōu)化（Layout Optimization）這兩項(xiàng)技術(shù)進(jìn)行專門優(yōu)化,。
消息合并（Message Fusion）：將發(fā)送給同一個(gè)設(shè)備的數(shù)據(jù)都合并到連續(xù)的發(fā)送緩沖區(qū)（send buffer）里，同理,，將從同一個(gè)設(shè)備接收的數(shù)據(jù)合并到連續(xù)的接收緩沖區(qū)（recv buffer）里，從而能夠合并多個(gè)p2p send/recv,，減少p2p kernel的調(diào)用次數(shù),，還能增大單次通信的數(shù)據(jù)量大小，提高帶寬利用率和通信效率,。此外,，通過NCCL提供的BatchedSendRecv原語，支持不同的send/recv buffer對應(yīng)的p2p send/recv并行傳輸,。
布局優(yōu)化（Layout Optimization）：為了避免引入contiguous算子和concat算子導(dǎo)致用kernel進(jìn)行訪存和數(shù)據(jù)搬運(yùn),，引入過高的拷貝開銷，這里將縱向切分的權(quán)重/梯度在布局上也按照橫向切分排布,。此時(shí)大部分的非連續(xù)的內(nèi)存訪問都可以轉(zhuǎn)化為連續(xù)的內(nèi)存訪問,，從而可以將大部分比較耗時(shí)的contiguous算子和concat算子直接轉(zhuǎn)化為訪存代價(jià)非常小的cudaMemCpy。具體計(jì)算時(shí),，只需要將改變gemm kernel的layout參數(shù)即可保證數(shù)學(xué)上的等價(jià)性,。
HotSPa是首個(gè)支持并行策略動(dòng)態(tài)熱切換的大模型分布式訓(xùn)練系統(tǒng),，相比現(xiàn)有的只支持靜態(tài)并行策略的系統(tǒng)（如Megatron-LM,,、DeepSpeed），HotSPa能更靈活地支持和適應(yīng)負(fù)載動(dòng)態(tài)變化的場景,，在現(xiàn)有的大多數(shù)長短序列分布不均衡的數(shù)據(jù)集中,，能夠獲得更高的訓(xùn)練吞吐。
實(shí)驗(yàn)設(shè)置：在實(shí)驗(yàn)中,，我們將HotSPa和現(xiàn)有的兩個(gè)分布式訓(xùn)練系統(tǒng)Megatron-LM（DP+TP+PP+SP）,、DeepSpeed（Zero1/2/3+Ulysses）在不同負(fù)載下進(jìn)行了比較。在實(shí)驗(yàn)環(huán)境上,，使用4臺(tái)GPU服務(wù)器,，每臺(tái)服務(wù)器上有8張A800-80G，機(jī)內(nèi)NVLink的通信帶寬為400GB/s,，機(jī)間IB通信帶寬為200GB/s,。在模型上選用了開源的LLaMA2，包括三種不同規(guī)模的參數(shù)量：7B,、13B和32B,。在數(shù)據(jù)集上選擇了兩個(gè)開源且被廣泛使用的數(shù)據(jù)集CommonCrawl和GitHub,。
端到端實(shí)驗(yàn)：在GPU數(shù)量為8卡~32卡，模型規(guī)模為LLaMA2-7B~32B,，最大序列長度為4k~32k的不同規(guī)模上進(jìn)行實(shí)驗(yàn),，在GitHub和CommonCrawl兩個(gè)數(shù)據(jù)集上，HotSPa相對于Megatron-LM分別取得最多1.5x和2.99x的加速比,，相對DeepSpeed分別取得最多2.6x和5x的加速比,。
熱切換代價(jià)實(shí)驗(yàn)：以LLaMA2-32B在32卡GPU上的測試結(jié)果為例，單次熱切換的時(shí)間開銷基本可以被優(yōu)化至1s以內(nèi),，相對于單個(gè)step的訓(xùn)練時(shí)間,，熱切換代價(jià)占比可以被忽略不計(jì)。
可擴(kuò)展性實(shí)驗(yàn)：Context Length: Megatron-LM的短序列被迫使用長序列的并行策略,，因此上下文長度越長,，HotSPa的加速比越大。GBS: Megatron-LM的時(shí)間基本隨全局批次大小線性增長,，而對HotSPa性能表現(xiàn)更優(yōu),，因?yàn)殡S著長序列數(shù)量增加，對應(yīng)分組的pp bubble減少,，會(huì)進(jìn)一步獲得加速,。GPUs：兩個(gè)系統(tǒng)都具有良好的擴(kuò)展性。
在這個(gè)工作里,，我們首次提出現(xiàn)有框架的靜態(tài)并行策略不適用于輸入序列長短變化的動(dòng)態(tài)場景的缺陷，并創(chuàng)新性地搭建了一套支持熱切換訓(xùn)練系統(tǒng)HotSPa,。HotSPa實(shí)現(xiàn)于PKU-DAIR實(shí)驗(yàn)室自研的分布式深度學(xué)習(xí)系統(tǒng)-河圖Hetu（https://github.com/PKU-DAIR/Hetu）,。除了性能上的優(yōu)勢，Hetu還有其他系統(tǒng)所不具備的高動(dòng)態(tài)性和高靈活性,。目前我們的系統(tǒng)已經(jīng)全面開源,，歡迎大家關(guā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),，長期從事數(shù)據(jù)庫系統(tǒng)、大數(shù)據(jù)管理與分析,、人工智能等領(lǐng)域的前沿研究,，在理論和技術(shù)創(chuàng)新以及系統(tǒng)研發(fā)上取得多項(xiàng)成果，已在國際頂級學(xué)術(shù)會(huì)議和期刊發(fā)表學(xué)術(shù)論文100余篇,，發(fā)布多個(gè)開源項(xiàng)目,。課題組同學(xué)曾數(shù)十次獲得包括CCF優(yōu)博、ACM中國優(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)化落地,。
