一、背景

本文中我們繼續(xù)介紹分布式訓(xùn)練中的張量并行（Tensor Parallelism，TP），包括 AlexNet、Google、Facebook、NVIDIA 以及 Colossal-AI 的一系列 Tensor Parallelism 方案。涉及 1D TP，2D 和 3D TP，也包含行切分和列切分等。

這里說的 Tensor Parallelism 和 Zero DP 以及 Pytorch FSDP 中的模型切分方式不一樣，Zero DP 和 FSDP 中的模型切分在實(shí)際使用的時(shí)候還會(huì)將相應(yīng)的參數(shù) AllGather 到當(dāng)前設(shè)備，使用全部的參數(shù)進(jìn)行計(jì)算。而 Tensor Parallelism 中的參數(shù)都始終在當(dāng)前設(shè)備，最終聚合的是結(jié)果（Activation）。當(dāng)然，在后續(xù)的文章中我們也會(huì)介紹 Zero 和 FSDP 相關(guān)方案。

二、分布式矩陣乘

Tensor Parallelism 的核心就是分布式矩陣乘法，其在高性能計(jì)算領(lǐng)域已經(jīng)有很長(zhǎng)的歷史，也已被廣泛研究。在 Tensor Parallelism 中有兩種常見的切分方式，Column Parallelism 和 Row Parallelism。如果從模型的的角度考慮，通常指的是 Y=XW 的形式，其中 X 為輸入，W 為權(quán)重參數(shù)，Y 為輸出。而具體是哪種切分方式也可以看 W 矩陣是在哪個(gè)維度切分的。（PS：這個(gè)小節(jié)的圖片來自 Tensor Parallelism — PyTorch Lightning 2.4.0 documentation）

2.1 Column Parallelism

如下圖所示為 Column Parallelism，其中的 Column 就是指權(quán)重參數(shù) W 按照 Column 維度切分。每個(gè) GPU 都包含一部分權(quán)重參數(shù)，并使用整個(gè)輸入 X 計(jì)算，得到 Y 的一部分，最后通過 AllGather 操作可以獲得全量結(jié)果。

2.2 Row Parallelism

如下圖所示為 Row Parallelism，其中的 Row 就是指權(quán)重參數(shù) W 按照 Row 維度切分。每個(gè) GPU 都包含一部分權(quán)重參數(shù)，并使用部分輸入 X 計(jì)算，結(jié)果和 Y 的 Shape 相同，但結(jié)果不完整，最后通過 AllReduce 操作可以獲得全量結(jié)果。因?yàn)?AllReduce 可以通過 ReduceScatter 和 AllGather 的方式實(shí)現(xiàn)，而 Column Parallelism 中的 AllGather 和 Row Parallelism 中 AllGather 通信量是一樣的，因此，總體來說 Column Parallelism 的通信量更少：

2.3 Column Parallelism + Row Parallelism

在 Transformer 等模型中會(huì)存在連續(xù)兩個(gè)矩陣乘法（Linear Layer）的情況，此時(shí)通常都會(huì)采用先 Column Parallelism，之后 Row Parallelism 的方式切分，可以在兩個(gè) Linear 之間減少一次通信操作。如下圖所示，W 是第一個(gè) Linear 權(quán)重，V 是第二個(gè) Linear 權(quán)重。只用在最后進(jìn)行一次 AllReduce 操作即可：

三、AlexNet

Tensor Parallelism 在 AI 模型中的應(yīng)用最早可以追溯到著名的論文（2012: ImageNet Classification with Deep Convolutional Neural Networks），也就是 AlexNet，作者是 Alex Krizhevsky，Ilya Sutskever 和 Geoffrey E. Hinton。如下圖所示為 AlexNet 模型的網(wǎng)絡(luò)結(jié)構(gòu)：其整個(gè)模型由 8 個(gè)可學(xué)習(xí)層組成，包括 5 個(gè) Convolutional 層 和 3 個(gè) Fully Connected 層；此外還有幾個(gè) Max Pooling 和 ReLU 層。

然而，當(dāng)時(shí)流行的消費(fèi)級(jí) GPU 為 GTX 580，只有 3GB 顯存，無法訓(xùn)練這么大的模型，因此作者采用 Tensor Parallelism 的方式，當(dāng)然，那個(gè)時(shí)候還沒叫 Tensor Parallelism。其切分也很簡(jiǎn)單，如下圖所示，Conv 層按照卷積核（Kernel）切分，而 FC 層按照神經(jīng)元（Neuron）切分。由于 Conv 層也可以等價(jià)于矩陣乘操作，所以 Conv 和 FC 的切分也都可以理解為矩陣乘中將 Weight 按行切分。

2012 年賈揚(yáng)清的 Caffe 還沒有誕生，Alex Krizhevsky 自己實(shí)現(xiàn)了一個(gè) cuda-convnet 庫(kù)（Google Code：cuda-convnet），為了支持 AlexNet 的 Tensor Parallelism，也繼續(xù)開發(fā)了 cuda-convnet2（Google Code：cuda-convnet2）。

PS：冷知識(shí)，NVIDIA 的 K20 GPU 是 2012 年發(fā)布的，只有 5GB 顯存；P100 是 2016 年發(fā)布的，最大也只有 16GB 顯存，也是在 P100 GPU 中首次引入 NVLink；NVIDIA 的集合通信庫(kù) NCCL 也是在 2016 年發(fā)布。作者本人也是在這一年用 CUDA 重構(gòu)了熊廠的圖像檢索系統(tǒng)，用的還是 K1200 GPU。

四、Google DistBelief

和 AlexNet 同年，Google 團(tuán)隊(duì)也發(fā)表了 DistBelief 論文（ 2012：Large Scale Distributed Deep Networks）。論文主要討論了如何使用大規(guī)模分布式計(jì)算集群（CPU）來訓(xùn)練具有數(shù)十億參數(shù)的深度神經(jīng)網(wǎng)絡(luò)。論文的核心貢獻(xiàn)是開發(fā)了 DistBelief 軟件框架，該框架能夠利用數(shù)千臺(tái)機(jī)器來訓(xùn)練大型模型（1.7B，現(xiàn)在的 1.7B 已經(jīng)是 Tiny Model 了 (⊙o⊙)…）。

如下圖所示為 DistBelief 中的模型并行（Model Parallelism）方案，這里已經(jīng)涉及 Tensor Parallelism 和 Pipeline Parallelism。下圖中是一個(gè) 5 層的深度神經(jīng)網(wǎng)絡(luò)，被切分到了 4 個(gè)機(jī)器上。由于每一層不是 Fully Connected 的，而是部分連接，因此只有切分的邊緣需要跨設(shè)備通信：

此外，模型訓(xùn)練也不是采用的集合通信方式，而是使用了 Parameter Server 架構(gòu)，如下圖所示。也就是說模型權(quán)重的更新都是在專有的 Parameter Server 上進(jìn)行，更新完后會(huì) Broadcast 給相應(yīng)的 Training Worker：

如下圖所示，作者在 2012 年已經(jīng)訓(xùn)練了 1.7B 參數(shù)量的模型。（PS：可能是模型結(jié)構(gòu)或者訓(xùn)練資源的制約，最終并沒有在 ImageNet 上大規(guī)模訓(xùn)練，和 AlexNet 同期但并沒有在效果上超越 2個(gè) GPU 上訓(xùn)練的 AlexNet）

五、Facebook TP + DP

Facebook AI 團(tuán)隊(duì)在 [1312.5853] Multi-GPU Training of ConvNets 中將數(shù)據(jù)并行（Data Parallelism）和 AlexNet 的 Tensor Parallelism 相結(jié)合。如下圖 Figure 5 所示，GPU 1 和 GPU 2 共同包含了一份完整的模型副本，并通過 Tensor Parallelism 的方式切分模型；GPU 3 和 GPU 4 包含另外一份模型副本，兩個(gè)模型部分使用不同的數(shù)據(jù)訓(xùn)練，并在每個(gè) Iter 進(jìn)行聚合同步： 

作者也對(duì)比了不同的分布式策略，最終使用 Tensor Parallelism + Data Parallelism 的方式獲得最優(yōu)的加速比：

六、Google TP + DP

Google 在 [1404.5997] One weird trick for parallelizing convolutional neural networks 中進(jìn)一步對(duì) TP + DP 進(jìn)行了優(yōu)化（PS：其實(shí)作者還是 Alex Krizhevsky，從 Google 離職后似乎不再進(jìn)行 AI 研究相關(guān)工作；現(xiàn)在 AlexNet 的二作 Ilya Sutskever 更為人熟知，尤其是 OpenAI 火了之后）。如下圖 Figure 1 所示，具體來說，作者發(fā)現(xiàn)其模型的大部分參數(shù)集中在最后的 Fully Connected 層，因此在前面的 Conv 層采用 DP，而在后面的 FC 層采用 TP。其中 DP 里每個(gè) Work 使用不同的 Batch 數(shù)據(jù)，而 TP 中使用相同的 Batch 數(shù)據(jù)：

如下圖 Table 1 所示，在 2GPU 和 4GPU 幾乎都獲得了線性加速：

七、UC Berkeley TP + DP

之前方案中討論了 Model Parallelism 和 Tensor Parallelism，其主要都是指權(quán)重參數(shù)按 Column 的切分方式。而在 [1712.04432] Integrated Model, Batch and Domain Parallelism in Training Neural Networks 中（PS：本文中作者用的是 Y=WX 的方式，所以雖然 Weight 按照 Row 切分了，但是實(shí)際對(duì)應(yīng) Y=XA 中的 Column Parallelism），UC Berkeley 等作者進(jìn)一步討論了輸入（Activation，X）相關(guān)的切分方式（作者稱作 Domain Parallelism），且與 Data Parallelism 和 Column Parallelism 一起分析（比如分析通信量）、整合，并證明 Data Parallelism 或 Tensor Parallelism 單獨(dú)使用都不是最優(yōu)的方案。

當(dāng)然，本文中作者的討論也有一定的局限性，比如主要針對(duì) AlexNet，并且假定所有設(shè)備之間都使用相同的拓?fù)溥B接。

如下圖 Figure 5 所示，作者首先評(píng)估了 Data Parallelism + Column Parallelism 的方案。其相當(dāng)于 TP=2，DP=3：

• 第一行：Forward 過程，采用的 Column Parallelism，所以計(jì)算完之后需要 AllGather 操作拿到完整數(shù)據(jù)。
• 第二行：Backward 權(quán)重梯度過程，相當(dāng)于 Row Parallelism，所以計(jì)算完之后需要 AllReduce 操作。
• 第三行：Backward 鏈?zhǔn)椒▌t（輸入梯度）過程，相當(dāng)于 Row Parallelism，所以計(jì)算完之后需要 AllReduce 操作。?

如下圖 Figure 3 就是其所述的 Domain Parallelism，因?yàn)槠鋽?shù)據(jù)實(shí)際是按 NCHW 方式存儲(chǔ)的，因此這里按照 Width 切分為 4 部分。這里實(shí)際就是對(duì)輸入的切分：

作者對(duì) Domain Parallelism 的通信也進(jìn)行了分析，對(duì)于 AlexNet，前積層都是 Conv 層，Activation 比較大，適合 Domain Parallelism（切分輸入）；后幾層是 FC 層，參數(shù)量比較大，適合切分 Weight。

八、Google Mesh-TensorFlow

Google 團(tuán)隊(duì)在 2017 年發(fā)表著名的 [1706.03762] Attention Is All You Need，此后模型規(guī)模不斷擴(kuò)大，Model Parallelism 已經(jīng)成為不可或缺的分布式訓(xùn)練方案。然而，當(dāng)時(shí)并沒有一個(gè)高效的 Model Parallelism 的訓(xùn)練框架，因此 Google 的作者提出了 Mesh-TensorFlow（[1811.02084] Mesh-TensorFlow: Deep Learning for Supercomputers）框架，它是一個(gè)通用的分布式 Tensor 計(jì)算框架，用戶可以在多維網(wǎng)格的任何維度上分割任何的張量，也就天然可以支持 Data Parallelism 以及 Tensor Parallelism（包括 Column Parallelism 和 Row Parallelism）。

此外，作者也在 512 個(gè) TPU 上訓(xùn)練了 5B 參數(shù)量的 Transformer 模型（PS：本文的一作也是 Attention Is All You Need 的二作 Noam Shazeer）。

也是在這篇文章中作者對(duì) Transformer 模型里這種 2 個(gè) FC 層相鄰的 Model Parallelism 切分方式進(jìn)行了介紹。如下圖所示，其中 w 和 v 是兩個(gè)權(quán)重矩陣，x 為輸入，y 為輸出，中間激活為 Relu：

針對(duì)上述 2 個(gè) FC 相鄰（比如 FFN，Attention）的計(jì)算，作者提出了可以第一個(gè) FC 的 Weight（w）采用列切，第二個(gè) FC 的 Weight（v）采用行切的方案。如下圖所示，這樣的好處是在第一個(gè)矩陣乘之后并不用通信，只用在最后的 y 這里 AllReduce 即可：

其實(shí)作者也探討了更復(fù)雜的 2D Parallelism 切分方式，如下圖 Figure 5 所示。不過最終在 GitHub - tensorflow/tensor2tensor 中實(shí)現(xiàn)的還是上面提到的方式：

九、NVIDIA Megatron-LM

Google 在 Mesh-TensorFlow 中已經(jīng)提供了比較豐富的 Tensor Parallelism 的支持，然而它是一種新的框架，用戶如果想要使用會(huì)有比較大的改造成本。因此 NVIDIA 在 [1909.08053] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism 中發(fā)布了 Megatron-LM 框架，只用對(duì)已有的 Pytorch Transformer 模型進(jìn)行少量修改就可以實(shí)現(xiàn)，相應(yīng)的更加簡(jiǎn)單、方便。

如下圖 （a）所示，其和 Mesh-Tensorflow 類似，MLP 層的兩個(gè) FC 采用先列切（A，Column Parallelism），然后行切（B，Row Parallelism）的方案，這樣兩個(gè) FC 之間不用通信：

如下圖（b）所示，由于每個(gè) Head 的 Attention，softmax 都是獨(dú)立的，因此可以采用按照 Head 的方式切分（等價(jià)于 Column Parallelism），然后對(duì)之后的 FC 采用行切分（B，Row Parallelism），這樣 Self-Attention 中間也不用通信：

需要說明的是，作者在其中引入了 f 和 g 模塊，其主要就是一個(gè)通信的抽象，比如如下圖所示，f 的 Forward 就是一個(gè) Identify 操作，Backward 是一個(gè) AllReduce 操作；而 g 正好相反，F(xiàn)orward 是一個(gè) AllReduce 操作，Backward 是一個(gè) Identify 操作。

如下圖 Figure 4 所示，采用這種 Tensor Parallelism 的 Transformer 模型的每一層在 Forward 和 Backward 都各只有 2 次 AllReduce 操作：

十、Colossal-AI 2D Tensor Parallelism

在 [2104.05343] An Efficient 2D Method for Training Super-Large Deep Learning Models 中，作者進(jìn)一步提出了 2D Tensor Parallelism 的方案 Optimus，它是一種高效的 2D Tensor Parallelism 并行范式，可以進(jìn)一步促進(jìn)無限大 LLM 的訓(xùn)練，Colossal-AI 作者也在 GitHub - hpcaitech/ColossalAI 中開源了相應(yīng)代碼。

作者提出 2D Tensor Parallelism 是因?yàn)橛^察到之前的方案（比如 Megatron-LM，1D）中雖然將模型的參數(shù)分配到了多個(gè)設(shè)備上，但是每個(gè)設(shè)備上 Forward 和 Backward 的 Activation 并沒有被有效切分，還會(huì)占用相應(yīng)的顯存空間。此外，雖然 Mesh-Tensorflow 中進(jìn)行了多維 Tensor Parallelism 的抽象，但是并沒有很好的實(shí)現(xiàn)。因此，作者提出了 Optimus 方案，以最大化通信效率，最小化顯存占用。

PS：這里有個(gè)約束，2D 的兩個(gè)維度需要是相等的，也就是說 2D 的數(shù)量只能是 2x2=4， 4x4=16， 6x6=36， 8x8=64 這些。

如下圖所示為 Megatron-LM 對(duì) Transformer 模型的切分方案，其中綠色為激活，藍(lán)色為權(quán)重?？梢钥闯觯瑱?quán)重都在各個(gè)設(shè)備上進(jìn)行了 1 維切分，而 Input 和 Output 在每個(gè)設(shè)備上都是保存的整個(gè)副本，也就是存在存儲(chǔ)上的冗余。作者將這種方式稱為 1D Tensor Parallelism。

對(duì)于 1D Tensor Parallelism，給定 P 個(gè)設(shè)備，其計(jì)算、內(nèi)存和通信成本如下所示：

如下圖所示為本文的 2D Tensor Parallelism 切分方案，可以看出，不管是權(quán)重還是 Input、Output 都按照設(shè)備進(jìn)行了切分，此時(shí)設(shè)備也是按照 2D Mesh （q，q）進(jìn)行排布。

對(duì)于 2D Tensor Parallelism，給定 P = q x q 個(gè)設(shè)備，其計(jì)算、內(nèi)存和通信成本如下所示：

如下圖 Table 1 所示，作者也統(tǒng)計(jì)了 Megatron-LM 和 Optimus 的通信量和計(jì)算量，可以看出，兩者計(jì)算量相同（分塊計(jì)算并不會(huì)增加計(jì)算量），而通信量有所差距，當(dāng) p 比較大時(shí) Optimus 的 2D 方式可以有一定優(yōu)勢(shì)。不過當(dāng)前通常都會(huì)將一個(gè) TP 分組放在一臺(tái)機(jī)器內(nèi)（通常不超過 8 個(gè) GPU），此時(shí) Optimus 在通信上并沒有什么優(yōu)勢(shì)，不過在顯存上會(huì)更有優(yōu)勢(shì)，當(dāng)然也會(huì)增加實(shí)現(xiàn)和調(diào)度的復(fù)雜度。

十一、Colossal-AI 3D Tensor Parallelism

在 [2105.14450] Maximizing Parallelism in Distributed Training for Huge Neural Networks 中，作者在 2D Tensor Parallelism 的基礎(chǔ)上進(jìn)一步提出了 3D Parallelism。通過實(shí)現(xiàn)完美的負(fù)載均衡，提出的方法提供比當(dāng)時(shí)最先進(jìn)的 1D 和 2D Tensor Parallelism 更小的內(nèi)存和通信成本。在 64 個(gè) V100 GPU 上的實(shí)驗(yàn)結(jié)果表明，提出的 3D 并行比 1D 和 2D 并行分別加速 2.32x 和 1.57x。

PS：這里同樣有個(gè)約束，3D 的三個(gè)維度需要是相等的，也就是說 3D 的數(shù)量只能是 2x2x2=8， 4x4x4=64 這些，無法支持 4 個(gè)設(shè)備的 Tensor Parallelism。

如下圖 Figure 2 所示為一個(gè) 3D 矩陣乘的示例，其中 A 和 B 的箭頭表示  Broadcast 的方向，C 中的箭頭是 Reduce 的方向：

如下圖 Figure 6 所示是 Transformer Layer 的 3D Tensor Parallelism 方案：

對(duì)于 3D Tensor Parallelism，給定 P = q x q x q 個(gè)設(shè)備，其計(jì)算、內(nèi)存和通信成本如下所示：

如下圖 Table 2 所示，作者對(duì) 1D，2D 和本文提出的 3D Tensor Parallelism 的方案進(jìn)行了驗(yàn)證，其中的 Average step time(s) 為 (Forward time + Backward time) / Batch Size。從中可以看出：

• 2D Tensor Parallelism相比 1D 只在 64 GPU 時(shí)略微有優(yōu)勢(shì)；而在更少的 GPU 時(shí)，反而不如 1D Tensor Parallelism。
• 3D Tensor Parallelism相比 1D 有比較大的提升。（PS：3D 更節(jié)約顯存，可以使用更大的 Batch Size，如果是 Zero-DP 或者 FSDP 的方式，通過增加 Batch Size 是否可以達(dá)到同樣的效果？）

PS：當(dāng)前除了 Zero-DP 和 FSDP 外，也有 Pipeline Parallelism，Sequence Parallelism，使用 2D 或者 3D Tensor Parallelism 的場(chǎng)景比較少；此外，GPU 算力越來越強(qiáng)，3D 方式會(huì)將矩陣計(jì)算切分得更加小，不一定能充分發(fā)揮 GPU 算力，還是需要綜合考慮。

十二、Megatron-LM TP+DP

NVIDIA 在上述的 Megatron-LM 中也介紹了 Tensor Parallelism 和 DP Parallelism 混合的方案。當(dāng)前大規(guī)模 GPU 集群廣泛采用 NVIDIA GPU 構(gòu)建，比如 H100、A100、V100，以及后續(xù)的 B200。通常每臺(tái)機(jī)器有 8 個(gè) GPU，每臺(tái)機(jī)器內(nèi)由 NVLink 高速互聯(lián)，通信帶寬更高。因此，在排布分布式并行策略時(shí)，通常將通信量比較大且有些時(shí)候無法有效 Overlap 的 Tensor Parallelism 放在一臺(tái)機(jī)器內(nèi)部，而不同機(jī)器之間采用 Data Parallelism 或 Pipeline Parallelism。

如下圖 Figure 8 所示為 Megatron-LM 中一個(gè) DP + TP 的混合分布式并行方案，總共采用 64 臺(tái) 8 GPU 機(jī)器，共 512 GPU。

• 每臺(tái)機(jī)器的 8 個(gè) GPU 組成一個(gè) Model Parallelism Group（TP），共 64 個(gè) TP Group；每個(gè)TP Group 內(nèi)的 GPU 包含不同的模型參數(shù)，并且使用相同的訓(xùn)練數(shù)據(jù)。
• 所有設(shè)備的同號(hào) GPU（比如 GPU 1，9，...，505）組成一個(gè) Data Parallelism Group（DP），共 8 個(gè) DP Group；每個(gè) DP Group 內(nèi)的 GPU 都有相同的模型參數(shù)，但是使用不同的訓(xùn)練數(shù)據(jù)。

基于以上的分布式排布方案，反向傳播和梯度更新的過程如下所示（當(dāng)然，下述兩個(gè)階段可以一定程度上 Overlap，以提升訓(xùn)練速度）：

1. 各個(gè) TP Group 獨(dú)立的進(jìn)行 Backward，Backward 涉及的 AllReduce 也只在 Group 內(nèi)部，也就是單個(gè)機(jī)器內(nèi)。每個(gè) GPU 上也只有對(duì)應(yīng)權(quán)重參數(shù)的梯度。
2. 各個(gè) DP Group 獨(dú)立地進(jìn)行梯度 AllReduce 以及權(quán)重參數(shù)更新。

如下圖所示，我們?cè)谥暗奈恼轮薪榻B過，當(dāng)前的大規(guī)模 GPU 訓(xùn)練集群的網(wǎng)絡(luò)拓?fù)渲校粋€(gè) Group 里不同機(jī)器的同號(hào) GPU 對(duì)應(yīng)的 NIC 會(huì)連接到同一個(gè) Leaf 交換機(jī)，比如所有機(jī)器的 2 號(hào) NIC 都連接到了 2 號(hào) Leaf 交換機(jī)。這也與上述的分布式排布方案相對(duì)應(yīng)，所有 TP Group 的通信都在機(jī)器內(nèi)，所有 DP Group 的通信都只用經(jīng)過 Leaf 交換機(jī)，此外在交換機(jī)上執(zhí)行規(guī)約計(jì)算也變成了一種可能。

十三、分布式 LLM 推理：DP+TP

在 LLM 的推理場(chǎng)景也會(huì)存在 DP + TP 的組合，比如為了支持更大的吞吐、考慮容災(zāi)等場(chǎng)景，通常會(huì)有多個(gè) Instance；每個(gè) Instance 都會(huì)有完整的模型參數(shù)，這些 Instance 可以獨(dú)立、并行的處理各自的 Request，這種方式可以理解為 DP，只是因?yàn)橥评碇挥?Forward，所以不涉及權(quán)重更新，各個(gè) Instance 之間也就不必通信。而單個(gè) GPU 的顯存或者算力可能不足，每個(gè) Instance 可能使用多個(gè) GPU 來推理，在 LLM 場(chǎng)景最常見的就是采用 TP 的方案。

然而，此時(shí)使用 TP 也可能有其局限性。如下圖所示（圖片來自 [2312.03134] A Hardware Evaluation Framework for Large Language Model Inference），我們之前已經(jīng)分析過，一個(gè) Transformer Layer 中會(huì)有兩次 AllReduce 通信，一次是 MHA 中的最后一個(gè) Linear，一次是 FFN 中的最后一個(gè) Linear。以 GPT-3 175B 模型為例，其共 96 個(gè) Layer，也就是說一次 Forward 要有 192 次 AllReduce（忽略非 Transformer Layer 之外的通信）。每次的通信量與 Hidden Dim 和 Batch Size 成正比，模型確定后 Hidden Dim 確定，其通信量就與 Batch Size 成正比。 

由于 LLM 推理通常會(huì)使用 Continuous Batching 的方式提升吞吐，隨著 Batch Size 增加，MHA 和 FFN 的 Kernel 計(jì)算時(shí)延不會(huì)明顯增加（序列不太長(zhǎng)）；而 AllReduce 的通信量卻線性增加，相應(yīng)的通信時(shí)延增加的更加明顯，以至于 AllReduce 通信可能成為瓶頸。GPU 間的通信時(shí)延與 GPU 之間的互聯(lián)方式有關(guān)，節(jié)點(diǎn)內(nèi)通常通過 PCIe 或 NVLink 互聯(lián)，在 PCIe 互聯(lián)方式下就需要密切關(guān)注 AllReduce 相應(yīng)的性能問題，尤其還會(huì)涉及跨 PCIe Switch 或 跨 NUMA 通過 UPI 通信的場(chǎng)景。

如下圖所示為 Batch size 等于 1 和 512 時(shí) LLM 中幾個(gè)主要 OP 的計(jì)算耗時(shí)，可以看出，將 Batch size 從 1 增加到 512，計(jì)算量增加 512 倍，但是其整體時(shí)間只增加到原來的 3 倍左右（圖片來自 openppl-public · GitHub）：

除此之外，Continuous Batching 的方式會(huì)希望組合盡可能大的 Batch Size，也就意味著同一時(shí)間可能只有一個(gè) Batch 的數(shù)據(jù)在計(jì)算，AllReduce 通信與計(jì)算無法充分 Overlap，出現(xiàn)算力的浪費(fèi)。針對(duì)這個(gè)問題，也有一些優(yōu)化方案，比如：

• Prefill 階段和 Decoding 階段的細(xì)粒度 Overlap，比如新到的 Request 在執(zhí)行 Prefill 計(jì)算時(shí)，可以執(zhí)行之前已經(jīng) Batching 的 Decoding 的通信；而 Prefill 在通信時(shí)，可以執(zhí)行 Decoding 的計(jì)算。當(dāng)然，實(shí)現(xiàn)這種細(xì)粒度的計(jì)算、通信 Overlap 的代價(jià)也比較高。
• Pipeline Parallelism、Sequence Parallelism 以及其組合方案，這里先不展開。

十四、參考鏈接

1. https://lightning.ai/docs/pytorch/stable/advanced/model_parallel/tp.html
2. https://proceedings.neurips.cc/paper_files/paper/2012/file/c399862d3b9d6b76c8436e924a68c45b-Paper.pdf
3. https://code.google.com/archive/p/cuda-convnet/
4. https://code.google.com/archive/p/cuda-convnet2/
5. https://www.cs.toronto.edu/~ranzato/publications/DistBeliefNIPS2012_withAppendix.pdf
6. https://arxiv.org/abs/1312.5853
7. https://arxiv.org/abs/1404.5997
8. https://arxiv.org/abs/1712.04432
9. https://arxiv.org/abs/1706.03762
10. https://arxiv.org/abs/1811.02084
11. https://github.com/tensorflow/tensor2tensor
12. https://arxiv.org/abs/1909.08053
13. https://arxiv.org/abs/2104.05343
14. https://github.com/hpcaitech/ColossalAI
15. https://arxiv.org/abs/2105.14450
16. https://arxiv.org/abs/2312.03134

本文轉(zhuǎn)載自 ??AI閑談??，作者： AI閑談