超大Transformer語言模型的分布式訓(xùn)練框架

NVIDIA Megatron 是一個基于 PyTorch 的框架，用于訓(xùn)練基于 Transformer 架構(gòu)的巨型語言模型。本系列文章將詳細介紹Megatron的設(shè)計和實踐，探索這一框架如何助力大模型的預(yù)訓(xùn)練計算。

大模型是大勢所趨

近年來，NLP 模型的發(fā)展十分迅速，模型的大小每年以1-2個數(shù)量級的速度在提升，背后的推動力當(dāng)然是大模型可以帶來更強大更精準(zhǔn)的語言語義理解和推理能力。

截止到去年，OpenAI發(fā)布的GPT-3模型達到了175B的大小，相比2018年94M的ELMo模型，三年的時間整整增大了1800倍之多。按此趨勢，預(yù)計兩年后，會有100 Trillion參數(shù)的模型推出。

另外一個特點是，自從18年 Google 推出 Attention is All You Need論文后，這幾年的模型架構(gòu)，不管是雙向的BERT，還是生成式的GPT，都是基于Transformer 架構(gòu)來構(gòu)建的，通常說的模型有多少層，指的便是有多少個Transformer塊來堆疊起來的。

而且，這類模型的計算量也主要來自于對Transformer塊的處理，其本質(zhì)上可以轉(zhuǎn)化成大量的矩陣操作，天然地適合NVIDIA GPU的并行架構(gòu)。

分布式是大模型訓(xùn)練的必須

大模型的預(yù)訓(xùn)練對計算、通信帶來的挑戰(zhàn)是不言而喻的。我們以GPT-3 175B 模型為例，分析預(yù)訓(xùn)練對計算量、顯存、通信帶來的挑戰(zhàn)。

GPT-3 175B模型的參數(shù)如下：網(wǎng)絡(luò)層（Number of layers）： 96

句子長度（Sequence length）： 2048

隱藏層大?。℉idden layer size）： 12288

詞匯表（Vocabulary size）：51200

總參數(shù)量：約175B

1. 對顯存的挑戰(zhàn)

175B的模型，一個原生沒有經(jīng)過優(yōu)化的框架執(zhí)行，各部分大概需要的顯存空間：

模型參數(shù)：700 GB （175B * 4bytes）

參數(shù)對應(yīng)的梯度：700 GB

優(yōu)化器狀態(tài)：1400 GB

所以，一個175B模型共需要大概2.8 TB的顯存空間，這對 GPU 顯存是巨大的挑戰(zhàn)：

1）模型在單卡、單機上存放不下。以 NVIDIA A100 80GB為例，存放此模型需要超過35塊。

2） 必須使用模型并行，并且需要跨機器。主流的A100 服務(wù)器是單機八卡，需要在多臺機器之間做模型切分。

2. 對計算的挑戰(zhàn)

基于Transformer 架構(gòu)的模型計算量主要來自于Transformer層和 logit 層里的矩陣乘，可以得出每個迭代步大致需要的計算量：

B： 批大小，S：句子長度，l：Transformer 層數(shù)，h：隱藏層大小，V：詞匯表大小

這是真實計算量的一個下限，但已是非常接近真實的計算量。關(guān)于此公式的詳細說明，請參考 NVIDIA Paper（https://arxiv.org/abs/2104.04473）里的附錄章節(jié)。

其中S=2048， l=96， h=12288， V=51200，在我們的實踐中，B = 1536，一共需要迭代大約95000次。代入這次參數(shù)到上述公式，可以得到：

一次迭代的計算量：4.5 ExaFLOPS.

完整訓(xùn)練的計算量：430 ZettaFLOPS （~95K 次迭代）

這是一個巨大的計算量，以最新的 NVIDIA A100 的FP16計算能力 312 TFLOPS來計算，即使不考慮計算效率和擴展性的情況，需要大概16K A100*days的計算量。直觀可以理解為16000塊A100一天的計算量，或者一塊A100 跑43.8年的計算量。

3. 對通信的挑戰(zhàn)

訓(xùn)練過程中GPU之間需要頻繁的通信，這些通信源于模型并行和數(shù)據(jù)并行的應(yīng)用，而不同的并行劃分策略產(chǎn)生的通信模式和通信量不盡相同。

對于數(shù)據(jù)并行來說，通信發(fā)生在后向傳播，用于梯度通信，通信類型為AllReduce，每次后向傳播中的通信量為每個GPU上的模型大小。

對于模型并行來說，稍微復(fù)雜些。模型并行通常有橫切和豎切兩種，比如把一個模型按網(wǎng)絡(luò)層從左到右橫著擺放，橫切即把每個網(wǎng)絡(luò)層切成多份（Intra-layer），每個GPU上計算網(wǎng)絡(luò)層的不同切塊，也稱為Tensor（張量）模型并行。豎切即把不同的網(wǎng)絡(luò)層切開（Inter-layer），每個GPU上計算不同的網(wǎng)絡(luò)層，也稱為Pipeline （流水線）模型并行。

對于Tensor模型并行，通信發(fā)生在每層的前向和后向傳播，通信類型為AllReduce，通信頻繁且通信量比較大。

對于Pipeline 模型并行，通信發(fā)生在相鄰的切分點，通信類型主要為P2P，每次通信數(shù)據(jù)量比較少但比較頻繁，而且會引入額外的GPU 空閑等待時間。

稍后會詳細闡述在Transformer 架構(gòu)上如何應(yīng)用這兩種模型劃分方式。

更為復(fù)雜的是，對于超大的語言模型，通常會采用數(shù)據(jù)并行 + Tensor 模型并行 + Pipeline 模型并行混合的方式，這使得通信方式錯綜復(fù)雜在一起，對系統(tǒng)連接拓撲提出更大的挑戰(zhàn)：能靈活滿足不同劃分策略、不同通信模式下，不同通信組里高效的通信。

總而言之，超大語言模型的預(yù)訓(xùn)練，采用多節(jié)點的分布式訓(xùn)練是必須，而且是基于模型并行的。這就對集群架構(gòu)和訓(xùn)練框架提出了嚴苛的設(shè)計要求，集群架構(gòu)要有優(yōu)化的互聯(lián)設(shè)計，訓(xùn)練框架更為重要：不僅僅是結(jié)合算法特點對模型做合理切割，更是需要做出結(jié)合系統(tǒng)架構(gòu)特點、軟硬一體的co-design。

為此，NVIDIA 分別提出了優(yōu)化的分布式框架NVIDIA Megatron 和優(yōu)化的分布式集群架構(gòu) NVIDIA DGX SuperPOD。

優(yōu)化的分布式框架：NVIDIA Megatron

Megatron設(shè)計就是為了支持超大的Transformer模型的訓(xùn)練的，因此它不僅支持傳統(tǒng)分布式訓(xùn)練的數(shù)據(jù)并行，也支持模型并行，包括Tensor并行和Pipeline并行兩種模型并行方式。

1. Tensor 模型并行

上面我們看到，對于一個Transformer塊，主要包括Masked Multi Self Attention和Feed Forward兩個部分，對于Tensor并行，需要把這兩部分都并行化。

對于Feed Forward部分，是由多個全連接層組成的MLP網(wǎng)絡(luò)，每個全連接層由矩陣乘和GeLU激活或Dropout組成，在Megatron中，F(xiàn)eed Forward采用兩層全連接層。對于一個全連接層，可以表示為：

其中X輸入，A為參數(shù)矩陣，Y為輸出，則可以有兩種并行方式。

一種是按行的方向把權(quán)重矩陣A切分開并按列的方向把輸入X切分開，即：則輸出：

其中括號中的每一項，可以在一個單獨的GPU上獨立的完成，再通過一次AllReduce完成求和操作。

另一種則是按列的方向把權(quán)重矩陣A切分開，而不切分輸入，即：

則可以得到同樣按列方向切分開的輸出：

方括號中每一項可以在一個單獨的GPU上獨立的完成，這樣每個GPU上得到部分的最終輸出，大家拼接在一起就是完整輸出，不需要再做AllReduce。

Megatron在計算MLP時采用了這兩種并行方式，具體如下圖所示：

整個MLP的輸入X先通過f放到每一塊GPU上，然后先使用上面提到的按列切分權(quán)重矩陣A的方式，在每塊GPU上得到第一層全連接的部分輸出Y1和Y2，然后采用按行切分權(quán)重矩陣B，按列切分Y的方式，其中前一層的輸出Y1和Y2剛好滿足Y的切分需求，因此可以直接和B的相應(yīng)部分做相應(yīng)的計算而不需要額外操作或通信。這樣得到了最終Z的部分、Z1和Z2，通過g做AllReduce得到最終的Z，再通過相應(yīng)的激活層或Dropout。

這樣就完成了MLP層的Tensor并行，對于Masked Multi Self Attention層，如下圖所示：

正如它的名字中提到的，它是由多個Self Attention組成的，因此很自然的并行方式就是可以把每個Self Attention分到不同的GPU上去進行計算，這樣每塊GPU上就能夠得到輸出的一部分，最后的Linear全連接層，由于每個GPU上已經(jīng)有部分輸出，因此可以采用上面全連接層的按行的方向切權(quán)重矩陣B并按列的方向切輸入Y的方式直接進行計算，再通過AllReduce操作g得到最終結(jié)果。

這樣我們就可以完成Transformer塊的Tensor并行。有了Tensor并行，我們可以把模型的每一層進行切分，分散到不同的GPU上，從而訓(xùn)練比較大的模型。由于Tensor并行會對每一層進行切分，并且需要通信，因此Tensor并行在同一臺機器上，并且有NVLink的加速情況下性能最好。如果模型進一步增大，大到一臺機器可能都放不下整個模型，這時就需要引入另一種并行方式，Pipeline并行。

2. Pipeline 模型并行

相對于Tensor并行的把模型的每一層內(nèi)部進行切分，Pipeline并行是會在模型的層之間進行切分，不同的層在不同的GPU或機器節(jié)點上進行計算。由于不同的層間有依賴關(guān)系，所以如果直接并行會像下圖所示，黑色部分是前向，綠色部分是反向計算，灰色部分是空閑，可以看出GPU的絕大部分時間是在等待。

為了解決這個問題，Megatron把每一個batch分成了更小的microbatch，如下圖所示，把batch 1分成了1a，1b，1c，1d四個microbatch，由于不同的microbatch間沒有數(shù)據(jù)依賴，因此互相可以掩蓋各自的等待時間，提高GPU利用率，提升整體的性能。

這就是Megatron 核心的兩種模型并行的設(shè)計，可以支撐超大的Transformer-based 語言模型，再結(jié)合經(jīng)典的數(shù)據(jù)并行方式，可以讓大模型的訓(xùn)練更快。

編輯：jq