Tensor-Parallel-Training

論文參考

張量並行

Megatron-LM 這篇論文主要貢獻是模型並行化（Model Parallelism）技術，用於在 多個 GPU 之間拆分 Transformer 層，主要用於訓練超大語言模型

張量並行主要透過切分 FFN(MLP)層 和 多頭自注意力層來實現張量並行

FFN(MLP)層

• $Y=GeLU(XA)$
• $Z=Dropout(YB)$
• $X$ 是輸入矩陣（batch size × hidden dimension）
• $A$ 是權重矩陣（hidden dimension × MLP hidden size）
• $B$ 是 第二個權重矩陣（MLP hidden size × hidden dimension）

前向傳播

方法 : 將

$A$ 沿著列方向切割

1. 將
   $A$ 沿著列方向劃分為兩個部分：
   
   $$A=[A_1,A_2]$$
   • 透過這樣切分，計算變成
     
     $$[Y_1,Y_2]=[\text{GeLU}(X{A}_1),\text{GeLU}(X{A}_2)]$$
   • 可以在每個 GPU 上獨立應用 GeLU ，避免了同步點的需求，使得計算更為高效
2. 進行Dropout，我們經過剛剛那步已經得到
   $$[Y_1,Y_2]=[\text{GeLU}(X{A}_1),\text{GeLU}(X{A}_2)]$$，代表第一個 GEMM 和 GeLU 的計算已經完成，接下來進行 第二個 GEMM 和 Dropout
   
   $$Z=Dropout(YB)$$
   • $B$ 沿行方向分割:
     
     $$B=\left[\begin{array}{c}{B}_1\\ B_2\end{array}\right]$$
   • 計算變為:
     
     $$Z=\text{Dropout}(Y_1{B}_1+Y_2{B}_2)$$
   • 這種處理方式會需要同步點(透過 All Reduce)，因為Dropout 不能作用在部分張量，必須先 All Reduce 相加兩張卡的數值，才能進行 Dropout

反向傳播

1. 對
   $Z$ 進行微分
   
   $${\delta }_Z=\frac{\partial \mathcal{L}}{\partial Z}$$
   • 前向傳播時，已經進行了 All-Reduce，每張GPU 都有完整的數值，所以梯度也是完整的
   • 由於 Dropout 作用在整個
     $YB$ 之後，反向傳播時，應用 Dropout 反向操作
     
     $${\delta }_{YB}={\delta }_Z\odot M$$
   • $M$ 是 Dropout 的 mask
2. 對
   $B$ 和
   $Y$ 進行微分

• 對
  $B$ 求梯度，因為是沿行方向分割
  
  $$B=\left[\begin{array}{c}{B}_1\\ B_2\end{array}\right]$$
  • 在不同 GPU 上:
    
    $$\frac{\partial \mathcal{L}}{\partial B_1}=Y_1^T{\delta }_{YB},\frac{\partial \mathcal{L}}{\partial B_2}=Y_2^T{\delta }_{YB}$$
  • 這些梯度可以在 各 GPU 上獨立計算，不需要 All-Reduce
• 對
  $Y$ 求梯度
  
  $$Z=YB$$
  
  $$\frac{\partial \mathcal{L}}{\partial Y}={\delta }_{YB}{B}^T$$
  
  $${\delta }_Y={\delta }_{YB}{B}^T$$
  • 因為
    $B$ 是沿行方向分割
    
    $${\delta }_{Y_1}={\delta }_{YB}{B}_1^T,{\delta }_{Y_2}={\delta }_{YB}{B}_2^T$$
  • 因為
    $A$ 是沿著列方向分割的:
    • $Y_1$ 和
      $Y_2$沒有加總，所以反向傳播時梯度可以在各自的 GPU 上獨立計算，不需要 All-Reduce

3. 對
   $A$ 和
   $X$ 進行微分

• 對
  $A$ 求梯度
  • 已知
    
    $$Y_1=\text{GeLU}(X{A}_1),Y_2=\text{GeLU}(X{A}_2)$$
  • 鏈式規則
    
    $${\delta }_{X{A}_1}={\delta }_{Y_1}\cdot {\text{GeLU}}^{\prime }(X{A}_1),{\delta }_{X{A}_2}={\delta }_{Y_2}\cdot {\text{GeLU}}^{\prime }(X{A}_2)$$
  • 對
    $A_1,A_2$求梯度
    
    $$\frac{\partial \mathcal{L}}{\partial A_1}=X^T{\delta }_{X{A}_1},\frac{\partial \mathcal{L}}{\partial A_2}=X^T{\delta }_{X{A}_2}$$
  • 這些梯度只涉及當前 GPU 上的
    $A_1,A_2$，不用 All-Reduce
• 對
  $X$ 求梯度
  
  $${\delta }_X={\delta }_{X{A}_1}{A}_1^T+{\delta }_{X{A}_2}{A}_2^T$$
  • $X$ 是所有GPU 上是共享的
  • 因為
    $A$ 是沿列進行分割，
    ${\delta }_{X{A}_1}$ 和
    ${\delta }_{X{A}_2}$ 是各自計算的，所以需要 All-Reduce 確保所有 GPU 都獲得完整的梯度
    ${\delta }_X$
    
    $${\delta }_X=\text{AllReduce}({\delta }_{X{A}_1}{A}_1^T+{\delta }_{X{A}_2}{A}_2^T)$$

多頭自注意力層

多頭注意力的張量並行

主要解決自注意力層中，查詢(

$Q$)、鍵(

$K$)跟值(

$V$)的大規模矩陣乘法，Megatron-LM 通過列方向(Column Parallel)分割來優化:

• 將
  $Q,K,V$ 矩陣乘法沿列方向劃分
• 每張 GPU 只負責計算部分注意力頭
• 注意力頭與頭之間在部分計算是無關的，不需要立即進行 GPU 間的通訊
  • GPU 間不用交換
    $Q,K,V$ 的部分結果，減少通訊開銷

自注意力後的線性變換並行

算完自注意力後，最後有一個線性變換層，Megatron-LM 透過將輸出線性層的矩陣乘法沿行方向（Row Parallel）分割

• 輸出線性層的矩陣沿行方向劃分（Row Parallel），接收前一層平行化的注意力輸出

範例 : GPU 數 : 2

• 輸入向量
  $X$，形狀為
  $(B,S,D)$
• 自注意力頭數為 8
• 每張 GPU 負責 4 個注意力頭計算(Megatron-LM 通過列方向分割來優化)

張量並行

1. 列方向分割

• 將
  $Q,K,V$ 沿列方向劃分
  • GPU-0：計算前 4 個注意力頭的
    $Q,K,V$
  • GPU-1：計算後 4 個注意力頭的
    $Q,K,V$
    每張 GPU，只需計算:
    
    $$Q_{\text{local}}=X{W}_Q^{\text{local}},K_{\text{local}}=X{W}_K^{\text{local}},V_{\text{local}}=X{W}_V^{\text{local}}$$

2. 各自計算局部注意力(跟原本一致)
   GPU 間不需要通信
   
   $${\text{Attention}}_{\text{local}}=\text{softmax}\left(\frac{Q_{\text{local}}{K}_{\text{local}}^T}{\sqrt{d_k}}\right)V_{\text{local}}$$
3. 行方向分割輸出線性層（Row Parallelism）
   假設有輸出線性層的權重矩陣
   $W_O$
   1. 前面假設 2 張 GPU，把
      $W_O$ 沿 行方向（Row Parallel） 劃分成 2 部分：
      
      $$W_O=\left[\begin{array}{c}{W}_{O_1}\\ W_{O_2}\end{array}\right]$$
   2. 每張 GPU 計算自己的部分輸出
      • GPU-0 計算
        
        $$Z_1=\text{Attention}{W}_{O_1}$$
      • GPU-1 計算
        
        $$Z_2=\text{Attention}{W}_{O_2}$$
4. 合併結果
   每張 GPU 只計算了一部分
   $Z$，需要一個 All-Reduce 操作來合併
   $Z_1$ 和
   $Z_2$
   • 讓 所有 GPU 共享
     $Z_1$ 和
     $Z_2$
   • 每張 GPU 都可以獲得完整
     $Z$

Input embedding層

1. 並行化輸入嵌入
   將輸入的 Token（詞彙索引）轉換為向量表示
   • 透過嵌入矩陣（Embedding Matrix），嵌入矩陣的形狀為:
     
     $$W_E\in {\mathbb{R}}^{V\times H}$$
   • $O\in {\mathbb{R}}^{V\times BS}$ 是 one-hot 編碼的輸入 Token
   • 進行行方向分割
     
     $$W_E=\left[\begin{array}{c}{W}_{E1}\\ W_{E2}\end{array}\right],W_{Ei}\in {\mathbb{R}}^{\frac{V}{2}\times H}$$
     • GPU 0 負責查找詞表範圍
       $0\sim \frac{V}{2}$
     • GPU 1 負責查找詞表範圍
       $\frac{V}{2}\sim V$
   • 嵌入主要是查找操作，可視為一個 矩陣乘法
     
     $$X=W_E^T\cdot O$$
   • 在每張 GPU 上計算局部嵌入
     
     $$X_i=W_{E_i}^T\cdot O_i,X_i\in {\mathbb{R}}^{H\times BS}$$
   • 因為每張 GPU 只負責一部份向量查找，所以需要透過 All Reduce 聚合結果
     
     $$X=\sum _{i=1}^N{X}_i,X\in {\mathbb{R}}^{H\times BS}$$

Output embedding層

一般並行做法

一般輸出輸出嵌入層（Output Embedding Layer） 的維度通常是 Hidden Size(

$H$) × Vocabulary Size(

$V$)

$$\text{Output Embedding}\in {\mathbb{R}}^{H\times V}$$
在 Megatron-LM 中，對其進行並行化處理，再者，因為 Transformer語言模型中輸入嵌入與輸出嵌入共用權重，因此都會修改計算方式

• 原本Output embedding層的形狀(與Input embedding層共享權重)
  
  $$W_O\in {\mathbb{R}}^{H\times V}$$
• 輸出層的矩陣運算為:
  
  $$Y=X{W}_O$$
  其中：
  • $X\in {\mathbb{R}}^{B\times S\times H}$（Transformer 最後一層的 hidden state）
  • $W_O\in {\mathbb{R}}^{H\times V}$
  • $Y\in {\mathbb{R}}^{B\times S\times V}$（最終 logits，對應詞彙表中的每個詞的預測分數）
• 將Output embedding層進行列方向切分
  • 將
    $W_O$ 沿著詞表維度
    $V$進行列方向切分
    
    $$W_O=[W_{O_1},W_{O_2}],W_{O_i}\in {\mathbb{R}}^{H\times V/2}$$
  • 每張 GPU 只對自己負責的詞表進行 logits 計算
    
    $$Y_i=X{W}_{O_i},Y_i\in {\mathbb{R}}^{B\times S\times V/2}$$
  • 每張 GPU 的輸出大小
    $(B,S,V/2)$，比完整的
    $(B,S,V)$ 小
    $2$ 倍
• 最終的 logits
  $Y_i$ 要完整的
  $V$ 才可以計算交叉熵損失，因此要 All Gather 來合併計算結果
  
  $$Y=\text{All-Gather}([Y_1,Y_2])$$

進一步優化(減少通信開銷) - Megatron-LM : 交叉熵融合(Loss Fusion)

避免 All-Gather，讓每張 GPU 直接計算部分 logits 的交叉熵，只通信標量損失值，而不是完整 logits
交叉熵計算公式:

$$L=-\sum _{j=1}^V{t}_j\log p_j$$

• $t_j$ 是目標 Token 的 one-hot 編碼，如果正確的詞是
  $j^{\ast }$，則
  $t_{j^{\ast }}=1$，其餘為 0
• $p_j$ 是 softmax 後的機率分布
  
  $$p_j=\frac{e^{Y_j}}{\sum _{k=1}^V{e}^{Y_k}}$$

直接在每張 GPU 上計算部分 logits 的交叉熵

• 每張 GPU 只存
  $W_O$ 的一部分，並計算對應的 logits
  $Y_i$
  
  $$Y_1=X{W}_{O_1},Y_2=X{W}_{O_2},Y_1,Y_2\in {\mathbb{R}}^{B\times S\times V/2}$$
  1. 每張 GPU 先找出他自己負責部分的最大值
     $Y_{max,i}$
     
     $$Y_{max,i}=max{Y}_i,Y_{max,i}\in {\mathbb{R}}^{B\times S}$$
  2. All-Reduce 操作獲得全局最大值
     
     $$Y_{max}=\text{All-Reduce}(max(Y_1,Y_2))$$
  3. 計算局部 softmax 分子(穩定數值修正，避免數值溢出)
     
     $${\tilde{Y}}_i=Y_i-Y_{max}$$
     • 這邊還是 partial logits
  4. 計算 softmax 分母
     
     $$Z_i=\sum _{j=1}^{V/2}{e}^{{\tilde{Y}}_{i,j}}$$
  5. All-Reduce，獲得完整的歸一化因子
     
     $$Z=\text{All-Reduce}(Z_1+Z_2)$$
  6. 計算部分 softmax 機率
     
     $$p_i=\frac{e^{{\tilde{Y}}_i}}{Z}$$
  7. 計算交叉熵損失
     
     $$L_i=-\sum _{j=1}^{V/2}{t}_j\log p_j$$
  8. All-Reduce，獲得總損失
     
     $$L=\text{All-Reduce}(L_1+L_2)$$

總結

步驟	計算內容	All-Reduce？
計算局部 logits	$Y_i=X{W}_{O_i}$	否
計算最大值	$Y_{max,i}=max(Y_i)$	是
計算 softmax 分母	$Z_i=\sum \exp ({\tilde{Y}}_i)$	是
計算 softmax 機率	$P_i=\frac{\exp ({\tilde{Y}}_i)}{Z}$	否
計算局部交叉熵	$L_i=-\log P(y^{\ast })$	否
最終合併損失	$L=\sum L_i$	是