一、背景

本文中我們繼續介紹分布式訓練中的張量并行（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 中的模型切分在實際使用的時候還會將相應的參數 AllGather 到當前設備，使用全部的參數進行計算。而 Tensor Parallelism 中的參數都始終在當前設備，最終聚合的是結果（Activation）。當然，在后續的文章中我們也會介紹 Zero 和 FSDP 相關方案。

二、分布式矩陣乘

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

2.1 Column Parallelism

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

2.2 Row Parallelism

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

2.3 Column Parallelism + Row Parallelism

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

三、AlexNet

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

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

2012 年賈揚清的 Caffe 還沒有誕生，Alex Krizhevsky 自己實現了一個 cuda-convnet 庫（Google Code：cuda-convnet），為了支持 AlexNet 的 Tensor Parallelism，也繼續開發了 cuda-convnet2（Google Code：cuda-convnet2）。

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

四、Google DistBelief

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

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

此外，模型訓練也不是采用的集合通信方式，而是使用了 Parameter Server 架構，如下圖所示。也就是說模型權重的更新都是在專有的 Parameter Server 上進行，更新完后會 Broadcast 給相應的 Training Worker：

如下圖所示，作者在 2012 年已經訓練了 1.7B 參數量的模型。（PS：可能是模型結構或者訓練資源的制約，最終并沒有在 ImageNet 上大規模訓練，和 AlexNet 同期但并沒有在效果上超越 2個 GPU 上訓練的 AlexNet）

五、Facebook TP + DP

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

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

六、Google TP + DP

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

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

七、UC Berkeley TP + DP

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

當然，本文中作者的討論也有一定的局限性，比如主要針對 AlexNet，并且假定所有設備之間都使用相同的拓撲連接。

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

• 第一行：Forward 過程，采用的 Column Parallelism，所以計算完之后需要 AllGather 操作拿到完整數據。
• 第二行：Backward 權重梯度過程，相當于 Row Parallelism，所以計算完之后需要 AllReduce 操作。
• 第三行：Backward 鏈式法則（輸入梯度）過程，相當于 Row Parallelism，所以計算完之后需要 AllReduce 操作。?

如下圖 Figure 3 就是其所述的 Domain Parallelism，因為其數據實際是按 NCHW 方式存儲的，因此這里按照 Width 切分為 4 部分。這里實際就是對輸入的切分：

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

八、Google Mesh-TensorFlow

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

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

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

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

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

九、NVIDIA Megatron-LM

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

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

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

需要說明的是，作者在其中引入了 f 和 g 模塊，其主要就是一個通信的抽象，比如如下圖所示，f 的 Forward 就是一個 Identify 操作，Backward 是一個 AllReduce 操作；而 g 正好相反，Forward 是一個 AllReduce 操作，Backward 是一個 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 中，作者進一步提出了 2D Tensor Parallelism 的方案 Optimus，它是一種高效的 2D Tensor Parallelism 并行范式，可以進一步促進無限大 LLM 的訓練，Colossal-AI 作者也在 GitHub - hpcaitech/ColossalAI 中開源了相應代碼。

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

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

如下圖所示為 Megatron-LM 對 Transformer 模型的切分方案，其中綠色為激活，藍色為權重。可以看出，權重都在各個設備上進行了 1 維切分，而 Input 和 Output 在每個設備上都是保存的整個副本，也就是存在存儲上的冗余。作者將這種方式稱為 1D Tensor Parallelism。

對于 1D Tensor Parallelism，給定 P 個設備，其計算、內存和通信成本如下所示：

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

對于 2D Tensor Parallelism，給定 P = q x q 個設備，其計算、內存和通信成本如下所示：

如下圖 Table 1 所示，作者也統計了 Megatron-LM 和 Optimus 的通信量和計算量，可以看出，兩者計算量相同（分塊計算并不會增加計算量），而通信量有所差距，當 p 比較大時 Optimus 的 2D 方式可以有一定優勢。不過當前通常都會將一個 TP 分組放在一臺機器內（通常不超過 8 個 GPU），此時 Optimus 在通信上并沒有什么優勢，不過在顯存上會更有優勢，當然也會增加實現和調度的復雜度。

十一、Colossal-AI 3D Tensor Parallelism

在 [2105.14450] Maximizing Parallelism in Distributed Training for Huge Neural Networks 中，作者在 2D Tensor Parallelism 的基礎上進一步提出了 3D Parallelism。通過實現完美的負載均衡，提出的方法提供比當時最先進的 1D 和 2D Tensor Parallelism 更小的內存和通信成本。在 64 個 V100 GPU 上的實驗結果表明，提出的 3D 并行比 1D 和 2D 并行分別加速 2.32x 和 1.57x。

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

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

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

對于 3D Tensor Parallelism，給定 P = q x q x q 個設備，其計算、內存和通信成本如下所示：

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

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

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

十二、Megatron-LM TP+DP

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

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

• 每臺機器的 8 個 GPU 組成一個 Model Parallelism Group（TP），共 64 個 TP Group；每個TP Group 內的 GPU 包含不同的模型參數，并且使用相同的訓練數據。
• 所有設備的同號 GPU（比如 GPU 1，9，...，505）組成一個 Data Parallelism Group（DP），共 8 個 DP Group；每個 DP Group 內的 GPU 都有相同的模型參數，但是使用不同的訓練數據。

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

1. 各個 TP Group 獨立的進行 Backward，Backward 涉及的 AllReduce 也只在 Group 內部，也就是單個機器內。每個 GPU 上也只有對應權重參數的梯度。
2. 各個 DP Group 獨立地進行梯度 AllReduce 以及權重參數更新。

如下圖所示，我們在之前的文章中介紹過，當前的大規模 GPU 訓練集群的網絡拓撲中，同一個 Group 里不同機器的同號 GPU 對應的 NIC 會連接到同一個 Leaf 交換機，比如所有機器的 2 號 NIC 都連接到了 2 號 Leaf 交換機。這也與上述的分布式排布方案相對應，所有 TP Group 的通信都在機器內，所有 DP Group 的通信都只用經過 Leaf 交換機，此外在交換機上執行規約計算也變成了一種可能。

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

在 LLM 的推理場景也會存在 DP + TP 的組合，比如為了支持更大的吞吐、考慮容災等場景，通常會有多個 Instance；每個 Instance 都會有完整的模型參數，這些 Instance 可以獨立、并行的處理各自的 Request，這種方式可以理解為 DP，只是因為推理只有 Forward，所以不涉及權重更新，各個 Instance 之間也就不必通信。而單個 GPU 的顯存或者算力可能不足，每個 Instance 可能使用多個 GPU 來推理，在 LLM 場景最常見的就是采用 TP 的方案。

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

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

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

除此之外，Continuous Batching 的方式會希望組合盡可能大的 Batch Size，也就意味著同一時間可能只有一個 Batch 的數據在計算，AllReduce 通信與計算無法充分 Overlap，出現算力的浪費。針對這個問題，也有一些優化方案，比如：

• Prefill 階段和 Decoding 階段的細粒度 Overlap，比如新到的 Request 在執行 Prefill 計算時，可以執行之前已經 Batching 的 Decoding 的通信；而 Prefill 在通信時，可以執行 Decoding 的計算。當然，實現這種細粒度的計算、通信 Overlap 的代價也比較高。
• 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

本文轉載自 ??AI閑談??，作者： AI閑談