人人都会AI编程

9.3 张量并行（TP）：模型横向切分的原理与通信逻辑

更新时间：2026-07-09

在 9.2 节中，数据并行（DP）让每个 GPU 都持有完整的模型副本，只拆分训练数据。它的前提很现实：单卡显存必须装得下整个模型。当模型参数量从 7B 增长到 70B、405B 时，单张 A100/H100 的 80 GB 显存早已成为瓶颈。此时，你必须把模型本身拆开。

张量并行（Tensor Parallelism, TP）就是解决这个问题的第一层手段。它通常被称为层内并行（Intra-layer Parallelism）或横向切分——把 Transformer 层内部的巨大矩阵计算分散到多张 GPU 上同时执行。

一、TP 解决的核心矛盾

单卡显存的硬约束可以用一个粗略公式概括：

单卡显存占用 ≈ 模型参数 + 优化器状态 + 梯度 + 激活值

以 175B 参数模型、FP16 混合精度训练为例：

  • 模型参数：约 350 GB（2 bytes/参数）
  • Adam 优化器状态+梯度：约 1050 GB（FP32 主权重 + 动量 + 二阶矩 + FP16 梯度）
  • 激活值：随序列长度和批量大小变化，百 GB 量级

总计远超单卡 80 GB。DP 对此无能为力，因为它只是复制。TP 的思路是：既然装不下，就把矩阵切成几块，每块 GPU 只算一部分。

二、切分原理：以 Megatron-LM 的经典 1D TP 为例

当前工业界最主流的 TP 实现源自 NVIDIA Megatron-LM 的 1D 张量并行。它的核心思想极其简洁：把线性层的权重矩阵按列或按行切开，让每张卡只负责一部分矩阵乘法，再通过通信拼接结果。

1. MLP 块的切分（列切 + 行切）

Transformer 中的 FFN（前馈网络）通常包含两个线性层：Y = GeLU(X · A) · B。

  • 第一层（A）按列切开：把权重矩阵 A 竖着切成 A₁ 和 A₂。每张卡各自用完整的输入 X 乘以自己那部分权重，得到 GeLU(X·A₁) 和 GeLU(X·A₂)。

结果：两张卡的输出在特征维度上天然互补，可以直接拼接，前向传播无需通信。

  • 第二层（B）按行切开：把权重矩阵 B 横着切成 B₁ 和 B₂。此时卡 1 上有 GeLU(X·A₁) · B₁，卡 2 上有 GeLU(X·A₂) · B₂。要得到最终正确输出，必须把两部分加起来：Output = (GeLU·A₁)·B₁ + (GeLU·A₂)·B₂。

结果：前向传播结尾需要一次 AllReduce（求和）。

反向传播时，梯度流动刚好相反：先有一次 AllReduce 聚合梯度，然后分别回传，中间激活值也按对应维度拆分存储。

2. Attention 块的切分（切头天然并行）

多头注意力（MHA）是 TP 最友好的结构。

  • Q/K/V 投影层按列切：把多个注意力头的 Q/K/V 权重矩阵按头维度切成 N 份。每张卡负责一部分头的计算。因为各头的计算本来就独立，卡与卡之间前向无需通信。
  • 输出投影层（Output Linear）按行切：与 MLP 的第二层同理，各卡算完各自头的结果后，需要通过 AllReduce 求和得到最终输出。

关键工程细节：TP 的切分数（TP size）通常必须整除注意力头的数量。例如 Llama 3 70B 有 64 个注意力头，TP=8 可以完美均分（每卡 8 个头），但 TP=7 就会引入复杂的 padding 和负载不均，主流框架会直接拒绝这种配置。

三、通信逻辑：TP 的“隐性成本”

TP 的通信模式非常规律，可以用“一次前向、一次反向、两次 AllReduce”来概括每个 Transformer 层。

| 阶段 | 操作位置 | 通信原语 | 通信量（每卡） |
|------|----------|----------|----------------|
| 前向 | Attention Output Linear 后 | AllReduce | batch_size × seq_len × hidden_size × sizeof(dtype) |
| 前向 | MLP 第二层后 | AllReduce | 同上 |
| 反向 | MLP 第二层梯度 | AllReduce | 同上 |
| 反向 | Attention Output 梯度 | AllReduce | 同上 |

通信量分析：

  • TP 的通信量只与激活值尺寸成正比，与模型参数量无关。这意味着当模型变宽（hidden size 增大）时，通信量线性增长，但不会因为参数量从 7B 涨到 70B 而爆炸。
  • 单次 AllReduce 的数据量通常只有几十到几百 MB（取决于 batch 和序列长度），远小于参数同步（如 DP 中的梯度 AllReduce 可能涉及数 GB）。

但 TP 对通信带宽极度敏感：

  • 因为 AllReduce 发生在每个 Transformer 层的每次前向+反向中，属于高频细粒度通信。
  • 如果 GPU 之间通过 PCIe（32 GB/s 量级）互联，TP 会成为严重瓶颈；只有通过 NVLink（900 GB/s 量级，如 NVLink4） 才能将通信延迟掩盖在计算之后。

这就是生产环境中的铁律：TP 几乎只部署在单节点内部。 例如，一台 8 卡 H100 服务器内部通过 NVLink/ NVSwitch 全互联，TP size 通常设为 8。跨节点的 TP 因为受限于 InfiniBand 或以太网带宽，延迟会拖垮训练效率，除非万不得已，否则不这么做。

四、TP 与相邻技术的边界

为了帮你在架构设计时快速决策，这里厘清 TP 与周边技术的关系：

TP vs DP（9.2 节）：

  • DP 切数据，模型完整；TP 切模型，数据完整（或部分复制）。
  • DP 的通信发生在反向传播末尾（梯度 AllReduce），频率低、数据量大；TP 的通信发生在每层内部，频率高、数据量小。
  • 标准组合：先用 TP 把模型塞进单节点（如 8 卡），再用 DP 把训练扩展到多节点（如 8 节点 × 8 卡 = 64 卡）。此时总并行度 = TP × DP。

TP vs PP（9.4 节 流水线并行）：

  • TP 是横向（层内切矩阵），PP 是纵向（层间切，每张卡负责若干完整 Transformer 层）。
  • TP 的通信是点对点的 AllReduce，PP 的通信是层间激活值点对点传输（P2P Send/Recv）。
  • TP 更适合解决“单层太宽”的问题，PP 更适合解决“网络太深”的问题。两者常结合使用。

TP vs ZeRO（9.5 节）：

  • ZeRO 本质还是 DP，只是把优化器状态、梯度、参数分片存储，计算时通过广播/散射临时凑齐。它不减少单步计算量。
  • TP 真正地把计算本身拆开，单步矩阵乘法规模减半（TP=2 时）。因此 TP 比 ZeRO 更激进地节省显存，但实现更复杂、通信更频繁。

五、显存节省效果：真实数字

以训练一个 70B 参数的 Dense 模型为例，FP16 混合精度：

  • 无并行：单卡需承载约 140 GB 参数 + 420 GB 优化器状态/梯度 + 激活值，远超 80 GB 显存，直接 OOM。
  • TP=8（单节点 8 卡）：
  • 每卡参数：140 GB / 8 ≈ 17.5 GB
  • 每卡优化器状态+梯度：420 GB / 8 ≈ 52.5 GB
  • 激活值：按层均分后显著降低
  • 总计可落入 80 GB 显存。这是 TP 最直接的价值。

注意：TP 节省的是每卡静态显存和每卡激活值峰值，但并不会降低整个集群的总显存消耗（毕竟模型总参数没变，只是分散了）。

六、配置与调优实战要点

  1. TP size 的选取：
  • 首选单节点 GPU 数（4/8）。主流训练框架（Megatron-LM、DeepSpeed、PyTorch ParallelTensor）默认 TP 组绑定在 NVLink 域内。
  • 不要超过单节点卡数，除非你做特殊网络拓扑（如 DGX SuperPOD 的 NVLink Domain 跨机柜），否则通信延迟会击穿 MFU（算力利用率）。
  1. 与 Sequence Parallelism（序列并行）结合：
  • 当序列长度超过 32K 甚至 128K 时，激活值的显存瓶颈会超过参数。此时可在 TP 基础上对 LayerNorm/Dropout 的激活值做序列维度切分（Megatron 的 SP），进一步摊薄显存。
  1. 通信与计算重叠：
  • 优秀的 TP 实现会把 AllReduce 与下一层的计算流水线重叠。框架层面的参数（如 delay_grad_reduce 或 bucket_size）会影响这一重叠效率，需要结合 profiling 调整。
  1. 调试信号：
  • 若发现 TP 开启后 GPU 利用率（Utilization）不低但迭代速度（Iteration Time）反而暴增，通常是通信瓶颈（NVLink 未启用、PCIe fallback、或 NCCL 拓扑识别错误）。
  • 用 nvidia-smi topo -m 检查 GPU 互联拓扑；用 Nsight Systems 查看 NCCL kernel 是否阻塞了 CUDA stream。

七、小结

张量并行（TP）是突破单卡显存天花板的第一道关口。它通过把 Transformer 层内的权重矩阵按列/行切开，让多张 GPU 共同完成一次前向/反向传播，并以高频、小数据量、低延迟的 AllReduce 同步中间结果。

核心 takeaway：

  • TP 切的是层内矩阵，通信方式是 AllReduce，通信量是激活值尺寸，与参数量无关；
  • TP 极度依赖高带宽互联，因此通常约束在单节点（NVLink 域）内部，TP size 一般等于单节点 GPU 数；
  • TP 常与 DP 叠加：TP 负责“装进单节点”，DP 负责“扩展到集群”；
  • Attention 的多头结构天生适合 TP，这是 Transformer 架构在工程可扩展性上的巨大红利。

在 9.4 节中，我们将转向另一种模型切分维度——流水线并行（PP）。如果说 TP 是“把一层拆开给多人做”，那么 PP 就是“把多层流水线串起来每人做几层”。两者结合，再加上 9.5 节的 ZeRO，才构成了今天千亿模型训练的工程基座。