近年來，隨著Transformer、MOE架構(gòu)的提出，使得深度學(xué)習(xí)模型輕松突破上萬億規(guī)模參數(shù)，傳統(tǒng)的單機(jī)單卡模式已經(jīng)無法滿足超大模型進(jìn)行訓(xùn)練的要求。因此，我們需要基于單機(jī)多卡、甚至是多機(jī)多卡進(jìn)行分布式大模型的訓(xùn)練。

而利用AI集群，使深度學(xué)習(xí)算法更好地從大量數(shù)據(jù)中高效地訓(xùn)練出性能優(yōu)良的大模型是分布式機(jī)器學(xué)習(xí)的首要目標(biāo)。為了實(shí)現(xiàn)該目標(biāo)，一般需要根據(jù)硬件資源與數(shù)據(jù)/模型規(guī)模的匹配情況，考慮對(duì)計(jì)算任務(wù)、訓(xùn)練數(shù)據(jù)和模型進(jìn)行劃分，從而進(jìn)行分布式存儲(chǔ)和分布式訓(xùn)練。因此，分布式訓(xùn)練相關(guān)技術(shù)值得我們進(jìn)行深入分析其背后的機(jī)理。

下面主要對(duì)大模型進(jìn)行分布式訓(xùn)練的并行技術(shù)進(jìn)行講解，本系列大體分九篇文章進(jìn)行講解。

大模型分布式訓(xùn)練并行技術(shù)（一）-概述

大模型分布式訓(xùn)練并行技術(shù)（二）-數(shù)據(jù)并行

大模型分布式訓(xùn)練并行技術(shù)（三）-流水線并行

大模型分布式訓(xùn)練并行技術(shù)（四）-張量并行

大模型分布式訓(xùn)練并行技術(shù)（五）-序列并行

大模型分布式訓(xùn)練并行技術(shù)（六）-多維混合并行

大模型分布式訓(xùn)練并行技術(shù)（七）-自動(dòng)并行

大模型分布式訓(xùn)練并行技術(shù)（八）-MOE并行

大模型分布式訓(xùn)練并行技術(shù)（九）-總結(jié)

本文為分布式訓(xùn)練并行技術(shù)的第一篇，對(duì)大模型進(jìn)行分布式訓(xùn)練常見的并行技術(shù)進(jìn)行簡要介紹。

數(shù)據(jù)并行

數(shù)據(jù)并行是最常見的并行形式，因?yàn)樗芎唵?。在?shù)據(jù)并行訓(xùn)練中，數(shù)據(jù)集被分割成幾個(gè)碎片，每個(gè)碎片被分配到一個(gè)設(shè)備上。這相當(dāng)于沿批次（Batch）維度對(duì)訓(xùn)練過程進(jìn)行并行化。每個(gè)設(shè)備將持有一個(gè)完整的模型副本，并在分配的數(shù)據(jù)集碎片上進(jìn)行訓(xùn)練。在反向傳播之后，模型的梯度將被全部減少，以便在不同設(shè)備上的模型參數(shù)能夠保持同步。典型的數(shù)據(jù)并行實(shí)現(xiàn)：PyTorch DDP。

image.png

模型并行

在數(shù)據(jù)并行訓(xùn)練中，一個(gè)明顯的特點(diǎn)是每個(gè) GPU 持有整個(gè)模型權(quán)重的副本。這就帶來了冗余問題。另一種并行模式是模型并行，即模型被分割并分布在一個(gè)設(shè)備陣列上。

通常有兩種類型的模型并行：張量并行和流水線并行。

張量并行是在一個(gè)操作中進(jìn)行并行計(jì)算，如：矩陣-矩陣乘法。

流水線并行是在各層之間進(jìn)行并行計(jì)算。

因此，從另一個(gè)角度來看，張量并行可以被看作是層內(nèi)并行，流水線并行可以被看作是層間并行。

張量并行

張量并行訓(xùn)練是將一個(gè)張量沿特定維度分成 N 塊，每個(gè)設(shè)備只持有整個(gè)張量的 1/N，同時(shí)不影響計(jì)算圖的正確性。這需要額外的通信來確保結(jié)果的正確性。

以一般的矩陣乘法為例，假設(shè)我們有 C = AB。我們可以將B沿著列分割成[B0 B1 B2 ... Bn]，每個(gè)設(shè)備持有一列。然后我們將 A 與每個(gè)設(shè)備上 B 中的每一列相乘，我們將得到[AB0 AB1 AB2 ... ABn]。此刻，每個(gè)設(shè)備仍然持有一部分的結(jié)果，例如，設(shè)備(rank=0)持有 AB0。為了確保結(jié)果的正確性，我們需要收集全部的結(jié)果，并沿列維串聯(lián)張量。通過這種方式，我們能夠?qū)埩糠植荚谠O(shè)備上，同時(shí)確保計(jì)算流程保持正確。

image.png

典型的張量并行實(shí)現(xiàn)：Megatron-LM（1D）、Colossal-AI（2D、2.5D、3D）。

流水線并行

流水線并行的核心思想是，模型按層分割成若干塊，每塊都交給一個(gè)設(shè)備。

在前向傳播過程中，每個(gè)設(shè)備將中間的激活傳遞給下一個(gè)階段。

在后向傳播過程中，每個(gè)設(shè)備將輸入張量的梯度傳回給前一個(gè)流水線階段。

這允許設(shè)備同時(shí)進(jìn)行計(jì)算，從而增加訓(xùn)練的吞吐量。

image.png

流水線并行訓(xùn)練的一個(gè)明顯缺點(diǎn)是訓(xùn)練設(shè)備容易出現(xiàn)空閑狀態(tài)（因?yàn)楹笠粋€(gè)階段需要等待前一個(gè)階段執(zhí)行完畢），導(dǎo)致計(jì)算資源的浪費(fèi)，加速效率沒有數(shù)據(jù)并行高。

image.png

典型的流水線并行實(shí)現(xiàn)：GPipe、PipeDream、PipeDream-2BW、PipeDream Flush（1F1B）。

優(yōu)化器相關(guān)的并行

目前隨著模型越來越大，單個(gè)GPU的顯存目前通常無法裝下那么大的模型了。那么就要想辦法對(duì)占顯存的地方進(jìn)行優(yōu)化。

通常來說，模型訓(xùn)練的過程中，GPU上需要進(jìn)行存儲(chǔ)的參數(shù)包括了模型本身的參數(shù)、優(yōu)化器狀態(tài)、激活函數(shù)的輸出值、梯度以及一些零時(shí)的Buffer。各種數(shù)據(jù)的占比如下圖所示：

image.png

可以看到模型參數(shù)僅占模型訓(xùn)練過程中所有數(shù)據(jù)的一部分，當(dāng)進(jìn)行混合精度運(yùn)算時(shí)，其中模型狀態(tài)參數(shù)(優(yōu)化器狀態(tài)+ 梯度+ 模型參數(shù)）占到了一大半以上。因此，我們需要想辦法去除模型訓(xùn)練過程中的冗余數(shù)據(jù)。

而優(yōu)化器相關(guān)的并行就是一種去除冗余數(shù)據(jù)的并行方案，目前這種并行最流行的方法是 ZeRO（即零冗余優(yōu)化器）。針對(duì)模型狀態(tài)的存儲(chǔ)優(yōu)化（去除冗余），ZeRO使用的方法是分片，即每張卡只存 1/N 的模型狀態(tài)量，這樣系統(tǒng)內(nèi)只維護(hù)一份模型狀態(tài)。ZeRO有三個(gè)不同級(jí)別，對(duì)模型狀態(tài)進(jìn)行不同程度的分片：

ZeRO-1 : 對(duì)優(yōu)化器狀態(tài)分片（Optimizer States Sharding）

ZeRO-2 : 對(duì)優(yōu)化器狀態(tài)和梯度分片（Optimizer States & Gradients Sharding）

ZeRO-3 : 對(duì)優(yōu)化器狀態(tài)、梯度分片以及模型權(quán)重參數(shù)分片（Optimizer States & Gradients & Parameters Sharding）

image.png

異構(gòu)系統(tǒng)并行

上述的方法中，通常需要大量的 GPU 來訓(xùn)練一個(gè)大型模型。然而，人們常常忽略一點(diǎn)，與 GPU 相比，CPU 的內(nèi)存要大得多。在一個(gè)典型的服務(wù)器上，CPU 可以輕松擁有幾百GB甚至上TB的內(nèi)存，而每張 GPU 卡通常只有 48 或 80 GB的內(nèi)存。這促使人們思考為什么 CPU 內(nèi)存沒有被用于分布式訓(xùn)練。

而最近的進(jìn)展是依靠 CPU 甚至是 NVMe 磁盤來訓(xùn)練大型模型。主要的想法是，在不使用張量時(shí)，將其卸載回 CPU 內(nèi)存或 NVMe 磁盤。

通過使用異構(gòu)系統(tǒng)架構(gòu)，有可能在一臺(tái)機(jī)器上容納一個(gè)巨大的模型。

image.png

多維混合并行

多維混合并行指將數(shù)據(jù)并行、模型并行和流水線并行等多種并行技術(shù)結(jié)合起來進(jìn)行分布式訓(xùn)練。

image.png

通常，在進(jìn)行超大規(guī)模模型的預(yù)訓(xùn)練和全參數(shù)微調(diào)時(shí)，都需要用到多維混合并行。

image.png

為了充分利用帶寬，通常情況下，張量并行所需的通信量最大，而數(shù)據(jù)并行與流水線并行所需的通信量相對(duì)來說較小。因此，同一個(gè)服務(wù)器內(nèi)使用張量并行，而服務(wù)器之間使用數(shù)據(jù)并行與流水線并行。

image.png

自動(dòng)并行

上面提到的數(shù)據(jù)并行、張量并行、流水線并行等多維混合并行需要把模型切分到多張AI加速卡上面，如果讓用戶手動(dòng)實(shí)現(xiàn)，對(duì)開發(fā)者來說難度非常大，需要考慮性能、內(nèi)存、通信、訓(xùn)練效果等問題，要是能夠?qū)⒛Ｐ桶此阕踊蛘甙磳幼詣?dòng)切分到不同的加速卡上，可以大大的降低開發(fā)者的使用難度。因此，自動(dòng)并行應(yīng)運(yùn)而生。

image.png

MOE并行 / 專家并行

通常來講，模型規(guī)模的擴(kuò)展會(huì)導(dǎo)致訓(xùn)練成本顯著增加，計(jì)算資源的限制成為了大規(guī)模密集模型訓(xùn)練的瓶頸。為了解決這個(gè)問題，一種基于稀疏 MoE 層的深度學(xué)習(xí)模型架構(gòu)被提出，即將大模型拆分成多個(gè)小模型(專家，expert)， 每輪迭代根據(jù)樣本決定激活一部分專家用于計(jì)算，達(dá)到了節(jié)省計(jì)算資源的效果；并引入可訓(xùn)練并確保稀疏性的門(gate)機(jī)制，以保證計(jì)算能力的優(yōu)化。

使用 MoE 結(jié)構(gòu)，可以在計(jì)算成本次線性增加的同時(shí)實(shí)現(xiàn)超大規(guī)模模型訓(xùn)練，為恒定的計(jì)算資源預(yù)算帶來巨大增益。而 MOE 并行，本質(zhì)上也是一種模型并行方法。下圖展示了一個(gè)有六個(gè)專家網(wǎng)絡(luò)的模型被兩路專家并行地訓(xùn)練。其中，專家1-3被放置在第一個(gè)計(jì)算單元上，而專家4-6被放置在第二個(gè)計(jì)算單元上。

image.png

結(jié)語

本文針對(duì)大模型進(jìn)行分布式訓(xùn)練常見的并行技術(shù)進(jìn)行了簡要的介紹。后續(xù)章節(jié)將針對(duì)常見并行技術(shù)的不同方案進(jìn)行詳細(xì)的講解。