10.3 张量并行（Tensor Parallelism）

数据并行通过"模型复制，数据分片"来提升吞吐量，但它有一个隐含前提：每张 GPU 必须能装下完整的模型。当单个 Transformer 层的权重矩阵大到超出单卡显存时，数据并行即宣告失效。流水线并行通过按层切分模型解决了部分问题，但它引入了难以消除的流水线气泡。张量并行（Tensor Parallelism, TP）则提供了一种更细粒度的切分策略——它不在层的边界上切分，而是深入到单个算子内部，将一次矩阵乘法本身拆分到多张 GPU 上并行执行。

本节将从线性层的矩阵切分原理出发，分别讲解列并行和行并行的数学机制与通信模式，再展示 Megatron-LM 如何将两者巧妙组合以最小化通信开销，最后讨论张量并行在实践中的配置选择与适用边界。

10.3.1 矩阵切分的两种基本方式

张量并行的核心对象是线性层的矩阵乘法 $Y = X W$ ，其中 $X \in R^{b \times m}$ 是输入， $W \in R^{m \times n}$ 是权重矩阵。根据对 $W$ 的切分维度不同，存在两种基本并行方式。以下假设使用 2 张 GPU 进行说明，推广到 $p$ 张 GPU 是直接的。

列并行（Column Parallelism）。 将权重矩阵 $W$ 沿列维度等分为两块：

W = [W_{1}, W_{2}], W_{1} \in R^{m \times n / 2}, W_{2} \in R^{m \times n / 2}

GPU 0 持有 $W_{1}$ ，GPU 1 持有 $W_{2}$ 。前向传播时，两张 GPU 各自独立地用完整的输入 $X$ 乘以自己的权重分片：

Y_{1} = X W_{1}, Y_{2} = X W_{2}

根据分块矩阵乘法的性质，完整结果可通过拼接得到： $Y = [Y_{1}, Y_{2}]$ 。此即列并行的核心洞察——前向传播不需要任何跨 GPU 通信，每张 GPU 独立计算出输出张量的一段列。代价出现在反向传播：计算输入 $X$ 的梯度时，需要将各 GPU 上的局部梯度通过一次 AllReduce 求和汇总。

用示意图表示列并行的数据流：

输入 X (完整, 两张 GPU 各持有一份副本)
  │                │
  ▼                ▼
[X × W_1]       [X × W_2]      ← 各 GPU 独立计算, 无通信
  │                │
  ▼                ▼
  Y_1              Y_2          ← 输出按列切分

行并行（Row Parallelism）。 将权重矩阵 $W$ 沿行维度等分：

W = [\begin{matrix} W_{1} \\ W_{2} \end{matrix}], W_{1} \in R^{m / 2 \times n}, W_{2} \in R^{m / 2 \times n}

行并行有一个前提：输入 $X$ 必须已经按列切分好，即 $X = [X_{1}, X_{2}]$ ，其中 $X_{i} \in R^{b \times m / 2}$ 。GPU 0 持有 $X_{1}$ 和 $W_{1}$ ，GPU 1 持有 $X_{2}$ 和 $W_{2}$ 。前向传播时：

Y = X W = X_{1} W_{1} + X_{2} W_{2}

每张 GPU 计算出一个部分和（ $X_{i} W_{i}$ ），要得到最终结果 $Y$ ，必须执行一次 AllReduce 将两个部分和累加。AllReduce 完成后，两张 GPU 上都持有相同的完整输出 $Y$ 。反向传播时则无需跨 GPU 通信。

输入 X_1 (切分)    输入 X_2 (切分)
  │                  │
  ▼                  ▼
[X_1 × W_1]       [X_2 × W_2]    ← 各 GPU 独立计算
  │                  │
  └──── AllReduce ───┘             ← 部分和累加
           │
           ▼
      Y (完整结果, 两张 GPU 均持有)

两种方式的对称性。 列并行的前向无通信、反向需通信；行并行的前向需通信、反向无通信。这种对称性为后续的组合设计提供了基础。

10.3.2 通信原语在张量并行中的角色

张量并行涉及的核心集合通信操作是 AllReduce 和 AllGather，它们分别出现在不同的场景中。

AllReduce 是张量并行中最核心的操作。它对所有 GPU 上的同形状张量执行逐元素求和（或其他规约操作），并将结果写回每一张 GPU。在行并行的前向传播中，各 GPU 的部分和 $X_{i} W_{i}$ 需要通过 AllReduce 累加得到完整的 $Y$ ；在列并行的反向传播中，各 GPU 的局部梯度也需要通过 AllReduce 汇总。AllReduce 的通信量为 $2 (p - 1) / p \cdot | Y |$ （Ring-AllReduce 实现），其中 $p$ 是 GPU 数量， $| Y |$ 是结果张量的大小。

AllGather 在张量并行的某些变体中出现。如果列并行的输出不进行拼接而是保持切分状态，下游操作可能需要一次 AllGather 将各 GPU 的片段拼接为完整张量。AllGather 的通信量为 $(p - 1) / p \cdot | Y |$ 。在 Megatron-LM 的经典 TP 实现中，通过巧妙的行列组合设计，主要使用的是 AllReduce 而非 AllGather。

Identity（恒等）操作。 这不是一个通信操作，而是一个概念标记：将完整输入复制到所有 GPU 上，或将输出不经修改地传递。它在 Megatron-LM 的符号体系中用 $f$ 表示，与 AllReduce 操作 $g$ 形成前向/反向的互补配对。

由于 AllReduce 需要所有参与的 GPU 进行同步和大量数据交换，张量并行对网络带宽和延迟极为敏感。这也是 TP 几乎只部署在单节点内部（利用 NVLink/NVSwitch 提供 600-900 GB/s 的双向带宽）而非跨节点使用的根本原因。

10.3.3 Megatron-LM 的张量并行设计

NVIDIA 的 Megatron-LM 框架是张量并行的经典实现。它的核心贡献在于：针对 Transformer 的 MLP 模块和自注意力模块，精心设计了列并行与行并行的组合方式，使得每个子模块在一次前向传播中仅需一次 AllReduce 通信。

MLP 模块的切分

一个标准的 Transformer MLP 由两个线性层和一个非线性激活组成：

Y = GeLU (X A) \cdot B

其中 $A \in R^{h \times 4 h}$ 是第一个线性层的权重， $B \in R^{4 h \times h}$ 是第二个线性层的权重（ $h$ 为隐藏维度）。Megatron-LM 对这两个矩阵采用互补的切分策略：

矩阵 $A$ 按列切分： $A = [A_{1}, A_{2}]$ 。GPU 0 持有 $A_{1}$ ，GPU 1 持有 $A_{2}$ 。
矩阵 $B$ 按行切分： $B = [B_{1}; B_{2}]$ （分号表示垂直堆叠）。GPU 0 持有 $B_{1}$ ，GPU 1 持有 $B_{2}$ 。

前向传播流程如下：

输入 X (完整) ── f (Identity, 复制到两张 GPU) ──┐
                                                 │
  GPU 0: Y_1 = GeLU(X × A_1)                    │  GPU 1: Y_2 = GeLU(X × A_2)
         Z_1 = Y_1 × B_1                        │         Z_2 = Y_2 × B_2
                │                                │                │
                └──────── g (AllReduce) ─────────┘────────────────┘
                                  │
                                  ▼
                         Z = Z_1 + Z_2 (完整结果)

整个过程的关键在于：

第一层用列并行： $A$ 按列切分，前向不需要通信。 $GeLU$ 是逐元素操作，直接作用于切分后的结果 $[X A_{1}, X A_{2}]$ 与作用于完整结果 $X A$ 再切分是等价的，即 $GeLU (X A_{i}) = [GeLU (X A)]_{i}$ 。这个性质至关重要——它保证了非线性激活不会破坏列切分的正确性。
第二层用行并行：列并行的输出天然按列切分为 $[Y_{1}, Y_{2}]$ ，恰好满足行并行的输入要求。 $B$ 按行切分后，各 GPU 计算部分和，最终通过一次 AllReduce 得到完整结果。

通过这种"列并行 → 行并行"的级联设计，整个 MLP 模块在前向传播中只需要一次 AllReduce，而非每个线性层各一次。

反向传播的对称性。 定义辅助算子 $f$ 和 $g$ ： $f$ 在前向为 Identity（复制），在反向为 AllReduce（梯度求和）； $g$ 在前向为 AllReduce，在反向为 Identity。MLP 模块前端是 $f$ ，后端是 $g$ ，因此前向传播执行一次 AllReduce（在 $g$ 处），反向传播也执行一次 AllReduce（在 $f$ 处）。每个 MLP 模块总共需要 2 次 AllReduce（一次前向，一次反向）。

自注意力模块的切分

多头注意力（Multi-Head Attention）天然适合张量并行，因为多个注意力头之间本就是独立计算的。Megatron-LM 的做法是：

将 $Q$ 、 $K$ 、 $V$ 三个投影矩阵按列切分（即按注意力头分组），每张 GPU 负责若干个完整的注意力头。
各 GPU 独立完成自己负责的注意力头的全部计算（包括 $Q K^{T}$ 矩阵乘法、Softmax、与 $V$ 的加权求和）。
输出投影矩阵按行切分，将各 GPU 的部分结果通过一次 AllReduce 汇总。

与 MLP 模块类似，自注意力模块也只需要前向一次 AllReduce + 反向一次 AllReduce。

单个 Transformer 层的总通信量

一个 Transformer 层包含一个自注意力模块和一个 MLP 模块，因此总共需要 4 次 AllReduce（每个模块前向+反向各 1 次）。在 Ring-AllReduce 实现下，每次 AllReduce 的数据量为 $2 (p - 1) / p \cdot b \cdot s \cdot h$ （其中 $b$ 为批大小， $s$ 为序列长度， $h$ 为隐藏维度， $p$ 为 TP 并行度）。

10.3.4 Megatron-LM 的代码与配置

在 Megatron-LM 的实际实现中，张量并行通过替换标准的 nn.Linear 为两种特殊的并行线性层来实现。

ColumnParallelLinear 和 RowParallelLinear 是两个核心模块。以 MLP 为例，其实现逻辑可概括为：

python

# 教学示例：展示核心逻辑，省略了部分 import 和辅助函数定义
class ParallelMLP(nn.Module):
    def __init__(self, hidden_size, ffn_hidden_size, tp_size):
        # 第一个线性层: 列并行, 将权重按列切分到 tp_size 张 GPU
        self.dense_h_to_4h = ColumnParallelLinear(
            hidden_size, ffn_hidden_size,
            gather_output=False  # 不拼接输出, 保持切分状态传给下一层
        )
        # 第二个线性层: 行并行, 将权重按行切分
        self.dense_4h_to_h = RowParallelLinear(
            ffn_hidden_size, hidden_size,
            input_is_parallel=True  # 输入已经是切分状态
        )

    def forward(self, x):
        # f: Identity(前向复制) / AllReduce(反向梯度汇总)
        intermediate = self.dense_h_to_4h(x)  # 列并行, 前向无通信
        intermediate = F.gelu(intermediate)
        # g: AllReduce(前向结果汇总) / Identity(反向梯度直传)
        output = self.dense_4h_to_h(intermediate)  # 行并行, 前向AllReduce
        return output

ColumnParallelLinear 内部做的事情：在初始化时，将完整的权重矩阵 $(h, 4 h)$ 按列切分为 $(h, 4 h / p)$ ，每个 TP rank 只存储和计算自己的分片。gather_output=False 表示输出保持切分状态，避免不必要的 AllGather。

RowParallelLinear 内部做的事情：将权重矩阵 $(4 h, h)$ 按行切分为 $(4 h / p, h)$ ，计算部分和后执行 AllReduce。

启动配置。 在 Megatron-LM 中，通过命令行参数指定张量并行度：

bash

python pretrain_gpt.py \
    --tensor-model-parallel-size 8 \    # TP 并行度, 通常 = 节点内 GPU 数
    --pipeline-model-parallel-size 4 \  # PP 并行度
    --micro-batch-size 2 \
    --global-batch-size 1024 \
    --hidden-size 12288 \
    --num-attention-heads 96 \
    --num-layers 96 \
    --seq-length 2048 \
    --distributed-backend nccl

其中 --tensor-model-parallel-size 控制 TP 的并行度。Megatron-LM 在初始化时会自动创建 TP 进程组——同一节点内的 GPU 被分为一个 TP 组，组内的 GPU 通过 NCCL 的 AllReduce 进行通信。

10.3.5 实践选择与适用边界

TP 并行度的选择。 Megatron-LM 团队的实验数据表明，8 路张量并行通常是性能的最佳平衡点。2 路和 4 路 TP 在减少单卡显存占用的同时通信开销较低；超过 8 路后，AllReduce 的通信量线性增长，而单卡计算量相应缩小，通信占比急剧升高，吞吐量反而下降。因此，TP 并行度通常等于单节点内的 GPU 数量（8 路对应一台 DGX A100/H100）。

与其他并行策略的对比。

维度	张量并行	流水线并行	数据并行
切分粒度	单个算子内部（矩阵列/行）	按层（模型深度）	不切分模型，切分数据
通信频率	每层多次（前向+反向各 1-2 次 AllReduce）	仅在层边界传递激活值	仅在梯度同步时通信
气泡问题	无	有（流水线启动/排空阶段）	无
对带宽的要求	极高（需 NVLink 级别）	中等	中等
典型部署范围	节点内（2/4/8 路）	跨节点	跨节点
对批大小的要求	无	需要较大批大小以减少气泡占比	需要较大全局批大小

表 10-2：三种主要并行策略的对比。

TP 的局限性。 张量并行无法处理 Transformer 中所有的显存瓶颈。LayerNorm 和 Dropout 等操作不涉及大矩阵乘法，它们的激活值在标准 TP 下并未被切分，仍然需要在每张 GPU 上保存完整的副本。Megatron-LM 后续通过**序列并行（Sequence Parallelism）**补全了这一短板——在 TP 组内，对 LayerNorm 和 Dropout 的输入/输出沿序列维度进行切分，使得几乎所有激活显存都能随 TP 并行度线性下降。

在 3D 并行中的位置。 在大规模训练的标准实践中，张量并行是 3D 并行拓扑 $(d p, p p, t p)$ 的最内层维度。它部署在硬件连接最紧密的位置——同一节点内通过 NVLink 互连的 GPU 之间。流水线并行占据跨节点的维度，数据并行则在最外层复制整个"TP+PP"组合体以扩展吞吐量。这种分层部署的原则是：通信频率越高的并行维度，越应当部署在带宽越高的硬件层级上。

10.3.6 本节小结

张量并行是训练超大模型时不可或缺的并行策略。它的核心思想可以归纳为三个要点：

矩阵切分的两种方式：列并行将权重按列切分，前向无通信、反向需 AllReduce；行并行将权重按行切分，前向需 AllReduce、反向无通信。两者在通信时机上互为对称。
Megatron-LM 的组合设计：通过将 MLP 和自注意力模块中的两个连续线性层分别采用列并行和行并行，实现了"一进一出"的级联配对，使每个子模块在前向和反向各仅需一次 AllReduce，将通信开销压缩到最低。
适用边界清晰：由于每层都需要多次 AllReduce 同步，张量并行对通信带宽极度敏感，几乎只适用于节点内部署（NVLink 互连），典型并行度为 2/4/8 路。它是 3D 并行拓扑的最内层维度，与流水线并行和数据并行形成互补，共同支撑起千亿至万亿参数模型的高效训练。

10.3 张量并行（Tensor Parallelism） ​

10.3.1 矩阵切分的两种基本方式 ​

10.3.2 通信原语在张量并行中的角色 ​

10.3.3 Megatron-LM 的张量并行设计 ​

MLP 模块的切分 ​

自注意力模块的切分 ​

单个 Transformer 层的总通信量 ​

10.3.4 Megatron-LM 的代码与配置 ​

10.3.5 实践选择与适用边界 ​

10.3.6 本节小结 ​