88 lines
3.5 KiB
Markdown
88 lines
3.5 KiB
Markdown
# 张量并行（Tensor Parallelism）
|
||
|
||
## 概述
|
||
|
||
张量并行（Tensor Parallelism，TP）是将单个计算任务（如矩阵乘法）的张量切分到多个GPU上并行执行的技术。主要用于训练参数量巨大的模型（如千亿参数LLM），通过将模型参数沿某个维度切分，使每个GPU只负责存储和计算的一部分。
|
||
|
||
## 核心切分策略
|
||
|
||
### 1. Column-Bert切分（列切分）
|
||
|
||
**原理**：将权重矩阵按列维度切分，每个GPU持有部分列的权重。
|
||
|
||
```
|
||
原始权重 W ∈ R^(d_hidden × d_ff)
|
||
切分后: W1 ∈ R^(d_hidden × d_ff/2), W2 ∈ R^(d_hidden × d_ff/2)
|
||
```
|
||
|
||
**计算方式**：
|
||
- 输入X保持完整，在每个GPU上分别计算 `Y1 = X @ W1`, `Y2 = X @ W2`
|
||
- 最终结果 `Y = [Y1, Y2]`（沿列维度拼接）
|
||
|
||
**适用场景**：适合 Linear 层中权重形状为 (hidden_size, intermediate_size) 的投影层（如 QKV 投影、FFN 第一个线性层）
|
||
|
||
**代表实现**：Megatron-LM 的 ColumnParallelLinear
|
||
|
||
### 2. Row-Bert切分（行切分）
|
||
|
||
**原理**：将权重矩阵按行维度切分，每个GPU持有部分行的权重。
|
||
|
||
```
|
||
原始权重 W ∈ R^(d_ff × d_hidden)
|
||
切分后: W1 ∈ R^(d_ff/2 × d_hidden), W2 ∈ R^(d_ff/2 × d_hidden)
|
||
```
|
||
|
||
**计算方式**：
|
||
- 各GPU分别计算 `Y1 = X @ W1^T`, `Y2 = X @ W2^T`
|
||
- 最终结果 `Y = Y1 + Y2`（AllReduce求和）
|
||
|
||
**适用场景**：适合 FFN 第二个线性层（intermediate_size → hidden_size），因为其输出需要求和
|
||
|
||
**代表实现**：Megatron-LM 的 RowParallelLinear
|
||
|
||
### 3. 通信分析
|
||
|
||
| 切分方式 | 通信操作 | 通信量 |
|
||
|---------|---------|--------|
|
||
| Column切分 | AllReduce | O(sequence_length × hidden_size) |
|
||
| Row切分 | AllReduce | O(sequence_length × hidden_size) |
|
||
| All-to-All（集合通信） | 用于跨层级联时 | 视实现而定 |
|
||
|
||
通信主要发生在每层的前向传播和反向传播中，由于通信带宽成为瓶颈，通常要求 NVLink 或高速互联。
|
||
|
||
## 3D并行中的角色
|
||
|
||
张量并行是 3D 并行（数据并行 + 流水线并行 + 张量并行）的一维：
|
||
|
||
- **Pipeline Parallelism（流水线并行）**：按模型层垂直切分，不同Stage的GPU负责不同层的计算
|
||
- **Tensor Parallelism（张量并行）**：对单层内部运算水平切分，同一Stage内不同GPU负责不同张量分片
|
||
- **Data Parallelism（数据并行）**：相同配置复制多份，处理不同batch数据
|
||
|
||
典型配置（如NVIDIA Megatron框架）：
|
||
```
|
||
tensor_parallelism = 8 # 每层在8个GPU间切分
|
||
pipeline_parallelism = 4 # 模型分为4个stage
|
||
data_parallelism = N/(8*4) # 剩余GPU用于数据并行
|
||
```
|
||
|
||
## 实现框架
|
||
|
||
| 框架 | 实现方式 |
|
||
|------|---------|
|
||
| Megatron-LM | 专门优化TP，支持all-reduce融合、异步通信 |
|
||
| DeepSpeed | 通过ZeRO-Inference支持TP，但主要focus在ZeRO |
|
||
| PyTorch FSDP | 采用分片策略，非传统TP |
|
||
| NeMo-Megatron | 基于Megatron封装的分布式训练框架 |
|
||
|
||
## 关键技术点
|
||
|
||
1. **AllReduce通信重叠**：利用计算掩盖通信
|
||
2. **梯度同步**：反向传播时需要对各TP分片的梯度做AllReduce
|
||
3. **模型一致性**：需要保证所有分片在相同输入上行为一致
|
||
4. **Memory efficiency**：激活值和分片权重在GPU间分布，降低单卡峰值内存
|
||
|
||
## 注意事项
|
||
|
||
- 张量并行度增加会显著提升通信开销，超过8卡通常收益递减
|
||
- GPU拓扑（NVLink vs PCIe）直接影响TP效率
|
||
- 某些操作（如BatchNorm、LayerNorm）需要在切分后做特殊处理保持语义一致 |