3.5 KiB
3.5 KiB
张量并行（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封装的分布式训练框架 |
关键技术点
- AllReduce通信重叠：利用计算掩盖通信
- 梯度同步：反向传播时需要对各TP分片的梯度做AllReduce
- 模型一致性：需要保证所有分片在相同输入上行为一致
- Memory efficiency：激活值和分片权重在GPU间分布，降低单卡峰值内存
注意事项
- 张量并行度增加会显著提升通信开销，超过8卡通常收益递减
- GPU拓扑（NVLink vs PCIe）直接影响TP效率
- 某些操作（如BatchNorm、LayerNorm）需要在切分后做特殊处理保持语义一致