2020 年，最轟動的 AI 新聞莫過于 OpenAI 發(fā)布的 GPT-3 了。它的1750億參數(shù)量及其在眾多NLP任務(wù)上超過人類的出眾表現(xiàn)讓大家堅信：大模型才是未來。但與之帶來的問題是，訓練超大模型所需的算力、存儲已不再是單機就能搞定的了（之前的 BERT 還是可以用 DGX-1/2 這樣的超級服務(wù)器訓練）。

NVIDIA 估算過，如果要訓練GPT-3 ，即使單個機器的顯存/內(nèi)存能裝得下，用 8 張 V100 的顯卡（一臺 DGX-1 的配置），訓練時長預(yù)計要 36 年；即使用 512 張 V100 ，訓練也需要將近 7 個月；如果你擁有 1024 張 80GB A100， 那么完整訓練 GPT-3 的時長可以縮減到 1 個月。

這意味著訓練大模型一定是一個分布式問題。對算力的需求還是一個相對容易解決的問題，因為擁有大集群的組織并不只 OpenAI 一家，而如何解決上千塊 GPU 的分布式訓練問題才是關(guān)鍵。

即使你是一位非常優(yōu)秀的數(shù)據(jù)科學家，知曉并能解決 Transformer 相關(guān)的所有算法問題，但如果你不是分布式專家，不知道如何解決分布式訓練時上百臺服務(wù)器之間的通信、拓撲、模型并行、流水并行等問題，你甚至都無法啟動這次訓練。這也解釋了為什么時隔一年，只有 NVIDIA 、微軟等大企業(yè)可以復(fù)現(xiàn) GPT-3 。

目前開源的 GPT 模型庫主要是 NVIDIA 的 Megatron-LM 和微軟的 DeepSpeed。其中，微軟的 DeepSpeed 的模型并行等內(nèi)核取自 Megatron，且 DeepSpeed 主打的是，在數(shù)據(jù)并行下如何以更少的機器去跑更大的模型 （ ZeRO 、 ZeRO-Offload 等都是用梯度切片、計算、內(nèi)存/硬盤換入換出來省顯存），所以我們本文主要介紹和對比 Megatron 。

簡單比較一下 NVIDIA 的 Megatron 和 微軟的 DeepSpeed：

DeepSpeed 本質(zhì)上是一種“節(jié)省顯存”的數(shù)據(jù)并行，即：數(shù)據(jù)并行的優(yōu)化版。DeepSpeed 假設(shè)了單層參數(shù)量可以在單張顯卡上放得下，如果不滿足這個假設(shè)，那么仍然需要使用模型并行，而且 DeepSpeed 的模型并行是通過調(diào)用 Megatron 來實現(xiàn)的。

根據(jù) NVIDIA 最新的那篇論文（鏈接：https://arxiv.org/abs/2104.04473，也是下面本文重點要介紹的），Megatron 在大規(guī)模訓練的效率是超過 DeepSpeed 不少的。

而 DeepSpeed 的論文一直強調(diào)：可以用更少機器訓練更大的模型，但沒有突出過在效率上的優(yōu)勢。

DeepSpeed 后來又出了一篇論文：ZeRO-Infinity（鏈接：https://arxiv.org/abs/2104.07857），當單層參數(shù)量在單張顯卡上放不下的時候，它通過對這一層算子切片，一片一片來執(zhí)行，使得單卡也能跑起來一個巨大的層，可以理解成一種 “時間”軸上展開的模型并行。

Megatron 和 DeepSpeed 都是基于 PyTorch ，分別由 NVIDIA 和微軟經(jīng)過深度定制開發(fā)，專門為支持 PyTorch 分布式訓練 GPT 而設(shè)計的。我們會簡單介紹一下 NVIDIA 如何使用 PyTorch 搞分布式訓練 GPT ，然后重點介紹 OneFlow 如何用一套通用設(shè)計非常簡單清晰地解決了這個難題，同時我們還在已有的測試規(guī)模上性能超過 NVIDIA。

相信讀完此文，你就會發(fā)現(xiàn) PyTorch 、 Megatron ( NVIDIA ) 、DeepSpeed ( Microsoft ) 都走了一個非常長的彎路，這條彎路從大方向上就走錯了，不僅是彎路，你還會發(fā)現(xiàn) Megatron 的代碼只能被 NVIDIA 的分布式訓練專家所復(fù)用，它對于 PyTorch 的算法工程師而言門檻極高，是非常難用的，以至于任何想要用 PyTorch 復(fù)現(xiàn)一個分布式大模型的算法工程師，都得先等 NVIDIA 開發(fā)完才能再使用 Megatron 提供的模型。同時，我們也通過這樣一個例子來證明：為什么一個分布式深度學習框架要像 OneFlow 這樣設(shè)計。

本文內(nèi)容較多，先列出主要目錄：

1.分布式訓練 GPT 的必要并行技術(shù)

流水并行

梯度累加

后向重計算

1F1B 策略

2.Megatron：PyTorch 分布式訓練的極限，痛點在哪兒？

流水并行 PyTorch 需要手寫專家級復(fù)雜調(diào)度器

模型并行 PyTorch 需要手工排線，在 kernel 里手寫特定的、經(jīng)過復(fù)雜推導的通信原語

3.OneFlow 用一致性視角（Consistent View）輕松填平分布式訓練難的鴻溝

流水并行 ，只需要配置 Placement 就夠了

數(shù)據(jù) + 模型的混合并行，只需要配置 Variable 的 SBP 就夠了

OneFlow：讓每一位算法工程師都有能力訓練 GPT

4.為什么分布式深度學習框架要像 OneFlow 這樣設(shè)計？

OneFlow 系統(tǒng)層面如何實現(xiàn)流水并行

OneFlow 系統(tǒng)層面如何實現(xiàn)混合（數(shù)據(jù) & 模型）并行

5.GPT 訓練性能對比: OneFlow vs Megatron

純數(shù)據(jù)并行性能對比

純模型并行性能對比

混合并行性能對比

流水并行 + 混合并行 性能對比

6.小結(jié)

分布式訓練 GPT 的必要并行技術(shù)

最近，NVIDIA 放出了一篇重量級的論文：Efficient Large-Scale Language Model Training on GPU Clusters ，用了 3072 張 80GB A100 訓練 GPT（ NVIDIA 也確實夠壕，這個集群的成本就不止 5 億了），最大規(guī)模的模型參數(shù)量達到了 1T（是 GPT-3 原版的 5 倍）。

NVIDIA 訓練 GPT-3 最大到 1T 參數(shù)規(guī)模

論文里 NVIDIA 介紹了分布式訓練超大規(guī)模模型的三種必須的并行技術(shù)：

• 數(shù)據(jù)并行（Data Parallelism）

• 模型并行（Tensor Model Parallelism）

• 流水并行（Pipeline Model Parallelism）

其中數(shù)據(jù)并行是大家都熟知的最常見的并行方式，而模型并行（NVIDIA 論文里叫做 "Tensor 級別的模型并行" ）是對某一層（如 Linear/Dense Layer 里的 Variable ）的模型 Tensor 切分，從而將大的模型 Tensor 分成多個相對較小的 Tensor 進行并行計算；流水并行（NVIDIA 論文里叫做流水線級別的模型并行），是將整個網(wǎng)絡(luò)分段（stage），不同段在不同的設(shè)備上，前后階段流水分批工作，通過一種“接力”的方式并行。

對于最大的 1T 規(guī)模的模型，NVIDIA 一共使用了 384 臺 DGX-A100 機器（每臺裝有 8 張 80GB A100 GPU），機器內(nèi)部各 GPU 間使用超高速 NVLink 和 NVSwitch 互聯(lián)，每臺機器裝有 8個 200Gbps 的 InfiniBand (IB) 網(wǎng)卡，可以說是硬件集群頂配中的頂配了。

那么，這些機器是如何協(xié)同工作的？GPT 網(wǎng)絡(luò)是由很多層 Transformer Layer 組成的，每一層內(nèi)部是一個由多層 MLP 和 attention 機制組成的子圖，對于參數(shù)規(guī)模 1T 的 GPT 而言就有 128 層的 Transformer Layer，這個超大超深的網(wǎng)絡(luò)被分割成了 64 個 stage ，每個 stage 跑在 6 臺 DGX-A100 上，其中 6 臺機器之間進行數(shù)據(jù)并行，每臺機器內(nèi)部的 8 張卡之間做模型并行，整個集群的 3072 張 A100 按照機器拓撲被劃分成了 [6 x 8 x 64] 的矩陣，同時使用數(shù)據(jù)并行 & 模型并行 & 流水并行 進行訓練。

3072 張 A100 集群拓撲

1.流水并行

從上述的機器拓撲中可以發(fā)現(xiàn)，流水并行是 3072 塊 A100 能訓練 GPT 的關(guān)鍵。因為無論是數(shù)據(jù)并行還是模型并行，都會在相應(yīng)的機器之間進行全連接的通信，當機器數(shù)量增大時，通信開銷和時延會大到難以忍受。而流水并行既解決了超大模型無法在單設(shè)備上裝下的難題，又很好解決了機器之間的通信開銷的問題，每個階段（stage） 和下一個階段之間僅有相鄰的某一個 Tensor 數(shù)據(jù)需要傳輸，每臺機器的數(shù)據(jù)傳輸量跟總的網(wǎng)絡(luò)大小、機器總數(shù)、并行規(guī)模無關(guān)。但流水并行為了多個階段之間可以流水起來，還依賴兩個重要的特性： 梯度累加（Gradient Accumulation） 和 亞線性內(nèi)存優(yōu)化（ Sublinear Memory Cost 2016, 陳天奇）。

近期，百度和華為相繼發(fā)了自己的千億級中文預(yù)訓練模型的宣傳文。其中，百度提出了 "4D混合并行"，本質(zhì)上是 Megatron 里的數(shù)據(jù)并行 + 模型并行 + 流水并行 + DeepSpeed 里的 ZeRO 優(yōu)化 ；華為文章中的 “5D混合并行”，是將重計算（Checkpointing, 亞線性內(nèi)存優(yōu)化的一種）作為了第5維 （其實百度也做了重計算，只是沒有將其列為多維并行中的一維）。

在介紹這兩個特性之前，我們先簡單解釋一下深度學習訓練和模型更新的兩種約束：BSP (Bulk Synchronous Parallel) 和 SSP (Stale Synchronous Parallel ) ，其中 BSP 是最常見的模型更新規(guī)則：每個 batch 的前向計算都需要使用最新的模型，那么就要求上一個 batch 的后向計算結(jié)束且在模型更新后，下一個 batch 的前向才能開始。如果使用 BSP 去做流水并行，我們就會發(fā)現(xiàn)每個階段的前向和后向是完全串行的，其中一個設(shè)備在工作時，其他所有設(shè)備都在等待，那么分布式的優(yōu)勢就完全沒有被發(fā)揮出來：

BSP 各個階段串行執(zhí)行

BSP 且沒有 Gradient Accumulation 下的流水并行。假設(shè)整個網(wǎng)絡(luò)被等分切成 4 個 stage，每個 stage 使用一個 device ，則在BSP下，各個設(shè)備串行執(zhí)行，中間有大段的氣泡。一般后向計算時間是前向計算的兩倍，如果算上 Checkpointing 的重計算部分，是前向計算的三倍。我們可以從上圖中看到，這種情況下有 70% 的時間設(shè)備是空閑的。

而 SSP 就是異步模型更新，允許前向計算時可以不使用最新的模型，而使用落后幾個版本之內(nèi)的模型。SSP 在 GPT-3 的訓練中并沒有被 NVIDIA 采用，其主要原因有以下幾點：

SSP 的模型收斂性并沒有被嚴格論證， 且有論文 GeePS 指出 SSP 的收斂效果不如 BSP ；

SSP 會在 GPU 上同時出現(xiàn)幾個不同版本的模型，而 GPT-3 又是一個非常大的模型網(wǎng)絡(luò)，多份模型所帶來的顯存開銷不可接受；

BSP 通過 Gradient Accumulation + Checkpointing 就可以很好的解決 GPT-3 中的流水并行問題。

另外， NVIDIA 的論文的分析前提就是 BSP 情況下， 根據(jù)嚴格的參數(shù)優(yōu)化器更新方式， 流水并行氣泡的占比是 Bubble time fraction = (p - 1) / m，其中 p 是 stage 數(shù)， m 是梯度累加的 micro-batch 數(shù)。如果采用 SSP，則 NVIDIA 整篇文章的理論基礎(chǔ)就沒有了。

"Pipeline parallelism comes in a few flavors: the mode discussed in this paper uses flushes to ensure exact strict optimizer semantics."

2.梯度累加

Gradient Accumulation 就是把一個大 Batch 拆分成多個 micro-batch ， 每個 micro-batch 前后向計算后的梯度累加，在最后一個micro-batch累加結(jié)束后，統(tǒng)一更新模型。

micro-batch 跟數(shù)據(jù)并行有高度的相似性：數(shù)據(jù)并行是空間上的， 數(shù)據(jù)被拆分成多個 tensor，同時喂給多個設(shè)備并行計算，然后將梯度累加在一起更新；而 micro-batch 是時間上的數(shù)據(jù)并行， 數(shù)據(jù)被拆分成多個 tensor， 按照時序依次進入同一個設(shè)備串行計算，然后將梯度累加在一起更新。當總的 batch size 一致，且數(shù)據(jù)并行的并行度和 micro-batch 的累加次數(shù)相等時，數(shù)據(jù)并行和 Gradient Accumulation 在數(shù)學上完全等價。Gradient Accumulation 通過多個 micro-batch的梯度累加使得下一個 micro-batch 的前向計算不需要依賴上一個 micro-batch 的反向計算，因此可以暢通無阻的進行下去（當然在一個大 batch 的最后一次 micro-batch 還是會觸發(fā)這個依賴）。

Gradient Accumulation 解決了很多問題：

在單卡下，Gradient Accumulation 可以將一個大的 batch size 拆分成等價的多個小 micro-batch ，從而達到節(jié)省顯存的目的。

在數(shù)據(jù)并行下，Gradient Accumulation 解決了反向梯度同步開銷占比過大的問題（隨著機器數(shù)和設(shè)備數(shù)的增加，梯度的 AllReduce 同步開銷也加大），因為梯度同步變成了一個稀疏操作，因此可以提升數(shù)據(jù)并行的加速比。

在流水并行下， Gradient Accumulation 使得不同 stage 之間可以并行執(zhí)行不同的 micro-batch， 從而讓各個階段的計算不阻塞，達到流水的目的。

單純通過 micro-batch，我們就實現(xiàn)了 GPipe （2018）論文中的流水并行，在 stage 數(shù)量為 4， micro-batch 數(shù)量為 8 （每個 batch 在計算 8 個 micro-batch 且 累加 8 次梯度后更新一次模型）下的時間線如下圖所示：

使用梯度累加后的 Pipeline 時間線

在 GPipe 的流水并行示例中，每個“時間點” 可以在多個階段（stage）上同時做不同的micro-batch，圖中每個方塊中的標號表示了第幾個 micro-batch；同一個 micro-batch 還是串行的經(jīng)過所有的 stage，在這種情況下，每個設(shè)備的空閑時間只有 25% 左右。

但這種流水并行存在一個問題：顯存占用太大。如果每個 micro-batch 前向計算的中間結(jié)果（activation）被后向計算所消費，則需要在顯存中緩存 8 份（梯度累加的次數(shù)）完整的前向 activation。這時就不得不用另一項重要的技術(shù)：重計算（Checkpointing）。