導(dǎo)讀

流水線性并行、張量并行、3D并行三種分布式訓(xùn)練方法的詳細(xì)解讀，從原理到具體方法案例。

流水線性并行和張量并行都是對(duì)模型本身進(jìn)行劃分，目的是利用有限的單卡顯存訓(xùn)練更大的模型。簡(jiǎn)單來說，流水線并行水平劃分模型，即按照層對(duì)模型進(jìn)行劃分；張量并行則是垂直劃分模型。3D并行則是將流行線并行、張量并行和數(shù)據(jù)并行同時(shí)應(yīng)用到模型訓(xùn)練中。

一、流水線并行

流水線并行的目標(biāo)是訓(xùn)練更大的模型。本小節(jié)先介紹符合直覺的樸素層并行方法，并分析其局限性。然后，介紹流水線并行算法GPipe和PipeDream。

1. 樸素層并行

當(dāng)一個(gè)模型大到單個(gè)GPU無法訓(xùn)練時(shí)，最直接的想法是對(duì)模型層進(jìn)行劃分，然后將劃分后的部分放置在不同的GPU上。下面以一個(gè)4層的序列模型為例，介紹樸素層并行：

將其按層劃分至兩個(gè)GPU上：

• GPU1負(fù)責(zé)計(jì)算：intermediate (input)) ;
• GPU2負(fù)責(zé)計(jì)算：output (intermediate;

動(dòng)圖

整個(gè)樸素層并行前向傳播和后向傳播的過程如上圖所示。GPU1執(zhí)行前向傳播, 并將激活 (activations)緩存下來。然后將層的輸出intermediate發(fā)送給GPU2, GPU2完成前向傳播和loss計(jì)算后, 開始反向傳播。當(dāng)GPU2完成反向傳播后, 會(huì)將的梯度返還給GPU1。GPU1完成最終的反向傳播。

根據(jù)上面的介紹，可以發(fā)現(xiàn)樸素層并行的缺點(diǎn)：

• 低GPU利用率。在任意時(shí)刻，有且僅有一個(gè)GPU在工作，其他GPU都是空閑的。
• 計(jì)算和通信沒有重疊。在發(fā)送前向傳播的中間結(jié)果(FWD)或者反向傳播的中間結(jié)果(BWD)時(shí)，GPU也是空閑的。
• 高顯存占用。GPU1需要保存整個(gè)minibatch的所有激活，直至最后完成參數(shù)更新。如果batch size很大，這將對(duì)顯存帶來巨大的挑戰(zhàn)。

2. GPipe

2.1 GPipe的原理

GPipe通過將minibatch劃分為更小且相等尺寸的microbatch來提高效率。具體來說，讓每個(gè)microbatch獨(dú)立的計(jì)算前后向傳播，然后將每個(gè)mircobatch的梯度相加，就能得到整個(gè)batch的梯度。由于每個(gè)層僅在一個(gè)GPU上，對(duì)mircobatch的梯度求和僅需要在本地進(jìn)行即可，不需要通信。

假設(shè)有4個(gè)GPU，并將模型按層劃分為4個(gè)部分。樸素層并行的過程為

Timestep	0	1	2	3	4	5	6	7
GPU3				FWD	BWD
GPU2			FWD			BWD
GPU1		FWD					BWD
GPU0	FWD							BWD

可以看到，在某一時(shí)刻僅有1個(gè)GPU工作。并且每個(gè)timesep花費(fèi)的時(shí)間也比較長，因?yàn)镚PU需要跑完整個(gè)minibatch的前向傳播。

GPipe將minibatch劃分為4個(gè)microbatch，然后依次送入GPU0。GPU0前向傳播后，再將結(jié)果送入GPU1，以此類推。整個(gè)過程如下表

Timestep	0	1	2	3	4	5	6	7	8	9	10	11	12	13
GPU3				F1	F2	F3	F4	B4	B3	B2	B1
GPU2			F1	F2	F3	F4			B4	B3	B2	B1
GPU1		F1	F2	F3	F4					B4	B3	B2	B1
GPU0	F1	F2	F3	F4							B4	B3	B2	B1

F1表示使用當(dāng)前GPU上的層來對(duì)microbatch1進(jìn)行前向傳播。在GPipe的調(diào)度中，每個(gè)timestep上花費(fèi)的時(shí)間要比樸素層并行更短，因?yàn)槊總€(gè)GPU僅需要處理microbatch。

2.2 GPipe的Bubbles問題

bubbles指的是流水線中沒有進(jìn)行任何有效工作的點(diǎn)。這是由于操作之間的依賴導(dǎo)致的。例如，在GPU3執(zhí)行完F1之前，GPU4只能等待。整個(gè)流水線過程中的bubbles如下圖所示。

bubbles浪費(fèi)時(shí)間的比例依賴于pipeline的深度和mincrobatch的數(shù)量。假設(shè)單個(gè)GPU上完成前向傳播或者后向傳播的面積為 1 (也就是上圖中的單個(gè)小方塊面積為 1 )。上圖中的總長度為, 寬度為, 總面積為。其中, 彩色小方塊占用的面積表示GPU執(zhí)行的時(shí)間, 為其。空白處面積的占比代表了浪費(fèi)時(shí)間的比較, 其值為

因此，增大microbatch的數(shù)量m，可以降低bubbles的比例。

2.3 GPipe的顯存需求

增大batch size就會(huì)線性增大需要被緩存激活的顯存需求。在GPipe中，GPU需要在前向傳播至反向傳播這段時(shí)間內(nèi)緩存激活(activations)。以GPU0為例，microbatch1的激活需要從timestep 0保存至timestep 13。

GPipe為了解決顯存的問題，使用了gradient checkpointing。該技術(shù)不需要緩存所有的激活，而是在反向傳播的過程中重新計(jì)算激活。這降低了對(duì)顯存的需求，但是增加了計(jì)算代價(jià)。

假設(shè)所有層都大致相等。每個(gè)GPU緩存激活所需要的顯存為

也就是與單個(gè)GPU上的層數(shù)以及batch size成正比。相反，使用gradient checkpointing僅需要緩存邊界層(需要跨GPU發(fā)送結(jié)果的層)的輸入。這可以降低每個(gè)GPU的顯存峰值需求

(batchsize) 是緩存邊界激活所需要的顯存。當(dāng)對(duì)給定的microbatch執(zhí)行反向傳播時(shí), 需要重新計(jì)算該microbatch梯度所需要的激活。對(duì)于每個(gè)GPU上的層需要的顯存空間。

3. PipeDream

GPipe需要等所有的microbatch前向傳播完成后，才會(huì)開始反向傳播。PipeDream則是當(dāng)一個(gè)microbatch的前向傳播完成后，立即進(jìn)入反向傳播階段。理論上，反向傳播完成后就可以丟棄掉對(duì)應(yīng)microbatch緩存的激活。由于PipeDream的反向傳播完成的要比GPipe早，因此也會(huì)減少顯存的需求。

下圖是PipeDream的調(diào)度圖，4個(gè)GPU和8個(gè)microbatchs。藍(lán)色的方塊表示前向傳播，綠色表示反向傳播，數(shù)字則是microbatch的id。

PipeDream在bubbles上與GPipe沒有區(qū)別，但是由于PipeDream釋放顯存的時(shí)間更早，因此會(huì)降低對(duì)顯存的需求。

4. 合并數(shù)據(jù)并行和流水線并行

數(shù)據(jù)并行和流水線并行是正交的，可以同時(shí)使用。

• 對(duì)于流水線并行。每個(gè)GPU需要與下個(gè)流水線階段(前向傳播)或者上個(gè)流水線階段(反向傳播)進(jìn)行通信。
• 對(duì)于數(shù)據(jù)并行。每個(gè)GPU需要與分配了相同層的GPU進(jìn)行通信。所有層的副本需要AllReduce對(duì)梯度進(jìn)行平均。

這將在所有GPU上形成子組，并在子組中使用集合通信。任意給定的GPU都會(huì)有兩部分的通信，一個(gè)是包含所有相同層的GPU(數(shù)據(jù)并行)，另一個(gè)與不同層的GPU(流水線并行)。下圖是流水線并行度為2且數(shù)據(jù)并行度為2的示例圖，水平方向是完整的一個(gè)模型，垂直方向是相同層的不同副本。

二、張量并行

Transformer中的主要部件是全連接層和注意力機(jī)制，其核心都是矩陣乘法。張量并行的核心就是將矩陣乘法進(jìn)行拆分，從而降低模型對(duì)單卡的顯存需求。

1. 1D張量并行

本小節(jié)以全鏈接層為例, 介紹張量并行。其中, 和是輸入和輸出向量,是權(quán)重矩陣,是非線性激活函數(shù)。總量來說張量并行可以分為列并行和行并行(以權(quán)重矩陣的分割方式命名), 上圖展示了兩種并行。

(1) 矩陣乘法角度

這里以矩陣乘法的方式輔助理解1D張量并行。

• 列并行

將矩陣行列劃分為n份(不一定必須相等大小)可以表示為，那么矩陣乘 法表示為

顯然，僅需要對(duì)權(quán)重進(jìn)行劃分。

• 行并行

對(duì)權(quán)重進(jìn)行劃分，那么必須對(duì)輸入矩陣也進(jìn)行劃分。假設(shè)要將A水平劃分為n份，則輸入矩陣X必須垂直劃分為n份，那么矩陣乘法表示為

(2) 激活函數(shù)與通信

顯然，只觀察上面的數(shù)據(jù)公式，無論是行并行還是列并行 ，都只需要在各個(gè)部分計(jì)算完后進(jìn)行一次通常。只不過列并行將通信的結(jié)果進(jìn)行拼接，而行并行則是對(duì)通信結(jié)果相加。

現(xiàn)在，我們將非線性激活GeLU加上，并模擬兩層的全鏈接層。設(shè)X是輸入，A和B則是兩個(gè)全鏈接層的權(quán)重。

• 列并行

通過上面的公式可以看到。當(dāng)我們將A和B提前劃分好后，就可以獨(dú)立進(jìn)行計(jì)算，在計(jì)算出后再進(jìn)行通信。也就是說，這個(gè)例子中雖然有兩個(gè)全鏈接層，但是僅需要在得到最終結(jié)果前進(jìn)行通信即可。

所以，多個(gè)全鏈接層堆疊時(shí)，僅需要在最終輸出時(shí)進(jìn)行一次通信即可(如上圖所示)。

• 行并行

由于是非線性的，所以

因此，行并行每一個(gè)全鏈接層都需要進(jìn)行通信來聚合最終的結(jié)果。

(3) 多頭注意力并行

多頭注意力并行不是1D張量并行，但是由于其是Megatron-LM中與1D張量并行同時(shí)提出的，所以這里也進(jìn)行簡(jiǎn)單的介紹。

由于多頭注意力的各個(gè)頭之間本質(zhì)上就是獨(dú)立的，因此各個(gè)頭完全可以并行運(yùn)算。

注意：張量并行(TP)需要非常快的網(wǎng)絡(luò)，因此不建議跨多個(gè)節(jié)點(diǎn)進(jìn)行張量并行。實(shí)際中，若一個(gè)節(jié)點(diǎn)有4個(gè)GPU，最高的張量并行度為4。

2. 2D、2.5D張量并行

在1D張量并行后，又逐步提出了2D、2.5D和3D張量并行。這里對(duì)2D和2.5D張量并行進(jìn)行簡(jiǎn)單介紹：

(1) 2D張量并行

1D張量并行并沒有對(duì)激活(activations，也就是模型中間層的輸出結(jié)果)進(jìn)行劃分，導(dǎo)致消耗大量的顯存。

這里仍然以矩陣乘法為例, 給定有個(gè)處理器。這里假設(shè), 則將X和A都劃分為的塊。

基于上面的矩陣乘法的變化，可以發(fā)現(xiàn)Y=XA可以分解為兩個(gè)矩陣相加。具體來說，兩個(gè)矩陣的結(jié)果仍然需要串行的計(jì)算。但是，單個(gè)矩陣中的4個(gè)子矩陣可以使用的處理器來并行計(jì)算。

當(dāng), 也就是第一步。對(duì)

進(jìn)行廣播, 所有的處理器均擁有這 4 個(gè)子矩陣。然后分別執(zhí)行。經(jīng)過這一步后就得到了第一個(gè)矩陣的結(jié)果。

當(dāng)。對(duì)

進(jìn)行廣播，各個(gè)處理器在分別計(jì)算。最終得到第二個(gè)矩陣的結(jié)果。

將兩個(gè)矩陣的結(jié)果相加。

(2) 2.5D張量并行

仍然是矩陣乘法, 并假設(shè)有個(gè)處理器, 將 X劃分為行和列。不妨設(shè), 將和分別劃分為

那么有

其中， concat 表示兩個(gè)矩陣的垂直拼接操作。

基于上面的推導(dǎo), 可以發(fā)現(xiàn)被拼接的兩個(gè)矩陣天然可以并行計(jì)算。即和可以并行計(jì)算。看到這里, 應(yīng)該就可以發(fā)現(xiàn)這兩個(gè)矩陣乘法就是上面的2D張量并行適用的形式。所以, 我們總計(jì)有個(gè)處理器, 每個(gè)處理器使用2D張量并行來處理對(duì)應(yīng)的矩陣乘法。最后, 將兩個(gè)2D張量并行的結(jié)果進(jìn)行拼接即可。

三、3D并行

總的來說，3D并行是由數(shù)據(jù)并行(DP)、張量并行(TP)和流水線并行(PP)組成。前面已經(jīng)分別介紹了TP和PP，ZeRO-DP是一種顯存高效的數(shù)據(jù)并行策略，原理見文章：

白強(qiáng)偉：【深度學(xué)習(xí)】【分布式訓(xùn)練】DeepSpeed：AllReduce與ZeRO-DP142

https://zhuanlan.zhihu.com/p/610587671

下面介紹如何將三種并行技術(shù)結(jié)合在一起，形成3D并行技術(shù)。

1. 一個(gè)3D并行的例子

假設(shè)有兩個(gè)節(jié)點(diǎn)Node1和Node2，每個(gè)節(jié)點(diǎn)有8個(gè)GPU，共計(jì)16個(gè)GPU。16個(gè)GPU的編號(hào)分別為Rank0、Rank1、...、Rank15。此外，假設(shè)用戶設(shè)置流水線并行度為4，張量并行度為2。

流水線并行。流水線并行會(huì)將整個(gè)模型劃分為4份，這里稱為sub_model_1至sub_model_4。每連續(xù)的4張GPU負(fù)責(zé)一個(gè)sub_model。即上圖右上角中，相同顏色的GPU負(fù)責(zé)相同的sub_model。

張量并行。張量并行會(huì)針對(duì)流水線并行中的sub_model來進(jìn)行張量的拆分。即Rank0和Rank1負(fù)責(zé)一份sub_model_1，Rank2和Rank3負(fù)責(zé)另一份sub_model_1；Rank4和Rank5負(fù)責(zé)sub_model_2，Rank6和Rank7負(fù)責(zé)另一份sub_model_2；以此類推。上圖右下角中，綠色線條表示單個(gè)張量并行組，每個(gè)張量并行組都共同負(fù)責(zé)某個(gè)具體的sub_model。

數(shù)據(jù)并行。數(shù)據(jù)并行的目的是要保證并行中的相同模型參數(shù)讀取相同的數(shù)據(jù)。經(jīng)過流水線并行和張量并行后，Rank0和Rank2負(fù)責(zé)相同的模型參數(shù)，所以Rank0和Rank2是同一個(gè)數(shù)據(jù)并行組。上圖左上角中的紅色線條表示數(shù)據(jù)并行組。

2. 3D并行分析

為什么3D并行需要按上面的方式劃分GPU呢？首先，模型并行是三種策略中通信開銷最大的，所以優(yōu)先將模型并行組放置在一個(gè)節(jié)點(diǎn)中，以利用較大的節(jié)點(diǎn)內(nèi)帶寬。其次，流水線并行通信量最低，因此在不同節(jié)點(diǎn)之間調(diào)度流水線，這將不受通信帶寬的限制。最后，若張量并行沒有跨節(jié)點(diǎn)，則數(shù)據(jù)并行也不需要跨節(jié)點(diǎn)；否則數(shù)據(jù)并行組也需要跨節(jié)點(diǎn)。

流水線并行和張量并行減少了單個(gè)顯卡的顯存消耗，提高了顯存效率。但是，模型劃分的太多會(huì)增加通信開銷，從而降低計(jì)算效率。ZeRO-DP不僅能夠通過將優(yōu)化器狀態(tài)進(jìn)行劃分來改善顯存效率，而且還不會(huì)顯著的增加通信開銷。