7.3 模型并行与张量并行：拆分权重的艺术

当单个 Transformer 层的参数量太大——即使用 ZeRO 也需要频繁的参数通信时——模型并行（Model Parallelism）提供了另一种思路：将模型的计算本身分配到多张 GPU 上。

7.3.1 张量并行的核心思想

张量并行（Tensor Parallelism，TP）由 NVIDIA 的 Megatron-LM 项目提出，核心思想是将单层内的权重矩阵切分到多张 GPU 上，协同完成矩阵运算。

以一个线性层 $Y = XW$ 为例，如果 $W \in \mathbb{R}^{d \times d}$，可以将 $W$ 按列切分为两半：$W = [W_1, W_2]$，分别放在两张 GPU 上。每张 GPU 计算 $Y_i = XW_i$，得到输出的一半。最后拼接结果：$Y = [Y_1, Y_2]$。

Multi-Head Attention 天然适合张量并行：不同的注意力头本来就是独立计算的，只需将不同的头分配到不同 GPU 上即可。

7.3.2 通信模式

张量并行的通信发生在层内——每个线性层的计算都需要在参与的 GPU 之间进行一次 AllReduce 或 AllGather 操作。由于 Transformer 每层包含多个线性层（注意力和 FFN 中各有两个），通信频率较高。

因此，张量并行通常只在同一节点内的 GPU 之间使用（利用 NVLink 的高带宽和低延迟），不跨节点部署。

7.3.3 与数据并行的互补

实际的大规模训练通常同时使用张量并行和数据并行：

• 节点内：使用张量并行（如在 8 张 GPU 之间进行 8 路 TP），解决单层权重太大的问题

• 节点间：使用数据并行（配合 ZeRO），扩展训练吞吐量

这种组合策略是当前主流大模型训练框架（如 Megatron-LM、DeepSpeed）的标配。