# 张量并行（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）需要在切分后做特殊处理保持语义一致