一文看懂AI大模型的并行訓(xùn)練方式（DP、PP、TP、EP）

大家都知道，AI計算（尤其是模型訓(xùn)練和推理），主要以并行計算為主。

AI計算中涉及到的很多具體算法（例如矩陣相乘、卷積、循環(huán)層、梯度運算等），都需要基于成千上萬的GPU，以并行任務(wù)的方式去完成。這樣才能有效縮短計算時間。

搭建并行計算框架，一般會用到以下幾種常見的并行方式：

Data Parallelism，數(shù)據(jù)并行

Pipeline Parallelism，流水線并行

Tensor Parallelism，張量并行

Expert Parallelism, 專家并行

接下來，我們逐一看看，這些并行計算方式的工作原理。

▉ DP（數(shù)據(jù)并行）

首先看看DP，數(shù)據(jù)并行（Data Parallelism）。

AI訓(xùn)練使用的并行，總的來說，分為數(shù)據(jù)并行和模型并行兩類。剛才說的PP（流水線并行）、TP（張量并行）和EP（專家并行），都屬于模型并行，待會再介紹。

這里，我們需要先大概了解一下神經(jīng)網(wǎng)絡(luò)的訓(xùn)練過程。簡單來說，包括以下主要步驟：

1、前向傳播：輸入一批訓(xùn)練數(shù)據(jù)，計算得到預(yù)測結(jié)果。

2、計算損失：通過損失函數(shù)比較預(yù)測結(jié)果與真實標(biāo)簽的差距。

3、反向傳播：將損失值反向傳播，計算網(wǎng)絡(luò)中每個參數(shù)的梯度。

4、梯度更新：優(yōu)化器使用這些梯度來更新所有的權(quán)重和偏置（更新參數(shù)）。

以上過程循環(huán)往復(fù)，直到模型的性能達到令人滿意的水平。訓(xùn)練就完成了。

我們回到數(shù)據(jù)并行。

數(shù)據(jù)并行是大模型訓(xùn)練中最為常見的一種并行方式（當(dāng)然，也適用于推理過程）。

它的核心思想很簡單，就是每個GPU都擁有完整的模型副本，然后，將訓(xùn)練數(shù)據(jù)劃分成多個小批次（mini-batch），每個批次分配給不同的GPU進行處理。

數(shù)據(jù)并行的情況下，大模型訓(xùn)練的過程是這樣的：

1、對數(shù)據(jù)進行均勻切割，發(fā)給不同的、并行工作的GPU（Worker）；

2、各GPU都擁有一樣的模型以及模型參數(shù)，它們各自獨立進行前向傳播、反向傳播，計算得到各自的梯度；

3、各GPU通過卡間通信，以All-Reduce的通信方式，將梯度推給一個類似管理者的GPU（Server）；

4、Server GPU對所有梯度進行求和或者平均，得到全局梯度；

5、Server GPU將全局梯度回傳（broadcast廣播）到每個Worker GPU，進行參數(shù)更新（更新本地模型權(quán)重）。更新后，所有worker GPU模型參數(shù)保持一致。

然后，再繼續(xù)重復(fù)這樣的過程，直至完成所有的訓(xùn)練。

再來一張圖，幫助理解：

從下往上看

這里提到的All-Reduce，也是一個AI領(lǐng)域的常見概念，字面意思是“全（All）-規(guī)約（Reduce）”，即：對所有節(jié)點的數(shù)據(jù)進行聚合（如求和、求最大值），并將最終結(jié)果分發(fā)到所有節(jié)點。（參考：）

數(shù)據(jù)并行的優(yōu)點，在于實現(xiàn)過程比較簡單，能夠顯著加速大規(guī)模數(shù)據(jù)的訓(xùn)練過程，尤其適用于數(shù)據(jù)量遠大于模型參數(shù)的場景。

數(shù)據(jù)并行的缺點，在于顯存的限制。因為每個GPU上都有完整的模型副本，而當(dāng)模型的規(guī)模和參數(shù)越大，所需要的顯存就越大，很可能超過單個GPU的顯存大小。

數(shù)據(jù)并行的通信開銷也比較大。不同GPU之間需要頻繁通信，以同步模型參數(shù)或梯度。而且，模型參數(shù)規(guī)模越大，GPU數(shù)量越多，這個通信開銷就越大。例如，對于千億參數(shù)模型，單次梯度同步需傳輸約2TB數(shù)據(jù)（FP16精度下）。

▉ ZeRO

這里要插播介紹一個概念——ZeRO（Zero Redundancy Optimizer，零冗余優(yōu)化器）。

在數(shù)據(jù)并行策略中，每個GPU的內(nèi)存都保存一個完整的模型副本，很占內(nèi)存空間。那么，能否每個GPU只存放模型副本的一部分呢？

沒錯，這就是ZeRo——通過對模型副本中的優(yōu)化器狀態(tài)、梯度和參數(shù)進行切分，來實現(xiàn)減少對內(nèi)存的占用。

ZeRO有3個階段，分別是：

ZeRO-1：對優(yōu)化器狀態(tài)進行劃分。

ZeRO-2：對優(yōu)化器狀態(tài)和梯度進行劃分

ZeRO-3：對優(yōu)化器狀態(tài)、梯度和參數(shù)進行劃分。（最節(jié)省顯存）

通過下面的圖和表，可以看得更明白些：

根據(jù)實測數(shù)據(jù)顯示，ZeRO-3在1024塊GPU上訓(xùn)練萬億參數(shù)模型時，顯存占用從7.5TB降至7.3GB/卡。

值得一提的是，DP還有一個DDP（分布式數(shù)據(jù)并行）。傳統(tǒng)DP一般用于單機多卡場景。而DDP能多機也能單機。這依賴于Ring-AllReduce，它由百度最先提出，可以有效解決數(shù)據(jù)并行中通信負載不均（Server存在瓶頸）的問題。

▉ PP（流水線并行）

再來看看模型并行。

剛才數(shù)據(jù)并行，是把數(shù)據(jù)分為好幾個部分。模型并行，很顯然，就是把模型分為好幾個部分。不同的GPU，運行不同的部分。（注意：業(yè)界對模型并行的定義有點混亂。也有的資料會將張量并行等同于模型并行。）

流水線并行，是將模型的不同層（單層，或連續(xù)的多層）分配到不同的GPU上，按順序處理數(shù)據(jù)，實現(xiàn)流水線式的并行計算。

例如，對于一個包含7層的神經(jīng)網(wǎng)絡(luò)，將1~2層放在第一個GPU上，3~5層放在第二個GPU上，6~7層放在第三個GPU上。訓(xùn)練時，數(shù)據(jù)按照順序，在不同的GPU上進行處理。

乍一看，流水并行有點像串行。每個GPU需要等待前一個GPU的計算結(jié)果，可能會導(dǎo)致大量的GPU資源浪費。

上面這個圖中，黃色部分就是Bubble （氣泡）時間。氣泡越多，代表GPU處于等待狀態(tài)（空閑狀態(tài)）越長，資源浪費越嚴重。

為了解決上述問題，可以將mini-batch的數(shù)據(jù)進一步切分成micro-batch數(shù)據(jù)。當(dāng)GPU 0處理完一個micro-batch數(shù)據(jù)后，緊接著開始處理下一個micro-batch數(shù)據(jù)，以此來減少GPU的空閑時間。如下圖（b）所示：

還有，在一個micro-batch完成前向計算后，提前調(diào)度，完成相應(yīng)的反向計算，這樣就能釋放部分顯存，用以接納新的數(shù)據(jù)，提升整體訓(xùn)練性能。如上圖（c）所示。

這些方法，都能夠顯著減少流水線并行的Bubble時間。

對于流水線并行，需要對任務(wù)調(diào)度和數(shù)據(jù)傳輸進行精確管理，否則可能導(dǎo)致流水線阻塞，以及產(chǎn)生更多的Bubble時間。

▉ TP（張量并行）

模型并行的另外一種，是張量并行。

如果說流水線并行是將一個模型按層「垂直」分割，那么，張量并行則是在一個層內(nèi)「橫向」分割某些操作。

具體來說，張量并行是將模型的張量（如權(quán)重矩陣）按維度切分到不同的GPU上運行的并行方式。

張量切分方式分為按行進行切分和按列進行切分，分別對應(yīng)行并行（Row Parallelism）（權(quán)重矩陣按行分割）與列并行（Column Parallelism）（權(quán)重矩陣按列分割）。

每個節(jié)點處理切分后的子張量。最后，通過集合通信操作（如All-Gather或All-Reduce）來合并結(jié)果。

張量并行的優(yōu)點，是適合單個張量過大的情況，可以顯著減少單個節(jié)點的內(nèi)存占用。

張量并行的缺點，是當(dāng)切分維度較多的時候，通信開銷比較大。而且，張量并行的實現(xiàn)過程較為復(fù)雜，需要仔細設(shè)計切分方式和通信策略。

放一張數(shù)據(jù)并行、流水線并行、張量并行的簡單對比：

▉ 專家并行

2025年初DeepSeek爆紅的時候，有一個詞也跟著火了，那就是MoE（Mixture of Experts，混合專家模型）。

MoE模型的核心是“多個專家層+路由網(wǎng)絡(luò)（門控網(wǎng)絡(luò)）”。

專家層的每個專家負責(zé)處理特定類型的token（如語法、語義相關(guān)）。路由網(wǎng)絡(luò)根據(jù)輸入token的特征，選擇少數(shù)專家處理這個token，其他專家不激活。

MoE實現(xiàn)了任務(wù)分工、按需分配算力，因此大幅提升了模型效率。

專家并行（Expert Parallelism），是MoE（混合專家模型）中的一種并行計算策略。它通過將專家（子模型）分配到不同的GPU上，實現(xiàn)計算負載的分布式處理，提高計算效率。

專家并行與之前所有的并行相比，最大的不同在于，輸入數(shù)據(jù)需要通過一個動態(tài)的路由選擇機制分發(fā)給相應(yīng)專家，此處會涉及到一個所有節(jié)點上的數(shù)據(jù)重分配的動作。

然后，在所有專家處理完成后，又需要將分散在不同節(jié)點上的數(shù)據(jù)按原來的次序整合起來。

這樣的跨片通信模式，稱為All-to-All。（再次參考：）

專家并行可能存在負載不均衡的問題。某個專家所接收到的輸入數(shù)據(jù)大于了其所能接收的范圍，就可能導(dǎo)致Tokens不被處理或不能被按時處理，成為瓶頸。

所以，設(shè)計合理的門控機制和專家選擇策略，是部署專家并行的關(guān)鍵。

▉ 混合并行

在實際應(yīng)用中，尤其是訓(xùn)練萬億參數(shù)級別的超大模型時，幾乎不會只使用單一的并行策略，而是采用多維度的混合并行（結(jié)合使用多種并行策略）。

例如：

數(shù)據(jù)并行+張量并行：數(shù)據(jù)并行處理批量樣本，張量并行處理單樣本的大矩陣計算。

流水線并行+專家并行：流水線并行劃分模型層，專家并行劃分層內(nèi)專家模塊。

更高級的，是3D并行，通過“數(shù)據(jù)并行+張量并行+流水線并行”，實現(xiàn)三重拆分，是超大模型訓(xùn)練的主流方案。

3D并行

▉ 最后的話

好啦，以上就是關(guān)于DP、PP、TP、EP等并行訓(xùn)練方式的介紹。大家都看懂了沒？

并行計算方式其實非常復(fù)雜，剛才我們只是做了最簡單的介紹。但在真實工作中，開發(fā)者無需了解具體的實現(xiàn)細節(jié)，因為業(yè)界提供了例如DeepSpeed（微軟開源，支持3D并行+ZeRO內(nèi)存優(yōu)化）、Megatron-LM（NVIDIA開源，3D并行的標(biāo)桿）、FSDP等開源軟件，能夠讓開發(fā)者直接進行大語言模型訓(xùn)練。

小棗君之所以要專門介紹并行訓(xùn)練方式，其實更多是為了幫助大家深入地理解算力集群架構(gòu)和網(wǎng)絡(luò)的設(shè)計。

大家可以看到，不同的并行訓(xùn)練方式，有著不同的通信流量特點。算力集群整體架構(gòu)和網(wǎng)絡(luò)設(shè)計，需要盡量去適配這些并行計算方式的流量特點，才能滿足模型訓(xùn)推任務(wù)的要求，實現(xiàn)更高的工作效率。

比如說，數(shù)據(jù)并行，由于需要頻繁同步梯度信息，對網(wǎng)絡(luò)帶寬要求較高，需要確保網(wǎng)絡(luò)帶寬能夠滿足大量梯度數(shù)據(jù)快速傳輸?shù)男枨螅苊庖驇挷蛔銓?dǎo)致通信延遲，影響訓(xùn)練效率。

流水線并行，大模型的每一段，在不同的服務(wù)器上以流水線的方式逐步計算，涉及到多個服務(wù)器“串起來”，就建議部署在比較靠近的服務(wù)器上（盡量部署在葉脊網(wǎng)絡(luò)的同一個leaf葉下）。

張量并行，通信數(shù)據(jù)量大，就建議部署在一臺服務(wù)器的多個GPU上進行計算。

專家并行中，不同專家分配在不同GPU上，GPU間需要交換中間計算結(jié)果等信息，其通信流量特點取決于專家的數(shù)量以及數(shù)據(jù)交互的頻率等，也需要合理規(guī)劃GPU間的連接方式和通信路徑。

總之，在GPU算卡性能越來越難以提升的背景下，深入研究并行計算的設(shè)計，從架構(gòu)和網(wǎng)絡(luò)上挖掘潛力，是業(yè)界的必然選擇。

隨著AI浪潮的繼續(xù)發(fā)展，以后是否還會出現(xiàn)其它的并行訓(xùn)練方式呢？讓我們拭目以待吧！