# <div align="center"> Tensor-Parallel-Training </div>
[論文參考](https://arxiv.org/abs/1909.08053)
## 張量並行
**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 = \begin{bmatrix} B_1 \\ B_2 \end{bmatrix}\\]
- 計算變為:
\\[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 = \begin{bmatrix} B_1 \\ B_2 \end{bmatrix}\\]
- 在不同 GPU 上:
\\[
\frac{\partial \mathcal{L}}{\partial B_1} = Y_1^T \delta_{YB}, \quad \frac{\partial \mathcal{L}}{\partial B_2} = Y_2^T \delta_{YB}
\\]
- 這些梯度可以在 各 GPU 上獨立計算，不需要 All-Reduce
- 對 \\(Y\\) 求梯度
\\[Z = Y B\\]
\\[\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, \quad \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), \quad Y_2 = \text{GeLU}(X A_2)\\]
- 鏈式規則
\\[
\delta_{X A_1} = \delta_{Y_1} \cdot \text{GeLU}^{\prime}(X A_1), \quad \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}, \quad \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}}, \quad K_{\text{local}} = X W_K^{\text{local}}, \quad 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 = \begin{bmatrix} W_{O_1} \\ W_{O_2} \end{bmatrix}\\]
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 B S} \\)** 是 one-hot 編碼的輸入 Token
- 進行行方向分割
\\[
W_E = \begin{bmatrix} W_{E1} \\ W
_{E2} \end{bmatrix}, \quad 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, \quad X_i \in \mathbb{R}^{H \times B S}
\\]
- 因為每張 GPU 只負責一部份向量查找，所以需要透過 All Reduce 聚合結果
\\[X = \sum_{i=1}^{N} X_i, \quad X \in \mathbb{R}^{H \times B S}\\]

## 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=XW_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}], \quad W_{O_i}\in \mathbb{R}^{H \times V /2}\\]
- 每張 GPU 只對自己負責的詞表進行 logits 計算
\\[
Y_i = X W_{O_i}, \quad 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^* \\)，則 \\( t_{j^*} = 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=XW_{O_1}, \quad Y_2=XW_{O_2}, \quad Y_1, Y_2 \in \mathbb{R}^{B \times S \times V/2}\\]
1. 每張 GPU 先找出他自己負責部分的最大值 \\(Y_{max, i}\\)
\\[Y_{\max, i} = \max Y_i, \quad 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^*) \\) | 否 |
| **最終合併損失** | \\( L = \sum L_i \\) | 是 |