擎天柱：訓練超大模型的高效二維方法

摘要

超大神經網路在實際應用中展現了空前的效能。然而，由於儲存限制，我們必須使用模型並行來儲存大型模型，否則單個裝置將無法容納這些模型。包括Megatron(威震天)在內的先前的工作將整個模型的引數劃分到多個裝置，然而每個裝置必須在正向和反向傳播中儲存冗餘的中間啟用。在這項工作中，我們提出了Optimus(擎天柱)，一個利用二維矩陣分割的高效可擴充套件模型並行正規化，它將有助於訓練無窮大的語言模型。在Optimus中，中間啟用也被分割並分佈到諸多裝置中，由此進一步減少了冗餘。就等效效率而言，Optimus顯著勝過了Megatron。在美國德州高階計算中心的64個GPU上，Optimus實現了1.48倍Megatron的訓練速度，1.78倍的推理速度，以及8倍的最大批處理量。Optimus在擴充套件效率上大大超過了Megatron。原文請參見https://arxiv.org/abs/2104.05343。程式碼請參見http://github.com/xuqifan897/Optimus。

介紹

大模型對硬體和軟體都帶來了巨大挑戰。為了應對這一挑戰，研究人員除了建造超級計算機之外還開發了新的演算法。諸如中間啟用存檔(activation checkpointing)，混合精度訓練(mixed precision training)，資料轉移(data offload)等方法從不同角度來輔助大模型的訓練。除這些工作之外，模型並行是另外一個重要的正規化。它把大模型的引數分佈到諸多裝置上。流水線並行(pipeline parallelism)是將整個模型按照層間結構劃分。在這種情況下，每個處理器從前一個處理器獲得輸入，然後將輸出結果傳給下一個處理器。另外一種並行方式是層內並行，即將一個運算分佈到多個處理器同時進行。Megatron就抓住了語言模型矩陣乘法的本質，將權重引數逐行或逐列地切分並分佈到不同處理器上。這樣計算就是分塊矩陣乘法了。這種方法雖然將引數分佈在了諸多處理器上，每個處理器卻需要儲存一整個中間啟用。所以在高並行度下，中間啟用會成為儲存的瓶頸。

值得注意的是，Megatron本質上採用了一維矩陣劃分，在這種情況下，每一步通訊都會涉及所有的處理器。然而，如果我們考慮諸如SUMMA和Cannon's algorithm的二維矩陣分割，就有可能加速大模型的推理和訓練。在這篇工作中，我們提出了Optimus，一種新型的模型並行正規化。它同時分佈了引數和啟用，實現了高超的通訊效率和前所未有的儲存效能。Optimus在PyTorch上實現，並且可以和前述訓練大模型的技巧協同使用。Optimus的等效率函式(isoefficiency function)可達到

，顯著優於(小於)Megatron的

。我們相信新型超級計算機的網格拓撲可以進一步解放Optimus的卓越效能。相對於Megatron，Optimus可以幫助我們將模型訓練擴充套件到無窮大。我們在美國德州高階計算中心的64個GPU上實驗演示了Optimus相對Megatron 1.48倍的訓練速度，1.78倍的推理速度和8倍的最大批處理量。證明了Optimus在強擴充套件效率和弱擴充套件效率上均超過Megatron。總的來說，我們的貢獻如下：

基於二維矩陣劃分的思想，我們設計和實現了一種用於超大模型的高效訓練系統。
我們提出了諸多高效能方法，例如二維分割梯度計算，高儲存效率啟用管理，以及系統緩衝技巧。
我們的方法可以基於PyTorch輕鬆實現，無需涉及編譯層。

背景

在這裡，我們用以下記號：

批處理大小(batch size): b
序列長度(sequence length): s
隱層大小(hidden size): h
注意力頭(attention heads)數量: n
詞彙量大小(vocabulary size): v
SUMMA維度: q
Transformer 層數: N
處理器個數: p

我們有

。此後我們用方括號表示張量的形狀。

Megatron的模型並行包括多層感知機(MLP)和自注意力元件(self-attention)，其結構如下圖所示.

Figure1 Megatron中MLP和MHA的多處理器對映

在MLP中有兩個引數矩陣，第一個按列劃分，第二個按行劃分。每個處理器都持有一份相同的輸入，並和同一處理器上的第一個子矩陣相乘，得到同樣按列劃分的中間結果。中間結果首先輸入啟用函式，然後再和第二個子矩陣相乘，得到了一個和輸入相同大小的輸出矩陣。不同處理器上輸出矩陣經歷一個歸約(all-reduce)操作，得到最終的輸出。

同樣地，在self-attention中也有兩個引數矩陣。第一個按列劃分，第二個按行劃分。相同的輸入矩陣被每個處理器所持有，並和第一個子矩陣相乘，得到中間結果。由於attention heads被分佈到不同的處理器，每個處理器就負責n/p個attention heads。中間結果被劃分為n/p組Q,K和V矩陣，每一組計算得到

。n/p個計算結果在一起重新排列，並和第二個子矩陣相乘，再歸約得到最終結果。同樣地，每個處理器都需要儲存一整個的輸入和輸出。

在這種正規化下，通訊僅發生在正向和反向傳播的歸約操作上。

Mesh-Tensorflow是谷歌提出的分散式機器學習的框架。其引入了張量分割的簡約表達。它首先將處理器排布為一個多維網格。對於每個網格維度，一個張量要麼在這個維度上分割，要麼在這個維度上覆制。雖然Mesh-Tensorflow允許多維分割，但其實質上還是採用了和Megatron相似的矩陣分割方式。事實上，其不可能完全消除矩陣複製造成的儲存冗餘，即不可能每個張量在每個維度上都沒有複製。

我們的方法

SUMMA (Scalable Universal Matrix Multiplication Algorithm)，或可擴充套件通用矩陣乘法，包含了四種矩陣乘法。在這裡，我們僅討論其中三個，即

和

。事實上，這三個乘法相互為對方的微分計算：

接下來，我們假設處理器組成一個

的網格，同時每個張量都被平均劃分為

的子矩陣，每個都分佈在相應的處理器上。

表示矩陣A的第i行第j列的子矩陣。我們注意到這種劃分方式和谷歌TPU的構型相符。

以

為例，SUMMA將矩陣乘法視為一組外積(outer product)的和。如果我們將A矩陣按照列劃分：

，將B矩陣按照行劃分：

。這裡我們用上標來表示列矩陣，用下標來表示行矩陣以示區分。接下來，我們只看外積

。我們可以推出，第j行第k列的處理器所持有的

子矩陣為

。也就是說，第j行的每個處理器都需要

，第k列的每個處理器都需要

。下圖演示了其細節：

Figure2 SUMMA演算法示意圖 i=2

我們以16個處理器為例，且僅展示i=2的情況。最開始，每個處理器持有其自己的A和B的子矩陣。接下來，第二列的所有處理器將自己的A子矩陣廣播到同一行的每個處理器；第二行的所有處理器將自己的B子矩陣廣播到同一列的每個處理器。然後，每個處理器在本地計算其持有的A和B的子矩陣相乘，並與i=1時的結果相加。

Optimus的MLP和self-attention由下圖所示：

Figure3 Optimus中MLP和MHA的多處理器對映

輸入和輸出矩陣形狀為[b,s,h]。在這裡，我們僅對b和h維度進行劃分，s維度保持不變。在圖中，s維度被省略。MLP為兩個連續的矩陣乘法，中間有一個非線性變換。由此，其可以直接表達為SUMMA形式。在self-attention中，第一個矩陣乘法可以表達為SUMMA形式。在這之後，每個處理器都擁有大小為[b/q,s, 3h/q]的子矩陣。這時，我們將其劃分為n/q組Q,K和V，形狀均為[b/q,s, h/n]。這樣，我們就可以在本地計算

。我們將n/q個計算結果重新排布為[b/q,s, h/q]的形狀，再按SUMMA的方法計算第二個矩陣乘法。

我們注意到，Transformer由N個結構相同的Transformer layer組成。這也就揭示了儲存的複用性，即每一層的計算可以使用相同一塊儲存空間。由此，除了中間啟用存檔(activation checkpointing)外，我們使用了緩衝工作區(workspace buffer)，正向傳播緩衝區(forward buffer)，反向傳播緩衝區(backward buffer)，引數梯度緩衝區(parameter gradient buffer)和連線緩衝區(conjunction buffer)。在SUMMA中，在每次通訊前需要複製本地的子矩陣。現在，我們從緩衝工作區裡分出一塊來存放複製的本地子矩陣。正向傳播緩衝區用於存放正向傳播的中間啟用，反向傳播緩衝區用於存放中間啟用的梯度，引數梯度緩衝區用於存放參數的梯度，連線緩衝區用來存放層間的中間啟用和梯度，以防止在重置緩衝區時丟失後續計算中需要的矩陣。

實驗結果

我們在美國德州高階計算中心的超級計算機上做了實驗，得到並比較了弱擴充套件效率，強擴充套件效率和儲存效能。結果如下圖所示：

Figure4 weak scaling和strong scaling

可以看到，儘管隨著GPU數量的增加，Optimus和Megatron的弱擴充套件效率都在下降，但是Optimus相對Megatron的優勢越來越顯著。在64結點上，Optimus達到了1.48倍Megatron的訓練速度和1.78倍Megatron的推理速度。在強擴充套件效率上，我們觀察到了隨著GPU數量的增加，Megatron的效率在下降，而Optimus的效率在上升。這是因為在SUMMA中，當矩陣大小不變時，處理器數量增加，通訊時間減少。

在測試弱擴充套件效率的相同條件下，我們測試了Optimus和Megatron的視訊記憶體效率。即我們改變b的大小，直到視訊記憶體耗盡(out-of-memory)。在下圖中，

表示在

時，可以執行，但當

時出現視訊記憶體耗盡。隨著GPU數量的增加，Megatron的批處理量在減少，而Optimus的批處理量在增加。在64GPU上，Optimus達到了Megatron 8倍的批處理量。

Figure5 記憶體效率

總結

在這篇工作中，我們提出了一種全新的模型並行正規化。該正規化成功地消除了儲存的瓶頸且提升了通訊效率。從而為更大語言模型的開發鋪平了道路。在近一段時間的工作中，混合專家模型(Mixtureof Experts,MOE)得到了越來越多的關注。許多強大的模型都是基於MOE實現。未來的工作可能會致力於解決MOE的通訊和儲存問題。此外，運算元融合也是一個新方向。該方法可以顯著減少中間結果的儲存，從而提高儲存效率。比如在我們的實現中，注意力得分(attention score)會佔據一個形狀為[b,n, s, s]的張量。而中間啟用的形狀為[b,s, h]。當n=16,s=512,h=1024時，注意力得分所佔據的儲存是中間啟用的8倍，但只佔總計算量的一小部分。所以，利用運算元融合，我們有希望解放這一部分儲存。隨著體量的增加，我們相信巨型模型將能夠勝任更多令人驚豔的任務。

作者簡介

許啟帆，加州大學洛杉磯分校一年級博士生，研究方向是高效能計算，醫學AI，機器學習。許啟帆本科畢業於中國科學技術大學，在讀期間入選拔尖計劃，於2019年入選加州大學洛杉磯分校CSST專案。2020-2021年在新加坡國立大學HPC-AI實驗室實習期間完成這項工作。

尤洋，壁仞科技顧問，是新加坡國立大學計算機系的校長青年教授。他從加州大學伯克利分校計算機係獲得了博士學位。他的導師是美國科學院院士，工程院院士，ACM/IEEE fellow，伯克利前計算機系主任以及EECS學院院長James Demmel教授。尤洋的研究興趣包括高效能計算，並行演算法，以及機器學習。他當前的研究重點是大規模深度學習訓練演算法的分散式最佳化。他曾創造ImageNet以及BERT訓練速度的世界紀錄，並被ScienceDaily，TheNextWeb，i-programmer等幾十家媒體廣泛報道。他設計的演算法被廣泛應用於谷歌，微軟，英特爾，英偉達等科技巨頭。尤洋近三年以第一作者身份在NIPS，ICLR，Supercomputing，IPDPS，ICS等國際重要會議或期刊上發表論文十餘篇。他曾以第一作者身份獲得了國際並行與分散式處理大會(IPDPS)的最佳論文(0.8%獲獎率)和國際並行處理大會(ICPP)的最佳論文(0.3%獲獎率)。尤洋曾獲清華大學優秀畢業生，北京市優秀畢業生，國家獎學金，以及當時清華大學計算機係數額最高的西貝爾獎學金。他還在2017年獲得美國計算機協會(ACM)官網上唯一頒給在讀博士生的ACM-IEEECS George Michael Memorial HPC Fellowships。他也獲得了頒發給伯克利優秀畢業生的LotfiA.Zadeh Prize。尤洋還被伯克利提名為ACM Doctoral Dissertation Award候選人(81名UC Berkeley EECS 2020博士畢業生中選2人)。2021年4月，尤洋入選福布斯亞洲30歲以下精英榜。更多資訊請檢視他的研究小組主頁(https://ai.comp.nus.edu.sg/)。

往期推薦

1、MLIR編譯技術應對硬體設計挑戰？—初探CIRCT專案

2、圖神經網路加速器深度調研（中）

3、圖神經網路加速器深度調研（下）