张量并行 (Tensor Parallelism, TP) 深度解析

张量并行 (Tensor Parallelism, TP) 深度解析

1. TP 只能用于 Transformer 吗？

答案是：不，但它在 Transformer 上用得最多，也最有效。

咱们从 CV 的角度来类比。TP 的核心思想是“拆分矩阵乘法”。任何包含巨大矩阵乘法 ($Y=X\cdot W$) 的模型，理论上都可以用 TP。比如一个超大的全连接网络 (MLP)，你完全可以把它那一层巨大的权重矩阵 $W$ 拆开存到不同卡上。

为什么大家都在说 TP + Transformer？

这是因为 Transformer 的结构（Multi-Head Attention 和 FFN）简直是为 TP 天造地设的。

Transformer 的每一层都包含两个主要的子模块：

1. Multi-Head Attention (MHA)
2. Feed-Forward Network (FFN)

这两个模块都有一个共同的神奇特性：可以把它们切成两半，让两张卡各自独立跑完大部分计算，只在最后一步进行一次通信就能得到正确结果。

这种特性极大地减少了通信频率，让 TP 变得实用。如果一个模型结构导致每做一步小计算都要通信一次，那通信开销会大到让并行失去意义。

所以，TP 是一种通用的技术，但 Transformer 的结构完美契合了 TP 的优势，成为了它的最佳拍档。

---

2. TP 的详细数据流是怎样的？

你问到了资料中没有展开的细节，这部分是理解 TP 的关键。资料里提到的“切分矩阵”、“算前16个头”、“立刻通信拼起来”，具体是怎么发生的呢？

我们以 2 张显卡 (GPU 0, GPU 1) 为例，分别看看 FFN 和 MHA 的数据流。

场景一：Feed-Forward Network (FFN) 的数据流

FFN 通常是一个两层的 MLP：先放大维度，再缩小回原维度。假设输入是 $X$，两个权重矩阵是 $W_1$ 和 $W_2$。

第一步：列切分 (Column Parallelism) 处理 $W_1$

1. 复制输入：GPU 0 和 GPU 1 拿到完全一样的输入数据 $X$。
2. 切分矩阵：巨大的 $W_1$ 被竖着切成两半。
   • GPU 0 拿左半边 $W_{1,L}$。
   • GPU 1 拿右半边 $W_{1,R}$。
3. 独立计算：
   • GPU 0 算 $Y_L=\text{ReLU}(X\cdot W_{1,L})$。
   • GPU 1 算 $Y_R=\text{ReLU}(X\cdot W_{1,R})$。
   • 关键点：到这里，两张卡各自拿着输出向量的一半，不需要通信。

第二步：行切分 (Row Parallelism) 处理 $W_2$

为了让最后的结果能加起来，第二个矩阵 $W_2$ 必须横着切。

1. 输入：GPU 0 的输入是 $Y_L$，GPU 1 的输入是 $Y_R$。
2. 切分矩阵：$W_2$ 被横着切成两半。
   • GPU 0 拿上半边 $W_{2,Up}$。
   • GPU 1 拿下半边 $W_{2,Down}$。
3. 独立计算：
   • GPU 0 算 $Z_L=Y_L\cdot W_{2,Up}$。
   • GPU 1 算 $Z_R=Y_R\cdot W_{2,Down}$。
4. 关键通信 (All-Reduce)：现在每张卡算出了最终结果的一部分。根据矩阵乘法规则，最终结果 $Z=Z_L+Z_R$。
   • 动作：两张卡进行一次 All-Reduce (Sum) 通信，把各自的 $Z_L$ 和 $Z_R$ 加起来。
   • 结果：通信结束后，GPU 0 和 GPU 1 都得到了完整的、正确的结果 $Z$。

FFN 总结：输入复制 -> 第一层列切分 -> 第二层行切分 -> 最后做一次 All-Reduce。

---

场景二：Multi-Head Attention (MHA) 的数据流

这个更直观，就是资料里说的“CV人的类比…撕成两半”。假设有 32 个头。

1. “头”的分配：

   • GPU 0 负责计算第 1-16 个头。
   • GPU 1 负责计算第 17-32 个头。

2. 切分 Q, K, V 投影层 (列切分)：

   • 计算 Attention 需要把输入 $X$ 投影成 Query, Key, Value。对应的权重 $W_Q,W_K,W_V$ 也是巨大的矩阵。
   • 我们像 FFN 第一层一样，把它们按列切分。GPU 0 只拿属于前 16 个头的权重，GPU 1 拿后 16 个头的。
   • 两张卡各自算出自己那 16 个头的 $Q,K,V$。

3. 独立计算 Attention：

   • GPU 0 用自己的 $Q,K,V$ 算它那 16 个头的注意力结果。
   • GPU 1 同理算它那 16 个头的结果。
   • 关键点：这一步完全独立，不需要通信。

4. 切分输出投影层 $W_O$ (行切分)：

   • Attention 的最后一步是把所有头的结果拼接起来，乘一个输出矩阵 $W_O$。
   • 为了配合前面的切分，这个 $W_O$ 必须按行切分。
   • GPU 0 用它那 16 个头的结果乘以 $W_O$ 的上半部分。
   • GPU 1 用它那 16 个头的结果乘以 $W_O$ 的下半部分。

5. 关键通信 (All-Reduce)：

   • 和 FFN 一样，最后需要把两部分结果加起来。
   • 动作：进行一次 All-Reduce (Sum) 通信。
   • 结果：两张卡都得到了完整的 MHA 输出。

MHA 总结：按“头”把 Q/K/V 权重列切分 -> 独立算 Attention -> 把输出权重 $W_O$ 行切分 -> 最后做一次 All-Reduce。

---

导师总结 (数据流)

现在回看资料里的那句话：“每一层计算结束，必须立刻通信，将两部分结果拼起来（或相加）”，是不是清晰多了？

• 在 Transformer 的每一层里，包含一个 MHA 块和一个 FFN 块。
• TP 的魔法流程是：
  1. 进入 MHA 块，大家分头干活，最后通信一次 (All-Reduce) 得到完整结果。
  2. 拿着这个结果进入 FFN 块，再分头干活，最后再通信一次 (All-Reduce) 得到完整结果。
  3. 这一层就算跑完了，带着完整的结果进入下一层。

这就是为什么 TP 的通信量巨大，因为它在一个 Transformer 层内部就要进行两次高频的 All-Reduce 通信。这也是为什么资料强调 TP “通常仅限于同一台机器内部（通过 NVLink 连接）使用” 的原因。

---

3. 模型是放在哪张卡上边的呢？

这是一个非常关键的问题！理解了这一点，你就真正掌握了张量并行（TP）的精髓。

在张量并行（TP）的模式下，完整的模型并不存在于任何单一的显卡上。

模型的权重参数是被“撕碎”了，分散存储在参与 TP 的所有显卡上的。

为了让你彻底明白，我们还是用你熟悉的 CV 领域和之前讲过的并行方式做对比：

3.1 对比三种方式的模型位置

• 数据并行 (DP)：

  • 模型位置：每张卡上都有一份完整的模型复制。
  • CV 类比：就像你和你的同学（不同的 GPU）每人都有一本一模一样的《新华字典》（完整的 ResNet 模型），你们各自查不同的字（处理不同的数据）。

• 流水线并行 (PP)：

  • 模型位置：每张卡上存储模型的一部分层 (Layers)。
  • CV 类比：就像工厂流水线，你负责装轮胎（前几层），他负责装发动机（中间层），最后一个人负责喷漆（最后几层）。大家手里都只有自己负责的那部分图纸。

• 张量并行 (TP)（你现在问的这个）：

  • 模型位置：每一层的权重矩阵都被切分，分散在不同卡上。没有任何一张卡拥有某一层完整的权重。
  • CV 类比：资料里那个比喻非常精准——“把一个巨大的卷积核撕成了两半，两张卡各拿一半去算”。
  • 具体来说：假设一个线性层的权重矩阵 $W$ 大小是 $4096\times 4096$。如果你用 2 张卡做 TP：
    • GPU 0 可能只存了这个矩阵的左半边（$4096\times 2048$）。
    • GPU 1 存了右半边（$4096\times 2048$）。

3.2 导师总结 (模型位置)

所以，回答你的问题：“模型是放在哪张卡上边的呢？”

答案是：模型被打散了，每个人手里都只有碎片。

• GPU 0：持有第 1 层矩阵的碎片 A，第 2 层矩阵的碎片 C，第 3 层矩阵的碎片 E…
• GPU 1：持有第 1 层矩阵的碎片 B，第 2 层矩阵的碎片 D，第 3 层矩阵的碎片 F…

这就是为什么资料的总结表里提到 TP 的显存节省是“极高 (单层参数被拆解)”。因为每张卡只需要存储原模型 $1/N$ 的参数量（$N$ 是 TP 的并行度）。

但也正因为大家手里都只有碎片，所以每次计算完，必须立刻通信，“将两部分结果拼起来”，才能凑出完整的结果往下传。