在上一篇的介紹中，我們介紹了以Google GPipe為代表的流水線并行范式。當模型太大，一塊GPU放不下時，流水線并行將模型的不同層放到不同的GPU上，通過切割mini-batch實現對訓練數據的流水線處理，提升GPU計算通訊比。同時通過re-materialization機制降低顯存消耗。

但在實際應用中，流水線并行并不特別流行，主要原因是模型能否均勻切割，影響了整體計算效率，這就需要算法工程師做手調。因此，今天我們來介紹一種應用最廣泛，最易于理解的并行范式：數據并行。

數據并行的核心思想是：在各個GPU上都拷貝一份完整模型，各自吃一份數據，算一份梯度，最后對梯度進行累加來更新整體模型。理念不復雜，但到了大模型場景，巨大的存儲和GPU間的通訊量，就是系統設計要考慮的重點了。在本文中，我們將遞進介紹三種主流數據并行的實現方式：

DP（Data Parallelism）：最早的數據并行模式，一般采用參數服務器(Parameters Server)這一編程框架。實際中多用于單機多卡

DDP（Distributed Data Parallelism）：分布式數據并行，采用Ring AllReduce的通訊方式，實際中多用于多機場景

ZeRO：零冗余優化器。由微軟推出并應用于其DeepSpeed框架中。嚴格來講ZeRO采用數據并行+張量并行的方式，旨在降低存儲。

本文將首先介紹DP和DDP，在下一篇文章里，介紹ZeRO。全文內容如下：

1、數據并行（DP）

1.1 整體架構

1.2 通訊瓶頸與梯度異步更

2、分布式數據并行（DDP）

2.1 圖解Ring-AllReduce

2.2 DP與DDP通訊分析

推薦閱讀： 圖解大模型訓練之：流水線并行，以GPipe為例

一、數據并行（DP）

1.1 整體架構

一個經典數據并行的過程如下：

若干塊計算GPU，如圖中GPU0~GPU2；1塊梯度收集GPU，如圖中AllReduce操作所在GPU。

在每塊計算GPU上都拷貝一份完整的模型參數。

把一份數據X（例如一個batch）均勻分給不同的計算GPU。

每塊計算GPU做一輪FWD和BWD后，算得一份梯度G。

每塊計算GPU將自己的梯度push給梯度收集GPU，做聚合操作。這里的聚合操作一般指梯度累加。當然也支持用戶自定義。

梯度收集GPU聚合完畢后，計算GPU從它那pull下完整的梯度結果，用于更新模型參數W。更新完畢后，計算GPU上的模型參數依然保持一致。

聚合再下發梯度的操作，稱為AllReduce。

前文說過，實現DP的一種經典編程框架叫“參數服務器”，在這個框架里，計算GPU稱為Worker，梯度聚合GPU稱為Server。在實際應用中，為了盡量減少通訊量，一般可選擇一個Worker同時作為Server。比如可把梯度全發到GPU0上做聚合。需要再額外說明幾點：

1個Worker或者Server下可以不止1塊GPU。

Server可以只做梯度聚合，也可以梯度聚合+全量參數更新一起做在參數服務器的語言體系下，DP的過程又可以被描述下圖：

1.2 通訊瓶頸與梯度異步更新

DP的框架理解起來不難，但實戰中確有兩個主要問題：

存儲開銷大。每塊GPU上都存了一份完整的模型，造成冗余。關于這一點的優化，我們將在后文ZeRO部分做講解。

通訊開銷大。Server需要和每一個Worker進行梯度傳輸。當Server和Worker不在一臺機器上時，Server的帶寬將會成為整個系統的計算效率瓶頸。

我們對通訊開銷再做詳細說明。如果將傳輸比作一條馬路，帶寬就是馬路的寬度，它決定每次并排行駛的數據量。例如帶寬是100G/s，但每秒卻推給Server 1000G的數據，消化肯定需要時間。那么當Server在搬運數據，計算梯度的時候，Worker們在干嘛呢？當然是在：

人類老板不愿意了：“打工系統里不允許有串行存在的任務！”，于是梯度異步更新這一管理層略誕生了。

上圖刻畫了在梯度異步更新的場景下，某個Worker的計算順序為：

在第10輪計算中，該Worker正常計算梯度，并向Server發送push&pull梯度請求。

但是，該Worker并不會實際等到把聚合梯度拿回來，更新完參數W后再做計算。而是直接拿舊的W，吃新的數據，繼續第11輪的計算。這樣就保證在通訊的時間里，Worker也在馬不停蹄做計算，提升計算通訊比。

當然，異步也不能太過份。只計算梯度，不更新權重，那模型就無法收斂。圖中刻畫的是延遲為1的異步更新，也就是在開始第12輪對的計算時，必須保證W已經用第10、11輪的梯度做完2次更新了。

參數服務器的框架下，延遲的步數也可以由用戶自己決定，下圖分別刻劃了幾種延遲情況：

(a) 無延遲

(b) 延遲但不指定延遲步數。也即在迭代2時，用的可能是老權重，也可能是新權重，聽天由命。

(c) 延遲且指定延遲步數為1。例如做迭代3時，可以不拿回迭代2的梯度，但必須保證迭代0、1的梯度都已拿回且用于參數更新。

總結一下，異步很香，但對一個Worker來說，只是等于W不變，batch的數量增加了而已，在SGD下，會減慢模型的整體收斂速度。異步的整體思想是，比起讓Worker閑著，倒不如讓它多吃點數據，雖然反饋延遲了，但只要它在干活在學習就行。

batch就像活，異步就像畫出去的餅，且往往不指定延遲步數，每個Worker干越來越多的活，但模型卻沒收斂取效。讀懂分布式訓練系統其實也不難。

二、分布式數據并行(DDP)

受通訊負載不均的影響，DP一般用于單機多卡場景。因此，DDP作為一種更通用的解決方案出現了，既能多機，也能單機。DDP首先要解決的就是通訊問題：將Server上的通訊壓力均衡轉到各個Worker上。實現這一點后，可以進一步去Server，留Worker。

前文我們說過，聚合梯度 + 下發梯度這一輪操作，稱為AllReduce。接下來我們介紹目前最通用的AllReduce方法：Ring-AllReduce。它由百度最先提出，非常有效地解決了數據并行中通訊負載不均的問題，使得DDP得以實現。

2.1 Ring-AllReduce

如下圖，假設有4塊GPU，每塊GPU上的數據也對應被切成4份。AllReduce的最終目標，就是讓每塊GPU上的數據都變成箭頭右邊匯總的樣子。

Ring-ALLReduce則分兩大步驟實現該目標：Reduce-Scatter和All-Gather。

Reduce-Scatter

定義網絡拓撲關系，使得每個GPU只和其相鄰的兩塊GPU通訊。每次發送對應位置的數據進行累加。每一次累加更新都形成一個拓撲環，因此被稱為Ring。看到這覺得困惑不要緊，我們用圖例把詳細步驟畫出來。

一次累加完畢后，藍色位置的數據塊被更新，被更新的數據塊將成為下一次更新的起點，繼續做累加操作。

3次更新之后，每塊GPU上都有一塊數據擁有了對應位置完整的聚合（圖中紅色）。此時，Reduce-Scatter階段結束。進入All-Gather階段。目標是把紅色塊的數據廣播到其余GPU對應的位置上。

All-Gather

如名字里Gather所述的一樣，這操作里依然按照“相鄰GPU對應位置進行通訊”的原則，但對應位置數據不再做相加，而是直接替換。All-Gather以紅色塊作為起點。

以此類推，同樣經過3輪迭代后，使得每塊GPU上都匯總到了完整的數據，變成如下形式：

建議讀者們手動推一次，加深理解(注：最后一圖箭頭畫錯，公眾號不許修改

2.2 Ring-AllReduce通訊量分析

假設模型參數W的大小為，GPU個數為。則梯度大小也為，每個梯度塊的大小為。

對單卡GPU來說：

Reduce-Scatter階段，通訊量為

All-Gather階段，通訊量為

總通訊量為，隨著N的增大，可以近似

而對前文的DP來說，它的Server承載的總通訊量也是。雖然通訊量相同，但搬運相同數據量的時間卻不一定相同。DDP把通訊量均衡負載到了每一時刻的每個Worker上，而DP僅讓Server做勤勞的搬運工。當越來越多的GPU分布在距離較遠的機器上時，DP的通訊時間是會增加的。

但這并不說明參數服務器不能打（有很多文章將參數服務器當作old dinosaur來看）。事實上，參數服務器也提供了多Server方法，如下圖：

在多Server的模式下，進一步，每個Server可以只負責維護和更新某一塊梯度（也可以某塊梯度+參數一起維護），此時雖然每個Server仍然需要和所有Worker通訊，但它的帶寬壓力會小非常多。經過調整設計后，依然可以用來做DDP。雖然這篇文章是用遞進式的方式來介紹兩者，但不代表兩者間一定要決出優劣。我想表達的觀點是，方法是多樣性的。對參數服務器有興趣的朋友，可以閱讀參考的第1個鏈接。

最后，請大家記住Ring-AllReduce的方法，因為在之后的ZeRO，Megatron-LM中，它將頻繁地出現，是分布式訓練系統中重要的算子。

三、總結

1、在DP中，每個GPU上都拷貝一份完整的模型，每個GPU上處理batch的一部分數據，所有GPU算出來的梯度進行累加后，再傳回各GPU用于更新參數。

2、DP多采用參數服務器這一編程框架，一般由若個計算Worker和1個梯度聚合Server組成。Server與每個Worker通訊，Worker間并不通訊。因此Server承擔了系統所有的通訊壓力。基于此DP常用于單機多卡場景。

3、異步梯度更新是提升計算通訊比的一種方法，延遲更新的步數大小決定了模型的收斂速度。

4、Ring-AllReduce通過定義網絡環拓撲的方式，將通訊壓力均衡地分到每個GPU上，使得跨機器的數據并行（DDP）得以高效實現。

5、DP和DDP的總通訊量相同，但因負載不均的原因，DP需要耗費更多的時間搬運數據。