张量并行切分

本文目录导读：

1. 核心思想
2. 核心切分策略
3. 核心通信模式
4. 常见变种与实现
5. 实际配置建议
6. 总结表

我们来详细拆解一下张量并行（Tensor Parallelism, TP） 中的切分策略，这是一个非常核心且复杂的概念，理解它对于掌握大模型分布式训练至关重要。

核心思想

张量并行的核心思想是：将模型某一层中的单个张量（如权重矩阵）沿着一个或多个维度切分成多个分片，并将这些分片分别放置在不同的GPU上。 这样，计算这个层的前向和反向传播时，每个GPU只需要处理其拥有的那部分数据，从而突破了单卡显存的限制。

与数据并行（Data Parallelism）不同，数据并行是每个GPU都有一份完整的模型副本，只切分数据；而张量并行是切分模型本身。

核心切分策略

最常见的张量并行切分方法,主要针对Transformer架构中的两个核心模块：自注意力层（Self-Attention） 和 多层感知机层（MLP / Feed-Forward Network）。

我们以 Megatron-LM 提出的经典“列并行”与“行并行”作为基础框架来讲解，这是目前最主流和高效的方式。

前馈神经网络层（MLP）的切分

一个标准的MLP层通常包含两个线性变换：Y = GeLU(X * A) 和 Z = Y * B。

• 第一步：列并行（Column-wise Parallelism）

  • 目标张量： 第一个权重矩阵 A，其形状通常为 [d_model, d_ff]（或 [d_model, 4 * d_model]）。
  • 切分方式： 沿着 d_ff（第二个维度）进行切分，假设有 t 个GPU，将 A 切分为 t 个 [d_model, d_ff/t] 的子矩阵 A_1, A_2, ..., A_t。
  • 计算过程：
    1. 每个GPU持有完整的输入 X（形状为 [batch_size, seq_len, d_model]）。
    2. 每个GPU计算自己的部分：Y_i = X * A_i，由于 A_i 的列数是原来的 1/t，计算量也降为 1/t。
    3. 输出 Y_i 的形状为 [batch_size, seq_len, d_ff/t]。
  • 结果： 经过这一步，t 个GPU上的 Y_i 是最终的输出 Y 在d_ff维上的一个分片，要将它们拼成完整的 Y，需要一次AllGather操作。

• 第二步：行并行（Row-wise Parallelism）

  • 目标张量： 第二个权重矩阵 B，其形状通常为 [d_ff, d_model]。
  • 切分方式： 沿着 d_ff（第一个维度）进行切分，将 B 切分为 t 个 [d_ff/t, d_model] 的子矩阵 B_1, B_2, ..., B_t。
  • 计算过程：
    1. 这一步的输入就是上一步输出的分片 Y_i（形状为 [batch_size, seq_len, d_ff/t]）。
    2. 每个GPU计算自己的部分：Z_i = Y_i * B_i，这里 Y_i 和 B_i 的矩阵乘法是合法的，因为 Y_i 的 d_ff/t 维与 B_i 的 d_ff/t 维匹配。
    3. 输出 Z_i 的形状为 [batch_size, seq_len, d_model]。
  • 关键点： t 个GPU上的 Z_i 并不是最终的输出 Z 的分片，而是求和的关系，因为 Z = Z_1 + Z_2 + ... + Z_t。
    • 为什么？因为全局的矩阵乘法 Z = Y * B 可以写成： Z = [Y_1, Y_2, ..., Y_t] * [B_1; B_2; ...; B_t] = Y_1 * B_1 + Y_2 * B_2 + ... + Y_t * B_t （注意这里 B 是沿着行切分，Y 是沿着列切分）
    • 只需要一次AllReduce（求和操作），就可以得到完整的输出 Z。

总结MLP切分：

• 列并行：将A沿列切分，每个GPU得到完整的X和A的一个分片。输出是分片。
• 行并行：将B沿行切分，每个GPU得到上一步的分片输出Y_i和B的一个分片。输出需要全局求和。

自注意力层（Self-Attention）的切分

注意力机制可以看作是多个“头”的并行计算，本质上，它也可以被规约为矩阵乘法：

Q = X * W_Q, K = X * W_K, V = X * W_V

• 切分方式（类似MLP的列并行）：

  • 目标张量： W_Q, W_K, W_V，它们的形状都是 [d_model, d_head * num_heads]。
  • 切分方式： 沿着 d_head * num_heads 维度进行切分，假设有 t 个GPU，将每个权重矩阵切为 t 份，这样，每个GPU上拥有 W_Q_i, W_K_i, W_V_i，它们对应的输出 Q_i, K_i, V_i 只包含了部分注意力头的计算。

• 计算过程：

  1. 每个GPU接收完整的输入 X。
  2. 每个GPU计算自己那部分注意力头的 Q_i, K_i, V_i。
  3. 每个GPU在自己的注意力头上独立执行注意力计算（softmax(Q_i * K_i^T / sqrt(d_k)) * V_i），得到输出 O_i。
  4. t 个GPU上的 O_i 是最终的注意力输出 O 沿着“头”维度的分片。

• 输出合并：

  • 需要将 O_i 通过一个输出投影矩阵 W_O（形状 [num_heads * d_head, d_model]）映射回 d_model 维度。
  • 这可以用行并行的方式处理：将 W_O 沿着 num_heads * d_head 维度（行）切分为 t 份 W_O_i（形状 [num_heads * d_head / t, d_model]）。
  • 每个GPU计算 Z_i = O_i * W_O_i，同样，最终的输出 Z 需要 AllReduce 求和。

总结注意力层切分：

• 注意力头的计算天然适合切分,将 W_Q, W_K, W_V 沿列切分，将 W_O 沿行切分。
• 计算过程类似于MLP的“列并行 -> 行并行”模式。

核心通信模式

张量并行的效率严重依赖于通信开销,主要涉及两种通信原语：

1. AllReduce：在行并行的最后一步，需要将各个GPU上的部分和进行全局求和，得到最终结果，这是一个同步操作，通信量很大（需要传输一个完整的[batch_size, seq_len, hidden_size]张量）。
2. AllGather：在列并行的第一步（MLP的GeLU前），需要将各个GPU上的分片输出收集起来，但这种情况在实际实现中通常会被优化掉（例如Megatron-LM的设计，使得下一个行并行的输入直接是分片，避免了AllGather）。

优化技巧： 将AllReduce与下一个层（例如LayerNorm或下一个Transformer Block）的计算进行重叠，减少通信对训练的阻塞。

常见变种与实现

• Megatron-LM (NVIDIA)： 使用上述的列/行并行策略，是工业界最广泛使用的标准实现。
• 序列并行 (Sequence Parallelism)： 在张量并行的基础上，进一步将输入的序列长度维度（seq_len）也切分到不同的GPU上，用于处理超长序列，同时减少LayerNorm和Dropout等操作对单个GPU显存的压力。
• TPT (Tensor Parallelism Transformer)： 有些工作尝试将注意力机制中的Q*K^T计算也并行化，但通常通信量更大，不如上述方法高效。

实际配置建议

常见组合：

• 3D并行： 数据并行 + 张量并行 + 流水线并行，在16台8卡机器上，可以在每台机器内部使用张量并行（利用NVLink高带宽），在机器之间使用流水线并行（减少跨机通信），最后对多个副本使用数据并行（打散数据）。

总结表

理解张量并行的切分策略,关键在于理解何时输出是分片，何时输出需要求和，以及如何通过精巧的切分设计最小化通信开销，希望这个解释对你有帮助。