張量並行（Tensor Parallelism）是什麼？

將單個張量（矩陣或更高維數組）的計算分割到多個 GPU 設備上，通過跨設備並行計算矩陣乘法等操作的分佈式訓練方法。

核心概念

張量並行（Tensor Parallelism）的核心思想是將單個計算操作的張量分割到多個 GPU 上。例如，考慮一個形狀為 (N, K) 的權重矩陣和形狀為 (B, N) 的輸入矩陣的乘法，其中 B 是批次大小，N 和 K 是維度。在張量並行中：

1. 列分割方案：權重矩陣按列分割到多個 GPU。每個 GPU 保有該權重矩陣的一部分列，計算該部分與輸入的乘積，得到輸出的一部分。
2. 行分割方案：權重矩陣按行分割到多個 GPU。每個 GPU 計算與輸入的乘積，得到輸出的一個完整部分。
3. 2D 分割方案：對於非常大的矩陣，可以同時按行和列分割到多個 GPU。

Flax 和 PyTorch 等框架提供了編程介面來方便地指定張量的分割方案（稱為分片規格或 mesh）。例如，可以指定某個矩陣的行在 GPU 維度上分割，列在另一個維度上分割。

張量並行的優勢是：

1. 通訊開銷相對較低。相比於流水線並行，張量並行通常只需要在計算完成後進行一次規約（reduce）或廣播（broadcast）操作。
2. 負載均衡較好。如果分割得當，所有 GPU 的計算量相近，不會出現某些 GPU 閒置的情況。
3. 實現相對簡單。許多框架已經提供了自動張量並行的支援。

張量並行的劣勢包括：

1. 要求設備間互連帶寬較高。分割的方式直接影響通訊量。
2. 不適合所有層類型。例如，對於嵌入層（embedding），張量並行的效果較差。
3. 可能引入複雜的通訊依賴。

運作原理

以 Transformer 的前饋層（FFN）為例，說明張量並行的運作。

FFN 層包含兩個線性變換： output = Linear2(ReLU(Linear1(input)))

其中 Linear1 的權重是 (hidden_size, intermediate_size)，Linear2 的權重是 (intermediate_size, hidden_size)。

在張量並行中，假設有 2 個 GPU：

1. Linear1 計算：

   • 權重按列分割：GPU0 保有第一部分列，GPU1 保有第二部分列。
   • 輸入複製到兩個 GPU。
   • 每個 GPU 計算其部分的矩陣乘積，得到輸出的一部分。
   • 輸出是分割的：GPU0 有第一部分，GPU1 有第二部分。

2. ReLU 激活：

   • 在分割的張量上逐元素應用，無需通訊。

3. Linear2 計算：

   • 權重按行分割：GPU0 保有第一部分行，GPU1 保有第二部分行。
   • 由於輸入已經分割，每個 GPU 可以計算其行對應的輸出。
   • 需要一次全規約（all-reduce）以綜合輸出。

整個過程中，通訊主要發生在矩陣乘法之間的同步點和規約操作。

實際應用

張量並行在訓練非常大的模型時廣泛應用。Google 的 PaLM、Meta 的 Llama 等大型模型訓練都採用了張量並行。

在分佈式訓練中，張量並行經常與流水線並行和數據並行結合使用，形成混合並行策略。例如，可以使用數據並行分割 batch，流水線並行分割模型層，張量並行分割每一層的計算。

在 GPU 集群訓練中，張量並行幫助利用高互連帶寬的 GPU 集群（如使用 NVLink 互連的 GPU 設備）的高通訊速度。

在推理服務中，張量並行也用於加速推理。系統將模型分割到多個 GPU，每個 GPU 計算其負責的部分，通過並行化提高吞吐量。

在微調預訓練模型時，張量並行允許研究人員在有限的 GPU 資源上微調非常大的模型。

常見誤區

一個常見的誤區是認為張量並行總是能帶來線性加速。實際上，加速因素取決於通訊開銷與計算開銷的比例。如果通訊開銷過高，加速效果不理想。

另一個誤區是認為張量並行適用於所有模型架構。實際上，張量並行對參數集中在少數幾層（如 Transformer 的注意力層和前饋層）的模型效果最好。對於其他架構，效果可能不同。

還有人誤以為張量並行只需要設置一個參數就能自動工作。實際上，選擇合適的分割方案需要考慮模型結構、GPU 設備拓撲和通訊特性。不同的分割方案會導致不同的性能。

與相關技術的比較

張量並行與流水線並行都是模型並行技術，但分割粒度不同。流水線並行在層级分割模型，張量並行在層內分割計算。

與數據並行相比，張量並行用於分割模型本身，而數據並行分割輸入數據。通常它們結合使用：數據並行分割 batch，張量並行分割層。

與序列並行相比，序列並行沿著序列長度維度分割（通常用於長序列），而張量並行沿著隱藏維度分割。