GPT-3模型為何難以復(fù)現(xiàn)?這也許是分布式AI框架的最優(yōu)設(shè)計(jì)(1)
2020 年,最轟動(dòng)的 AI 新聞莫過(guò)于 OpenAI 發(fā)布的 GPT-3 了。它的1750億參數(shù)量及其在眾多NLP任務(wù)上超過(guò)人類的出眾表現(xiàn)讓大家堅(jiān)信:大模型才是未來(lái)。但與之帶來(lái)的問(wèn)題是,訓(xùn)練超大模型所需的算力、存儲(chǔ)已不再是單機(jī)就能搞定的了(之前的 BERT 還是可以用 DGX-1/2 這樣的超級(jí)服務(wù)器訓(xùn)練)。
NVIDIA 估算過(guò),如果要訓(xùn)練GPT-3 ,即使單個(gè)機(jī)器的顯存/內(nèi)存能裝得下,用 8 張 V100 的顯卡(一臺(tái) DGX-1 的配置),訓(xùn)練時(shí)長(zhǎng)預(yù)計(jì)要 36 年;即使用 512 張 V100 ,訓(xùn)練也需要將近 7 個(gè)月;如果你擁有 1024 張 80GB A100, 那么完整訓(xùn)練 GPT-3 的時(shí)長(zhǎng)可以縮減到 1 個(gè)月。
這意味著訓(xùn)練大模型一定是一個(gè)分布式問(wèn)題。對(duì)算力的需求還是一個(gè)相對(duì)容易解決的問(wèn)題,因?yàn)閾碛写蠹旱慕M織并不只 OpenAI 一家,而如何解決上千塊 GPU 的分布式訓(xùn)練問(wèn)題才是關(guān)鍵。
即使你是一位非常優(yōu)秀的數(shù)據(jù)科學(xué)家,知曉并能解決 Transformer 相關(guān)的所有算法問(wèn)題,但如果你不是分布式專家,不知道如何解決分布式訓(xùn)練時(shí)上百臺(tái)服務(wù)器之間的通信、拓?fù)洹⒛P筒⑿?、流水并行等?wèn)題,你甚至都無(wú)法啟動(dòng)這次訓(xùn)練。這也解釋了為什么時(shí)隔一年,只有 NVIDIA 、微軟等大企業(yè)可以復(fù)現(xiàn) GPT-3 。
目前開(kāi)源的 GPT 模型庫(kù)主要是 NVIDIA 的 Megatron-LM 和微軟的 DeepSpeed。其中,微軟的 DeepSpeed 的模型并行等內(nèi)核取自 Megatron,且 DeepSpeed 主打的是,在數(shù)據(jù)并行下如何以更少的機(jī)器去跑更大的模型 ( ZeRO 、 ZeRO-Offload 等都是用梯度切片、計(jì)算、內(nèi)存/硬盤(pán)換入換出來(lái)省顯存),所以我們本文主要介紹和對(duì)比 Megatron 。
簡(jiǎn)單比較一下 NVIDIA 的 Megatron 和 微軟的 DeepSpeed:
DeepSpeed 本質(zhì)上是一種“節(jié)省顯存”的數(shù)據(jù)并行,即:數(shù)據(jù)并行的優(yōu)化版。DeepSpeed 假設(shè)了單層參數(shù)量可以在單張顯卡上放得下,如果不滿足這個(gè)假設(shè),那么仍然需要使用模型并行,而且 DeepSpeed 的模型并行是通過(guò)調(diào)用 Megatron 來(lái)實(shí)現(xiàn)的。
根據(jù) NVIDIA 最新的那篇論文(鏈接:https://arxiv.org/abs/2104.04473,也是下面本文重點(diǎn)要介紹的),Megatron 在大規(guī)模訓(xùn)練的效率是超過(guò) DeepSpeed 不少的。
而 DeepSpeed 的論文一直強(qiáng)調(diào):可以用更少機(jī)器訓(xùn)練更大的模型,但沒(méi)有突出過(guò)在效率上的優(yōu)勢(shì)。
DeepSpeed 后來(lái)又出了一篇論文:ZeRO-Infinity(鏈接:https://arxiv.org/abs/2104.07857),當(dāng)單層參數(shù)量在單張顯卡上放不下的時(shí)候,它通過(guò)對(duì)這一層算子切片,一片一片來(lái)執(zhí)行,使得單卡也能跑起來(lái)一個(gè)巨大的層,可以理解成一種 “時(shí)間”軸上展開(kāi)的模型并行。
Megatron 和 DeepSpeed 都是基于 PyTorch ,分別由 NVIDIA 和微軟經(jīng)過(guò)深度定制開(kāi)發(fā),專門為支持 PyTorch 分布式訓(xùn)練 GPT 而設(shè)計(jì)的。我們會(huì)簡(jiǎn)單介紹一下 NVIDIA 如何使用 PyTorch 搞分布式訓(xùn)練 GPT ,然后重點(diǎn)介紹 OneFlow 如何用一套通用設(shè)計(jì)非常簡(jiǎn)單清晰地解決了這個(gè)難題,同時(shí)我們還在已有的測(cè)試規(guī)模上性能超過(guò) NVIDIA。
相信讀完此文,你就會(huì)發(fā)現(xiàn) PyTorch 、 Megatron ( NVIDIA ) 、DeepSpeed ( Microsoft ) 都走了一個(gè)非常長(zhǎng)的彎路,這條彎路從大方向上就走錯(cuò)了,不僅是彎路,你還會(huì)發(fā)現(xiàn) Megatron 的代碼只能被 NVIDIA 的分布式訓(xùn)練專家所復(fù)用,它對(duì)于 PyTorch 的算法工程師而言門檻極高,是非常難用的,以至于任何想要用 PyTorch 復(fù)現(xiàn)一個(gè)分布式大模型的算法工程師,都得先等 NVIDIA 開(kāi)發(fā)完才能再使用 Megatron 提供的模型。同時(shí),我們也通過(guò)這樣一個(gè)例子來(lái)證明:為什么一個(gè)分布式深度學(xué)習(xí)框架要像 OneFlow 這樣設(shè)計(jì)。
本文內(nèi)容較多,先列出主要目錄:
1.分布式訓(xùn)練 GPT 的必要并行技術(shù)
流水并行
梯度累加
后向重計(jì)算
1F1B 策略
2.Megatron:PyTorch 分布式訓(xùn)練的極限,痛點(diǎn)在哪兒?
流水并行 PyTorch 需要手寫(xiě)專家級(jí)復(fù)雜調(diào)度器
模型并行 PyTorch 需要手工排線,在 kernel 里手寫(xiě)特定的、經(jīng)過(guò)復(fù)雜推導(dǎo)的通信原語(yǔ)
3.OneFlow 用一致性視角(Consistent View)輕松填平分布式訓(xùn)練難的鴻溝
流水并行 ,只需要配置 Placement 就夠了
數(shù)據(jù) + 模型的混合并行,只需要配置 Variable 的 SBP 就夠了
OneFlow:讓每一位算法工程師都有能力訓(xùn)練 GPT
4.為什么分布式深度學(xué)習(xí)框架要像 OneFlow 這樣設(shè)計(jì)?
OneFlow 系統(tǒng)層面如何實(shí)現(xiàn)流水并行
OneFlow 系統(tǒng)層面如何實(shí)現(xiàn)混合(數(shù)據(jù) & 模型)并行
5.GPT 訓(xùn)練性能對(duì)比: OneFlow vs Megatron
純數(shù)據(jù)并行性能對(duì)比
純模型并行性能對(duì)比
混合并行性能對(duì)比
流水并行 + 混合并行 性能對(duì)比
6.小結(jié)
分布式訓(xùn)練 GPT 的必要并行技術(shù)
最近,NVIDIA 放出了一篇重量級(jí)的論文:Efficient Large-Scale Language Model Training on GPU Clusters ,用了 3072 張 80GB A100 訓(xùn)練 GPT( NVIDIA 也確實(shí)夠壕,這個(gè)集群的成本就不止 5 億了),最大規(guī)模的模型參數(shù)量達(dá)到了 1T(是 GPT-3 原版的 5 倍)。
NVIDIA 訓(xùn)練 GPT-3 最大到 1T 參數(shù)規(guī)模
論文里 NVIDIA 介紹了分布式訓(xùn)練超大規(guī)模模型的三種必須的并行技術(shù):
數(shù)據(jù)并行(Data Parallelism)
模型并行(Tensor Model Parallelism)
流水并行(Pipeline Model Parallelism)
其中數(shù)據(jù)并行是大家都熟知的最常見(jiàn)的并行方式,而模型并行(NVIDIA 論文里叫做 "Tensor 級(jí)別的模型并行" )是對(duì)某一層(如 Linear/Dense Layer 里的 Variable )的模型 Tensor 切分,從而將大的模型 Tensor 分成多個(gè)相對(duì)較小的 Tensor 進(jìn)行并行計(jì)算;流水并行(NVIDIA 論文里叫做流水線級(jí)別的模型并行),是將整個(gè)網(wǎng)絡(luò)分段(stage),不同段在不同的設(shè)備上,前后階段流水分批工作,通過(guò)一種“接力”的方式并行。
對(duì)于最大的 1T 規(guī)模的模型,NVIDIA 一共使用了 384 臺(tái) DGX-A100 機(jī)器(每臺(tái)裝有 8 張 80GB A100 GPU),機(jī)器內(nèi)部各 GPU 間使用超高速 NVLink 和 NVSwitch 互聯(lián),每臺(tái)機(jī)器裝有 8個(gè) 200Gbps 的 InfiniBand (IB) 網(wǎng)卡,可以說(shuō)是硬件集群頂配中的頂配了。
那么,這些機(jī)器是如何協(xié)同工作的?GPT 網(wǎng)絡(luò)是由很多層 Transformer Layer 組成的,每一層內(nèi)部是一個(gè)由多層 MLP 和 attention 機(jī)制組成的子圖,對(duì)于參數(shù)規(guī)模 1T 的 GPT 而言就有 128 層的 Transformer Layer,這個(gè)超大超深的網(wǎng)絡(luò)被分割成了 64 個(gè) stage ,每個(gè) stage 跑在 6 臺(tái) DGX-A100 上,其中 6 臺(tái)機(jī)器之間進(jìn)行數(shù)據(jù)并行,每臺(tái)機(jī)器內(nèi)部的 8 張卡之間做模型并行,整個(gè)集群的 3072 張 A100 按照機(jī)器拓?fù)浔粍澐殖闪?[6 x 8 x 64] 的矩陣,同時(shí)使用數(shù)據(jù)并行 & 模型并行 & 流水并行 進(jìn)行訓(xùn)練。
3072 張 A100 集群拓?fù)?/p>
1.流水并行
從上述的機(jī)器拓?fù)渲锌梢园l(fā)現(xiàn),流水并行是 3072 塊 A100 能訓(xùn)練 GPT 的關(guān)鍵。因?yàn)闊o(wú)論是數(shù)據(jù)并行還是模型并行,都會(huì)在相應(yīng)的機(jī)器之間進(jìn)行全連接的通信,當(dāng)機(jī)器數(shù)量增大時(shí),通信開(kāi)銷和時(shí)延會(huì)大到難以忍受。而流水并行既解決了超大模型無(wú)法在單設(shè)備上裝下的難題,又很好解決了機(jī)器之間的通信開(kāi)銷的問(wèn)題,每個(gè)階段(stage) 和下一個(gè)階段之間僅有相鄰的某一個(gè) Tensor 數(shù)據(jù)需要傳輸,每臺(tái)機(jī)器的數(shù)據(jù)傳輸量跟總的網(wǎng)絡(luò)大小、機(jī)器總數(shù)、并行規(guī)模無(wú)關(guān)。但流水并行為了多個(gè)階段之間可以流水起來(lái),還依賴兩個(gè)重要的特性: 梯度累加(Gradient Accumulation) 和 亞線性內(nèi)存優(yōu)化( Sublinear Memory Cost 2016, 陳天奇)。
近期,百度和華為相繼發(fā)了自己的千億級(jí)中文預(yù)訓(xùn)練模型的宣傳文。其中,百度提出了 "4D混合并行",本質(zhì)上是 Megatron 里的數(shù)據(jù)并行 + 模型并行 + 流水并行 + DeepSpeed 里的 ZeRO 優(yōu)化 ;華為文章中的 “5D混合并行”,是將重計(jì)算(Checkpointing, 亞線性內(nèi)存優(yōu)化的一種)作為了第5維 (其實(shí)百度也做了重計(jì)算,只是沒(méi)有將其列為多維并行中的一維)。
在介紹這兩個(gè)特性之前,我們先簡(jiǎn)單解釋一下深度學(xué)習(xí)訓(xùn)練和模型更新的兩種約束:BSP (Bulk Synchronous Parallel) 和 SSP (Stale Synchronous Parallel ) ,其中 BSP 是最常見(jiàn)的模型更新規(guī)則:每個(gè) batch 的前向計(jì)算都需要使用最新的模型,那么就要求上一個(gè) batch 的后向計(jì)算結(jié)束且在模型更新后,下一個(gè) batch 的前向才能開(kāi)始。如果使用 BSP 去做流水并行,我們就會(huì)發(fā)現(xiàn)每個(gè)階段的前向和后向是完全串行的,其中一個(gè)設(shè)備在工作時(shí),其他所有設(shè)備都在等待,那么分布式的優(yōu)勢(shì)就完全沒(méi)有被發(fā)揮出來(lái):
BSP 各個(gè)階段串行執(zhí)行
BSP 且沒(méi)有 Gradient Accumulation 下的流水并行。假設(shè)整個(gè)網(wǎng)絡(luò)被等分切成 4 個(gè) stage,每個(gè) stage 使用一個(gè) device ,則在BSP下,各個(gè)設(shè)備串行執(zhí)行,中間有大段的氣泡。一般后向計(jì)算時(shí)間是前向計(jì)算的兩倍,如果算上 Checkpointing 的重計(jì)算部分,是前向計(jì)算的三倍。我們可以從上圖中看到,這種情況下有 70% 的時(shí)間設(shè)備是空閑的。
而 SSP 就是異步模型更新,允許前向計(jì)算時(shí)可以不使用最新的模型,而使用落后幾個(gè)版本之內(nèi)的模型。SSP 在 GPT-3 的訓(xùn)練中并沒(méi)有被 NVIDIA 采用,其主要原因有以下幾點(diǎn):
SSP 的模型收斂性并沒(méi)有被嚴(yán)格論證, 且有論文 GeePS 指出 SSP 的收斂效果不如 BSP ;
SSP 會(huì)在 GPU 上同時(shí)出現(xiàn)幾個(gè)不同版本的模型,而 GPT-3 又是一個(gè)非常大的模型網(wǎng)絡(luò),多份模型所帶來(lái)的顯存開(kāi)銷不可接受;
BSP 通過(guò) Gradient Accumulation + Checkpointing 就可以很好的解決 GPT-3 中的流水并行問(wèn)題。
另外, NVIDIA 的論文的分析前提就是 BSP 情況下, 根據(jù)嚴(yán)格的參數(shù)優(yōu)化器更新方式, 流水并行氣泡的占比是 Bubble time fraction = (p - 1) / m,其中 p 是 stage 數(shù), m 是梯度累加的 micro-batch 數(shù)。如果采用 SSP,則 NVIDIA 整篇文章的理論基礎(chǔ)就沒(méi)有了。
"Pipeline parallelism comes in a few flavors: the mode discussed in this paper uses flushes to ensure exact strict optimizer semantics."
2.梯度累加
Gradient Accumulation 就是把一個(gè)大 Batch 拆分成多個(gè) micro-batch , 每個(gè) micro-batch 前后向計(jì)算后的梯度累加,在最后一個(gè)micro-batch累加結(jié)束后,統(tǒng)一更新模型。
micro-batch 跟數(shù)據(jù)并行有高度的相似性:數(shù)據(jù)并行是空間上的, 數(shù)據(jù)被拆分成多個(gè) tensor,同時(shí)喂給多個(gè)設(shè)備并行計(jì)算,然后將梯度累加在一起更新;而 micro-batch 是時(shí)間上的數(shù)據(jù)并行, 數(shù)據(jù)被拆分成多個(gè) tensor, 按照時(shí)序依次進(jìn)入同一個(gè)設(shè)備串行計(jì)算,然后將梯度累加在一起更新。當(dāng)總的 batch size 一致,且數(shù)據(jù)并行的并行度和 micro-batch 的累加次數(shù)相等時(shí),數(shù)據(jù)并行和 Gradient Accumulation 在數(shù)學(xué)上完全等價(jià)。Gradient Accumulation 通過(guò)多個(gè) micro-batch的梯度累加使得下一個(gè) micro-batch 的前向計(jì)算不需要依賴上一個(gè) micro-batch 的反向計(jì)算,因此可以暢通無(wú)阻的進(jìn)行下去(當(dāng)然在一個(gè)大 batch 的最后一次 micro-batch 還是會(huì)觸發(fā)這個(gè)依賴)。
Gradient Accumulation 解決了很多問(wèn)題:
在單卡下,Gradient Accumulation 可以將一個(gè)大的 batch size 拆分成等價(jià)的多個(gè)小 micro-batch ,從而達(dá)到節(jié)省顯存的目的。
在數(shù)據(jù)并行下,Gradient Accumulation 解決了反向梯度同步開(kāi)銷占比過(guò)大的問(wèn)題(隨著機(jī)器數(shù)和設(shè)備數(shù)的增加,梯度的 AllReduce 同步開(kāi)銷也加大),因?yàn)樘荻韧阶兂闪艘粋€(gè)稀疏操作,因此可以提升數(shù)據(jù)并行的加速比。
在流水并行下, Gradient Accumulation 使得不同 stage 之間可以并行執(zhí)行不同的 micro-batch, 從而讓各個(gè)階段的計(jì)算不阻塞,達(dá)到流水的目的。
單純通過(guò) micro-batch,我們就實(shí)現(xiàn)了 GPipe (2018)論文中的流水并行,在 stage 數(shù)量為 4, micro-batch 數(shù)量為 8 (每個(gè) batch 在計(jì)算 8 個(gè) micro-batch 且 累加 8 次梯度后更新一次模型)下的時(shí)間線如下圖所示:
使用梯度累加后的 Pipeline 時(shí)間線
在 GPipe 的流水并行示例中,每個(gè)“時(shí)間點(diǎn)” 可以在多個(gè)階段(stage)上同時(shí)做不同的micro-batch,圖中每個(gè)方塊中的標(biāo)號(hào)表示了第幾個(gè) micro-batch;同一個(gè) micro-batch 還是串行的經(jīng)過(guò)所有的 stage,在這種情況下,每個(gè)設(shè)備的空閑時(shí)間只有 25% 左右。
但這種流水并行存在一個(gè)問(wèn)題:顯存占用太大。如果每個(gè) micro-batch 前向計(jì)算的中間結(jié)果(activation)被后向計(jì)算所消費(fèi),則需要在顯存中緩存 8 份(梯度累加的次數(shù))完整的前向 activation。這時(shí)就不得不用另一項(xiàng)重要的技術(shù):重計(jì)算(Checkpointing)。
*博客內(nèi)容為網(wǎng)友個(gè)人發(fā)布,僅代表博主個(gè)人觀點(diǎn),如有侵權(quán)請(qǐng)聯(lián)系工作人員刪除。