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

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

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

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

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

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

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

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

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

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

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

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

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

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

输入：GPU 0 的输入是YLY_LYL，GPU 1 的输入是YRY_RYR。
切分矩阵：W2W_2W2被横着切成两半。
- GPU 0 拿上半边W2,UpW_{2,Up}W2,Up。
- GPU 1 拿下半边W2,DownW_{2,Down}W2,Down。
独立计算：
- GPU 0 算ZL=YL⋅W2,UpZ_L = Y_L \cdot W_{2,Up}ZL=YL⋅W2,Up。
- GPU 1 算ZR=YR⋅W2,DownZ_R = Y_R \cdot W_{2,Down}ZR=YR⋅W2,Down。
关键通信 (All-Reduce)：现在每张卡算出了最终结果的一部分。根据矩阵乘法规则，最终结果Z=ZL+ZRZ = Z_L + Z_RZ=ZL+ZR。
- 动作：两张卡进行一次All-Reduce (Sum)通信，把各自的ZLZ_LZL和ZRZ_RZR加起来。
- 结果：通信结束后，GPU 0 和 GPU 1 都得到了完整的、正确的结果ZZZ。

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

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

“头”的分配：
- GPU 0 负责计算第 1-16 个头。
- GPU 1 负责计算第 17-32 个头。
切分 Q, K, V 投影层 (列切分)：
- 计算 Attention 需要把输入XXX投影成 Query, Key, Value。对应的权重WQ,WK,WVW_Q, W_K, W_VWQ,WK,WV也是巨大的矩阵。
- 我们像 FFN 第一层一样，把它们按列切分。GPU 0 只拿属于前 16 个头的权重，GPU 1 拿后 16 个头的。
- 两张卡各自算出自己那 16 个头的Q,K,VQ, K, VQ,K,V。
独立计算 Attention：
- GPU 0 用自己的Q,K,VQ, K, VQ,K,V算它那 16 个头的注意力结果。
- GPU 1 同理算它那 16 个头的结果。
- 关键点：这一步完全独立，不需要通信。
切分输出投影层WOW_OWO(行切分)：
- Attention 的最后一步是把所有头的结果拼接起来，乘一个输出矩阵WOW_OWO。
- 为了配合前面的切分，这个WOW_OWO必须按行切分。
- GPU 0 用它那 16 个头的结果乘以WOW_OWO的上半部分。
- GPU 1 用它那 16 个头的结果乘以WOW_OWO的下半部分。
关键通信 (All-Reduce)：
- 和 FFN 一样，最后需要把两部分结果加起来。
- 动作：进行一次All-Reduce (Sum)通信。
- 结果：两张卡都得到了完整的 MHA 输出。

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

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

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

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

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

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

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

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

数据并行 (DP)：
- 模型位置：每张卡上都有一份完整的模型复制。
- CV 类比：就像你和你的同学（不同的 GPU）每人都有一本一模一样的《新华字典》（完整的 ResNet 模型），你们各自查不同的字（处理不同的数据）。
流水线并行 (PP)：
- 模型位置：每张卡上存储模型的一部分层 (Layers)。
- CV 类比：就像工厂流水线，你负责装轮胎（前几层），他负责装发动机（中间层），最后一个人负责喷漆（最后几层）。大家手里都只有自己负责的那部分图纸。
张量并行 (TP)（你现在问的这个）：
- 模型位置：每一层的权重矩阵都被切分，分散在不同卡上。没有任何一张卡拥有某一层完整的权重。
- CV 类比：资料里那个比喻非常精准——“把一个巨大的卷积核撕成了两半，两张卡各拿一半去算”。
- 具体来说：假设一个线性层的权重矩阵WWW大小是4096×40964096 \times 40964096×4096。如果你用 2 张卡做 TP：
  - GPU 0可能只存了这个矩阵的左半边（4096×20484096 \times 20484096×2048）。
  - GPU 1存了右半边（4096×20484096 \times 20484096×2048）。