4、张量并行原理：张量切分方式（行并行、列并行）、矩阵乘法中的切分与聚合、通信模式

张量并行（Tensor Parallelism）是大模型分布式训练与推理中最核心的模型并行策略之一。与流水线并行按层切分不同，张量并行是在单个算子（尤其是矩阵乘法）内部对权重张量进行切分，将计算分散到多个设备上，从而解决单卡显存无法容纳超大权重矩阵的问题。理解张量并行的本质，需要掌握三个关键点：张量如何切分（行并行与列并行）、切分后的矩阵乘法如何计算与聚合、以及切分带来的通信模式与开销。本章将从这三个维度展开，并结合实际案例说明工程实现中的要点与陷阱。

张量并行切分与通信模式总览 权重 W [Out, In] 例如 [4096, 4096] 列切分 W₁ [Out, In/2] W₂ [Out, In/2] Y = [X·W₁, X·W₂] 拼接聚合 无需通信 行切分 W₁ [Out/2, In] W₂ [Out/2, In] Y₁ = X·W₁ Y₂ = X·W₂ AllReduce 求和 Y = Y₁ + Y₂ 通信模式对比 • 列并行：前向无需通信，反向需要 AllReduce（梯度聚合） • 行并行：前向需要 AllReduce（部分和聚合），反向需要 AllReduce（梯度同步） • 通信量级：每次 AllReduce 传输数据量 ≈ 激活值大小 × 2（数据并行中为梯度 × 2）

上图展示了张量并行中两种最基本的切分方式及其对应的通信模式。接下来，我们将深入剖析每种切分方式的数学原理、实现细节与工程权衡。

4.1 列并行（Column Parallelism）：按输出维度切分

4.1.1 核心原理

列并行是将权重矩阵 W（形状为 [Out, In]）沿输出维度（列方向）切分为多块。假设有 P 个设备，则第 i 个设备持有 W_i，形状为 [Out, In/P]。输入 X（形状为 [Batch, In]）在每个设备上保持完整副本。每个设备独立计算 Y_i = X · W_i，得到形状为 [Batch, Out] 的部分结果。最后将所有设备的 Y_i 在输出维度上拼接（Concat），得到完整的输出 Y = [Y_1, Y_2, ..., Y_P]。

关键点：列并行在前向传播中不需要任何通信，因为每个设备只计算输出的一部分，拼接操作是纯本地逻辑（在内存中按列拼接即可）。但在反向传播中，需要对输入梯度 dX 进行 AllReduce 聚合，因为 dX = dY · W^T，而 W^T 被切分在多个设备上，每个设备计算得到 dX_i = dY_i · W_i^T，需要将所有 dX_i 求和才能得到完整的 dX。

4.1.2 判断方法与适用场景

列并行适用于以下场景：

  • MLP 的第一层（输入投影层）：输入特征维度较大，输出维度相对较小，列并行可以避免输入在不同设备间复制。
  • Attention 中的 Q/K/V 投影：将 Q、K、V 的权重矩阵按列切分到不同设备，每个设备独立计算注意力头的一部分。
  • 需要减少前向通信延迟的场景：由于前向无通信，列并行在推理场景中具有明显优势。

4.1.3 案例拆解： 的 MLP 层列并行

以 的 MLP 层为例，其结构为 Linear(12288 → 49152) → GELU → Linear(49152 → 12288)。在 8 卡张量并行下：

  1. 第一个 Linear 层：权重 [49152, 12288] 按列切分为 8 块，每块 [49152, 1536]。每个设备持有完整的输入 X（形状 [Batch, 12288]），计算得到 Y_i（形状 [Batch, 49152]）。
  2. GELU 激活：逐元素操作，无需通信。
  3. 第二个 Linear 层：权重 [12288, 49152] 按行切分（见 4.2 节），每个设备持有 [1536, 49152]，接收第一个 Linear 层输出的部分列。

实际测量中，8 卡列并行相比单卡，计算吞吐提升约 7.2 倍（受限于通信开销），显存占用降低为原来的 1/8（仅权重部分）。

常见误区：初学者常误以为列并行前向无通信就代表整体无通信。实际上，反向传播中的 AllReduce 通信量与前向激活值大小成正比。在深层网络中，激活值可能非常大（例如 Batch=64, Seq=2048, Hidden=12288 时，单层激活值约 6GB），此时反向通信开销可能成为瓶颈。

4.2 行并行（Row Parallelism）：按输入维度切分

4.2.1 核心原理

行并行是将权重矩阵 W 沿输入维度（行方向）切分。第 i 个设备持有 W_i，形状为 [Out/P, In]。输入 X 也被切分为 X_i（形状 [Batch, In/P]），每个设备持有输入的一部分。每个设备计算 Y_i = X_i · W_i，得到形状为 [Batch, Out/P] 的部分结果。由于每个设备只计算了输出的一部分（按行切分），需要将所有设备的 Y_i 进行AllReduce 求和，得到完整的输出 Y = Σ Y_i。

关键点：行并行在前向传播中需要一次 AllReduce 通信，通信数据量为输出激活值的大小。在反向传播中，需要对权重梯度 dW 进行 AllReduce（因为 dW = X^T · dY，而 X 被切分），同时输入梯度 dX 的聚合方式取决于具体实现（通常也需要 AllReduce）。

4.2.2 判断方法与适用场景

行并行适用于：

  • MLP 的第二层（输出投影层）：输入维度较大（来自前一层的输出），输出维度较小。
  • Attention 的输出投影层：将多头注意力的结果合并回原始维度。
  • Embedding 层与输出层：当词表维度极大时，可将 Embedding 矩阵按行切分，每个设备负责一部分词向量。

4.2.3 案例拆解：LLaMA-70B 的 Embedding 行并行

LLaMA-70B 的词表大小为 32000，Embedding 维度为 8192。在 4 卡张量并行下：

  1. 切分：Embedding 矩阵 [32000, 8192] 按行切分为 4 块，每块 [8000, 8192]。
  2. 前向：输入 ID 序列，每个设备只处理落在自己词表范围内的 。若 ID 不在本设备范围内，则对应输出为 0 向量。
  3. 聚合：将 4 个设备的输出进行 AllReduce 求和，得到完整的 Embedding 输出。

实测中，4 卡行并行 Embedding 的吞吐为单卡的 3.5 倍（受限于 AllReduce 通信），显存占用降低 75%。但需要注意，当 Batch 较小时，通信开销占比上升，加速比可能降至 2.8 倍。

经验提醒：在实际工程中，行并行和列并行通常成对出现。例如 MLP 层：第一层用列并行，第二层用行并行，这样可以在两层之间省去一次 AllReduce（第一层的输出拼接后直接作为第二层的输入切分，无需额外通信）。这种设计被称为"交错并行"（Interleaved Parallelism），是 Megatron-LM 的核心优化之一。

4.3 矩阵乘法中的切分与聚合策略

4.3.1 切分维度选择的原则

矩阵乘法 Y = X · W 中，X 的形状为 [M, K]，W 的形状为 [K, N]，Y 的形状为 [M, N]。张量并行切分维度选择遵循以下原则：

切分方式 权重切分 输入切分 输出聚合 前向通信 反向通信
列并行 按 N 维切分 完整复制 拼接（Concat） 无 AllReduce（dX）
行并行 按 K 维切分 按 K 维切分 求和（AllReduce） AllReduce（Y） AllReduce（dW）
二维并行 按 K 和 N 同时切分 按 K 切分 先局部拼接再 AllReduce AllReduce（部分和） 两次 AllReduce

4.3.2 聚合操作的工程实现

聚合操作有两种实现方式：

  • 拼接（Concat）：适用于列并行。每个设备输出 [M, N/P]，在内存中按列拼接为 [M, N]。该操作不涉及跨设备通信，但需要额外的内存拷贝。
  • AllReduce 求和：适用于行并行。每个设备输出 [M, N/P]，需要将所有设备的结果逐元素相加。AllReduce 的通信量 = M × N/P × 4 字节（FP32）。
// 伪代码：行并行前向聚合
// 每个设备 i 计算得到 local_output_i [M, N/P]
// 执行 AllReduce 求和
AllReduce(local_output_i, SUM);
// 聚合后每个设备都持有完整的 output [M, N]

4.3.3 案例拆解：Transformer 层中的切分组合

一个标准的 Transformer 层包含两个子层：Attention 和 MLP。在 Megatron-LM 的张量并行实现中：

  1. Attention 层：QKV 投影使用列并行（每个设备计算部分注意力头），输出投影使用行并行（合并所有头的结果）。
  2. MLP 层：第一个 Linear 使用列并行（扩展维度），第二个 Linear 使用行并行（压缩回原维度）。
  3. 层间通信：每个子层结束后，输出通过 AllReduce 同步到所有设备，作为下一子层的输入。

这种设计使得每个子层只需要一次 AllReduce 通信（在子层输出处），将通信次数从 4 次（每个 Linear 一次）减少到 2 次，显著降低通信开销。

4.4 通信模式与开销分析

4.4.1 通信原语与数据量

张量并行主要依赖两种通信原语：

  • AllReduce：对所有设备的数据进行规约（求和/求平均），并将结果广播给所有设备。通信量 = 2 × 数据量 × (P-1)/P（基于 Ring AllReduce 实现）。
  • AllGather：收集所有设备的数据并拼接，每个设备获得完整数据。通信量 = 数据量 × (P-1)/P。

以 8 卡、激活值 1GB 为例：

操作 通信量 耗时（NVLink 300GB/s） 耗时（RDMA 25GB/s）
AllReduce（1GB） ~1.75GB ~5.8ms ~70ms
AllGather（1GB） ~0.875GB ~2.9ms ~35ms

4.4.2 通信与计算重叠

为了隐藏通信延迟，现代框架（如 Megatron-LM、DeepSpeed）采用以下策略：

  • 分块通信：将大张量切分为多个小块，每计算完一块就立即通信，实现通信与计算的重叠。
  • 流水线通信：在反向传播中，先计算梯度再通信，将通信操作放入独立的 CUDA 流中。
  • 梯度压缩：对梯度进行量化或稀疏化，减少通信量（但可能影响收敛精度）。

风险提示：张量并行并非卡数越多越好。当设备数量超过 8-16 张时，通信开销呈指数增长，计算加速比趋于饱和。例如，在 64 卡上使用张量并行，通信占比可能超过 60%，实际加速比仅为单卡的 20 倍左右。此时应考虑结合流水线并行或数据并行。

4.4.3 案例拆解：通信瓶颈定位与优化

背景：某团队在 32 卡 A100（NVLink 全互联）上训练 175B 模型，采用 8 路张量并行 + 4 路流水线并行。训练吞吐仅为理论峰值的 35%。

现象：通过 Nsight Systems 分析，发现每个 Transformer 层的 AllReduce 耗时占层计算时间的 45%。进一步定位发现，AllReduce 的通信量远超预期——激活值大小为 2.5GB，但实际通信量达到 4.4GB。

分析：检查代码发现，框架在列并行反向传播时，对 dX 执行了两次 AllReduce（一次在 Attention 输出投影，一次在 MLP 第一层），而这两次 AllReduce 可以合并为一次。此外，AllReduce 使用了默认的 Ring 算法，在 8 卡场景下效率低于 Tree 算法。

优化措施：

  1. 合并相邻层的 AllReduce 操作，减少通信次数。
  2. 将 AllReduce 算法从 Ring 切换为 Tree（8 卡时 Tree 算法通信量减少 30%）。
  3. 启用通信与计算重叠：将 AllReduce 拆分为 4 个 1.1GB 的块，每计算完 1/4 激活值就启动通信。

结果：优化后，AllReduce 耗时降低 62%，训练吞吐提升至理论峰值的 68%。

复盘经验：张量并行通信优化应优先检查：①是否存在重复通信；②通信算法是否匹配硬件拓扑；③是否实现了通信与计算重叠。这三项通常能带来 2-3 倍的性能提升。

4.5 常见误区与总结复盘

4.5.1 常见误区

  • 误区一：认为列并行比行并行更优。实际上两者各有优劣，列并行前向无通信但反向通信量大，行并行前向有通信但反向通信量小。选择哪种方式取决于具体层的位置和激活值大小。
  • 误区二：忽略通信与计算的平衡。张量并行度越高，单卡计算量越小，但通信量不变（甚至增加）。当通信时间超过计算时间时，增加并行度反而降低性能。
  • 误区三：所有层使用相同的并行策略。不同层的最佳并行度不同。例如，Embedding 层通常使用行并行，而 Attention 层使用列并行+行并行组合。应根据每层的计算/通信比动态调整。
  • 误区四：忽视内存布局。切分后的权重矩阵需要连续内存存储，否则会导致访存效率下降。使用 torch.chunk 或 split 时需注意内存对齐。

4.5.2 总结复盘

张量并行的核心设计决策可以归纳为以下清单：

决策点 选项 推荐场景 注意事项
切分维度 行 / 列 / 二维 列并行用于输入投影，行并行用于输出投影 二维并行适合超大矩阵，但通信复杂度高
并行度 P 2-64 8 卡以内通信开销可控 超过 16 卡需结合流水线并行
通信算法 Ring / Tree / 分块 小规模用 Tree，大规模用 Ring 需根据硬件拓扑选择
内存优化 权重分片 / 激活重计算 显存不足时使用 激活重计算增加 30% 计算量

在实际工程中，建议先通过性能分析工具（如 PyTorch Profiler、Nsight Systems）定位通信热点，再针对性地调整切分策略。张量并行是"计算换通信"的典型权衡，只有深入理解矩阵乘法的数学本质和硬件通信特性，才能设计出高效的并行方案。

← 上一章 📖 返回目录 下一章 →