跳到主要内容

张量并行（Tensor Parallelism）

张量并行（TP）由 Megatron-LM 提出（Shoeybi et al., 2019），是推理场景中最常用的并行策略之一，主要用于在同一节点内的高带宽互联设备（如 NVLink）上切分单层的权重矩阵。

基本原理​

TP 将每层的权重矩阵分为列并行和行并行两种切分方式，交替使用以最小化通信次数。

MLP 块（以两层 Linear 为例）：

Input X
  → [Column-parallel Linear W1]   # 各节点持有 W1 的列分片，独立计算
  → GeLU
  → [Row-parallel Linear W2]      # 各节点持有 W2 的行分片，独立计算部分积
  → AllReduce                     # 合并各节点的部分积，得到完整输出
  → Output

Attention 块：

Input X
  → [Column-parallel QKV projection]   # Q/K/V 各按头切分
  → Scaled Dot-Product Attention       # 各节点独立计算自己分配的头
  → [Row-parallel Output projection]   # 按行切分，计算部分积
  → AllReduce                          # 合并部分积
  → Output

每个 Transformer 层需要 2 次 AllReduce（MLP 和 Attention 各一次）。

通信时机​

阶段AllReduce 次数说明
推理前向每层 2 次Attention 输出 + MLP 输出各一次
训练前向每层 2 次同上
训练反向每层 2 次梯度也需要 AllReduce 合并
训练合计每层 4 次前向 + 反向

通信量分析​

每次 TP AllReduce 的消息大小等于该层激活值张量的大小：

$M_{\text{TP}} = b \times s \times h \times \text{dtype\_size}$

其中：

  • $b$：micro-batch size
  • $s$：序列长度（sequence length）
  • $h$：hidden dimension
  • $\text{dtype\_size}$：数据类型字节数（BF16 = 2 字节，FP32 = 4 字节）

典型模型的通信量（$b=1$，$s=4096$，BF16）：

模型$h$每次 AllReduce 消息大小
LLaMA-70B819267.1 MB
DeepSeek-V3716858.7 MB
Qwen-72B819267.1 MB

这些均属于大消息（远超 $\alpha$-$\beta$ 模型的带宽-延迟交叉点），应使用 Ring AllReduce 算法。

对延迟的影响​

TP AllReduce 位于关键路径上——每层的计算必须等待 AllReduce 完成才能进入下一层，无法在推理前向传播中与同层计算重叠。

TP 通信特征汇总：

特征值
通信原语AllReduce
消息大小10 ~ 100 MB（大消息）
通信组大小通常 2 ~ 8（同 board/node 内）
频率推理每层 2 次，训练每层 4 次
延迟敏感性高（在关键路径上，推理无法 overlap）
推荐算法Ring AllReduce（大消息带宽最优）
推荐拓扑高带宽互联（NVLink / C2C）

TP 度越大通信开销越高的原因​

Ring AllReduce 的延迟公式：

$T_{\text{Ring-AR}} = 2(N-1)\alpha + \frac{2(N-1)}{N} \cdot \frac{M}{\beta}$

当 TP 度 $N$ 增大时：

  • 延迟项 $2(N-1)\alpha$ 线性增长：$N=2$ 时 $2\alpha$，$N=8$ 时 $14\alpha$
  • 带宽项 $\frac{2(N-1)}{N}$ 趋向 2（渐近最优），带宽开销增长放缓

实际瓶颈：大消息场景下带宽项主导，每次 AllReduce 的时间约为 $\frac{2M}{\beta}$，与 $N$ 关系不大；但更高的 TP 度意味着每个设备分配到的权重分片更小，单设备算力利用率下降，通信/计算比上升。

推理 vs 训练的影响差异​

推理训练
AllReduce 次数/层2（前向）4（前向 + 反向）
batch size通常较小较大
消息大小中等～大大
实时性要求高（影响 TTFT/TPOT）较低
Overlap 机会少（前向无法 overlap）有（反向可与下一层梯度 overlap）

参考文献​