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

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

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

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

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

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

DeepSpeed 本質上是一種“節省顯存”的數據并行，即：數據并行的優化版。DeepSpeed 假設了單層參數量可以在單張顯卡上放得下，如果不滿足這個假設，那么仍然需要使用模型并行，而且 DeepSpeed 的模型并行是通過調用 Megatron 來實現的。

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

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

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

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

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

本文內容較多，先列出主要目錄：

1.分布式訓練 GPT 的必要并行技術

流水并行

梯度累加

后向重計算

1F1B 策略

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

流水并行 PyTorch 需要手寫專家級復雜調度器

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

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

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

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

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

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

OneFlow 系統層面如何實現流水并行

OneFlow 系統層面如何實現混合（數據 & 模型）并行

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

純數據并行性能對比

純模型并行性能對比

混合并行性能對比

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

6.小結

分布式訓練 GPT 的必要并行技術

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

NVIDIA 訓練 GPT-3 最大到 1T 參數規模

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

• 數據并行（Data Parallelism）

• 模型并行（Tensor Model Parallelism）

• 流水并行（Pipeline Model Parallelism）

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

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

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

3072 張 A100 集群拓撲

1.流水并行

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

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

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

BSP 各個階段串行執行

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

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

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

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

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

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

"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累加結束后，統一更新模型。

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

Gradient Accumulation 解決了很多問題：

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

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

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

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

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

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

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