Tensor Parallel

Tensor Parallelism, TP，是把单层内部的张量计算切分到多张 GPU 上。它不同于 Data Parallel 的“每张 GPU 都有完整模型”，也不同于 Pipeline Parallel 的“不同 GPU 负责不同层”。TP 关注的是同一层里的矩阵乘法、attention heads 或 MLP hidden dimension 如何拆分。

TP 主要解决单层参数和激活太大、单卡无法高效计算的问题，是大规模预训练和 Megatron-style 训练的核心组件。

线性层切分

Transformer 中大量计算来自线性层：

Y = X W

其中 $X \in R^{B \times H}$ ， $W \in R^{H \times K}$ 。Tensor parallel 可以按列或按行切分 $W$ 。

Column Parallel

按输出维度切分：

W = [W_{1}, W_{2}, \dots, W_{p}]

每张 GPU 计算：

Y_{i} = X W_{i}

最后 $Y = [Y_{1}, \dots, Y_{p}]$ 。如果下一层可以消费分片输出，就不必立即 AllGather。

Row Parallel

按输入维度切分：

W = W_{1} W_{2} \dots W_{p}, X = [X_{1}, \dots, X_{p}]

每张 GPU 计算 partial output：

Y_{i} = X_{i} W_{i}

然后通过 AllReduce 求和：

Y = i \sum Y_{i}

Megatron-LM 的 MLP 常用 column-parallel + row-parallel 配对，以减少不必要通信。

Attention 与 MLP 切分

在 Transformer 中：

attention 的 Q/K/V projection 可以按 heads 或 hidden dimension 切分；
attention output projection 可以用 row parallel；
MLP 的 up/gate projection 常用 column parallel；
MLP 的 down projection 常用 row parallel。

这种设计让每个 GPU 只处理部分 heads 或部分 intermediate dimension，降低单卡计算和参数压力。

通信模式

TP 的通信发生在层内部，通常比 data parallel 更频繁。常见通信：

AllReduce：合并 row-parallel partial outputs；
AllGather：收集分片 activation；
ReduceScatter：合并并分发结果；
broadcast / scatter：准备输入分片。

因为 TP 通信频繁，最好放在同节点高速互联 GPU 上，例如 NVLink。跨节点做大 TP 往往通信开销很高。

Sequence Parallel

Sequence Parallelism 常与 tensor parallel 配合，把某些 activation 沿 sequence dimension 切分，减少 activation memory。它通常用于 LayerNorm、Dropout、residual 等不需要完整 hidden dimension 的部分。

直观上：

TP 切 hidden / heads / MLP dimension；
sequence parallel 切 sequence dimension；
二者配合降低单卡 activation 和通信压力。

具体实现依赖框架，Megatron-Core 中 sequence parallel 是重要优化。

与其他并行的关系

TP 解决单层太大的问题；PP 解决层数/模型深度分布问题；DP/FSDP/ZeRO 解决数据并行冗余状态问题。

常见组合：

WorldSize = TP \times PP \times D P

如果一个 1024 GPU 训练使用 $TP = 8, PP = 8$ ，则 $D P = 16$ 。每个并行维度有不同通信组。

适用场景

TP 适合：

单层矩阵乘法太大；
attention heads 很多；
hidden size / MLP intermediate size 很大；
希望提高单 step 计算并行度；
大规模预训练中同节点 GPU 高速互联充足。

TP 不适合过小模型或低带宽跨节点环境，因为通信可能超过计算收益。

常见失败模式

TP degree 过大：每卡矩阵太小，kernel efficiency 下降。
跨节点 TP：频繁通信导致严重瓶颈。
通信未重叠：AllReduce/AllGather 等待时间变长。
与 PP/DP 组网不匹配：GPU 拓扑没有按通信密集度布置。
实现细节错误：分片权重、checkpoint 合并和随机种子处理复杂。

🍊 Latent Atlas 🍉

探索