大家都知道,AI計算(尤其是模型訓練和推理),主要以并行計算為主。
AI計算中涉及到的很多具體算法(例如矩陣相乘、卷積、循環層、梯度運算等),都需要基于成千上萬的GPU,以并行任務的方式去完成。這樣才能有效縮短計算時間。
搭建并行計算框架,一般會用到以下幾種常見的并行方式:
Data Parallelism,數據并行
Pipeline Parallelism,流水線并行
Tensor Parallelism,張量并行
Expert Parallelism, 專家并行
接下來,我們逐一看看,這些并行計算方式的工作原理。
DP(數據并行)
首先看看DP,數據并行(Data Parallelism)。
AI訓練使用的并行,總的來說,分為數據并行和模型并行兩類。剛才說的PP(流水線并行)、TP(張量并行)和EP(專家并行),都屬于模型并行,待會再介紹。

這里,我們需要先大概了解一下神經網絡的訓練過程。簡單來說,包括以下主要步驟:

1. 前向傳播:輸入一批訓練數據,計算得到預測結果。
2. 計算損失:通過損失函數比較預測結果與真實標簽的差距。
3. 反向傳播:將損失值反向傳播,計算網絡中每個參數的梯度。
4. 梯度更新:優化器使用這些梯度來更新所有的權重和偏置(更新參數)。
以上過程循環往復,直到模型的性能達到令人滿意的水平。訓練就完成了。
我們回到數據并行。
數據并行是大模型訓練中最為常見的一種并行方式(當然,也適用于推理過程)。
它的核心思想很簡單,就是每個GPU都擁有完整的模型副本,然后,將訓練數據劃分成多個小批次(mini-batch),每個批次分配給不同的GPU進行處理。
數據并行的情況下,大模型訓練的過程是這樣的:

1、對數據進行均勻切割,發給不同的、并行工作的GPU(Worker);
2、各GPU都擁有一樣的模型以及模型參數,它們各自獨立進行前向傳播、反向傳播,計算得到各自的梯度;
3、各GPU通過卡間通信,以All-Reduce的通信方式,將梯度推給一個類似管理者的GPU(Server);
4、Server GPU對所有梯度進行求和或者平均,得到全局梯度;
5、Server GPU將全局梯度回傳(broadcast廣播)到每個Worker GPU,進行參數更新(更新本地模型權重)。更新后,所有worker GPU模型參數保持一致。
然后,再繼續重復這樣的過程,直至完成所有的訓練。
再來一張圖,幫助理解:

從下往上看
這里提到的All-Reduce,也是一個AI領域的常見概念,字面意思是“全(All)-規約(Reduce)”,即:對所有節點的數據進行聚合(如求和、求最大值),并將最終結果分發到所有節點。(參考:到底什么是All-Reduce、All-to-All?)
數據并行的優點,在于實現過程比較簡單,能夠顯著加速大規模數據的訓練過程,尤其適用于數據量遠大于模型參數的場景。
數據并行的缺點,在于顯存的限制。因為每個GPU上都有完整的模型副本,而當模型的規模和參數越大,所需要的顯存就越大,很可能超過單個GPU的顯存大小。
數據并行的通信開銷也比較大。不同GPU之間需要頻繁通信,以同步模型參數或梯度。而且,模型參數規模越大,GPU數量越多,這個通信開銷就越大。例如,對于千億參數模型,單次梯度同步需傳輸約2TB數據(FP16精度下)。
ZeRO
這里要插播介紹一個概念——ZeRO(Zero Redundancy Optimizer,零冗余優化器)。
在數據并行策略中,每個GPU的內存都保存一個完整的模型副本,很占內存空間。那么,能否每個GPU只存放模型副本的一部分呢?
沒錯,這就是ZeRo——通過對模型副本中的優化器狀態、梯度和參數進行切分,來實現減少對內存的占用。
ZeRO有3個階段,分別是:
ZeRO-1:對優化器狀態進行劃分。
ZeRO-2:對優化器狀態和梯度進行劃分
ZeRO-3:對優化器狀態、梯度和參數進行劃分。(最節省顯存)
通過下面的圖和表,可以看得更明白些:


根據實測數據顯示,ZeRO-3在1024塊GPU上訓練萬億參數模型時,顯存占用從7.5TB降至7.3GB/卡。
值得一提的是,DP還有一個DDP(分布式數據并行)。傳統DP一般用于單機多卡場景。而DDP能多機也能單機。這依賴于Ring-AllReduce,它由百度最先提出,可以有效解決數據并行中通信負載不均(Server存在瓶頸)的問題。

PP(流水線并行)
再來看看模型并行。
剛才數據并行,是把數據分為好幾個部分。模型并行,很顯然,就是把模型分為好幾個部分。不同的GPU,運行不同的部分。(注意:業界對模型并行的定義有點混亂。也有的資料會將張量并行等同于模型并行。)
流水線并行,是將模型的不同層(單層,或連續的多層)分配到不同的GPU上,按順序處理數據,實現流水線式的并行計算。

例如,對于一個包含7層的神經網絡,將1~2層放在第一個GPU上,3~5層放在第二個GPU上,6~7層放在第三個GPU上。訓練時,數據按照順序,在不同的GPU上進行處理。
乍一看,流水并行有點像串行。每個GPU需要等待前一個GPU的計算結果,可能會導致大量的GPU資源浪費。

上面這個圖中,黃色部分就是Bubble (氣泡)時間。氣泡越多,代表GPU處于等待狀態(空閑狀態)越長,資源浪費越嚴重。
為了解決上述問題,可以將mini-batch的數據進一步切分成micro-batch數據。當GPU 0處理完一個micro-batch數據后,緊接著開始處理下一個micro-batch數據,以此來減少GPU的空閑時間。如下圖(b)所示:

還有,在一個micro-batch完成前向計算后,提前調度,完成相應的反向計算,這樣就能釋放部分顯存,用以接納新的數據,提升整體訓練性能。如上圖(c)所示。
這些方法,都能夠顯著減少流水線并行的Bubble時間。
對于流水線并行,需要對任務調度和數據傳輸進行精確管理,否則可能導致流水線阻塞,以及產生更多的Bubble時間。
TP(張量并行)
模型并行的另外一種,是張量并行。
如果說流水線并行是將一個模型按層「垂直」分割,那么,張量并行則是在一個層內「橫向」分割某些操作。

具體來說,張量并行是將模型的張量(如權重矩陣)按維度切分到不同的GPU上運行的并行方式。
張量切分方式分為按行進行切分和按列進行切分,分別對應行并行(Row Parallelism)(權重矩陣按行分割)與列并行(Column Parallelism)(權重矩陣按列分割)。

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

張量并行的優點,是適合單個張量過大的情況,可以顯著減少單個節點的內存占用。
張量并行的缺點,是當切分維度較多的時候,通信開銷比較大。而且,張量并行的實現過程較為復雜,需要仔細設計切分方式和通信策略。
放一張數據并行、流水線并行、張量并行的簡單對比:

專家并行
2025年初DeepSeek爆紅的時候,有一個詞也跟著火了,那就是MoE(Mixture of Experts,混合專家模型)。
MoE模型的核心是“多個專家層+路由網絡(門控網絡)”。

專家層的每個專家負責處理特定類型的token(如語法、語義相關)。路由網絡根據輸入token的特征,選擇少數專家處理這個token,其他專家不激活。
MoE實現了任務分工、按需分配算力,因此大幅提升了模型效率。
專家并行(Expert Parallelism),是MoE(混合專家模型)中的一種并行計算策略。它通過將專家(子模型)分配到不同的GPU上,實現計算負載的分布式處理,提高計算效率。
專家并行與之前所有的并行相比,最大的不同在于,輸入數據需要通過一個動態的路由選擇機制分發給相應專家,此處會涉及到一個所有節點上的數據重分配的動作。
然后,在所有專家處理完成后,又需要將分散在不同節點上的數據按原來的次序整合起來。
這樣的跨片通信模式,稱為All-to-All。(再次參考:到底什么是All-Reduce、All-to-All?)
專家并行可能存在負載不均衡的問題。某個專家所接收到的輸入數據大于了其所能接收的范圍,就可能導致Tokens不被處理或不能被按時處理,成為瓶頸。
所以,設計合理的門控機制和專家選擇策略,是部署專家并行的關鍵。
混合并行
在實際應用中,尤其是訓練萬億參數級別的超大模型時,幾乎不會只使用單一的并行策略,而是采用多維度的混合并行(結合使用多種并行策略)。
例如:
數據并行+張量并行:數據并行處理批量樣本,張量并行處理單樣本的大矩陣計算。
流水線并行+專家并行:流水線并行劃分模型層,專家并行劃分層內專家模塊。
更高級的,是3D并行,通過“數據并行+張量并行+流水線并行”,實現三重拆分,是超大模型訓練的主流方案。

3D并行
最后的話
好啦,以上就是關于DP、PP、TP、EP等并行訓練方式的介紹。大家都看懂了沒?

并行計算方式其實非常復雜,剛才我們只是做了最簡單的介紹。但在真實工作中,開發者無需了解具體的實現細節,因為業界提供了例如DeepSpeed(微軟開源,支持3D并行+ZeRO內存優化)、Megatron-LM(NVIDIA開源,3D并行的標桿)、FSDP等開源軟件,能夠讓開發者直接進行大語言模型訓練。
小棗君之所以要專門介紹并行訓練方式,其實更多是為了幫助大家深入地理解算力集群架構和網絡的設計。
大家可以看到,不同的并行訓練方式,有著不同的通信流量特點。算力集群整體架構和網絡設計,需要盡量去適配這些并行計算方式的流量特點,才能滿足模型訓推任務的要求,實現更高的工作效率。
比如說,數據并行,由于需要頻繁同步梯度信息,對網絡帶寬要求較高,需要確保網絡帶寬能夠滿足大量梯度數據快速傳輸的需求,避免因帶寬不足導致通信延遲,影響訓練效率。
流水線并行,大模型的每一段,在不同的服務器上以流水線的方式逐步計算,涉及到多個服務器“串起來”,就建議部署在比較靠近的服務器上(盡量部署在葉脊網絡的同一個leaf葉下)。
張量并行,通信數據量大,就建議部署在一臺服務器的多個GPU上進行計算。
專家并行中,不同專家分配在不同GPU上,GPU間需要交換中間計算結果等信息,其通信流量特點取決于專家的數量以及數據交互的頻率等,也需要合理規劃GPU間的連接方式和通信路徑。

總之,在GPU算卡性能越來越難以提升的背景下,深入研究并行計算的設計,從架構和網絡上挖掘潛力,是業界的必然選擇。
隨著AI浪潮的繼續發展,以后是否還會出現其它的并行訓練方式呢?讓我們拭目以待吧!
文章來源于鮮棗課堂,作者小棗君
-
TP
+關注
關注
0文章
81瀏覽量
31456 -
DP
+關注
關注
1文章
241瀏覽量
42235 -
pp
+關注
關注
0文章
8瀏覽量
8500 -
ep
+關注
關注
0文章
16瀏覽量
16498 -
AI大模型
+關注
關注
0文章
398瀏覽量
998
發布評論請先 登錄
基于Transformer做大模型預訓練基本的并行范式
圖解大模型訓練之:數據并行上篇(DP, DDP與ZeRO)
基于PyTorch的模型并行分布式訓練Megatron解析
一文看懂AI大模型的并行訓練方式(DP、PP、TP、EP)
評論