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