1 min read#parallelism#tensor-parallelism#distributed-training#megatron

Tensor Parallelism

Apr 7, 2026

张量并行(Tensor Parallelism, TP) 将单层 Transformer 内部的线性层参数按维度切分到多个 GPU 上，使得一次矩阵乘法由多张卡协作完成。TP 的设计哲学是”为计算而切分”——切分方式严格遵循线性代数规则，保证计算可以直接在分片上进行。

列并行与行并行

TP 的核心是两种互补的切分方式：

列并行(Column Parallel)

将权重矩阵 $W \in \mathbb{R}^{d \times h}$ 沿列方向切分为 $N$ 份： $W = [W_0 | W_1 | \cdots | W_{N-1}]$ 。

每个 rank $i$ 存储 $W_i \in \mathbb{R}^{d \times h/N}$ ，独立计算 $Y_i = XW_i$ 。结果 $[Y_0 | Y_1 | \cdots]$ 是最终输出的列方向分片，无需通信。

行并行(Row Parallel)

将权重矩阵 $W$ 沿行方向切分。输入也对应分片（通常来自前一层的列并行输出）。每个 rank 计算 $Z_i = Y_i W_i$ ，得到的是最终输出的部分和，需要通过 all-reduce 聚合。

交替使用的精妙之处

在 Transformer 的 FFN 中，TP 交替使用列并行和行并行：

第一个线性层（ $W_{up}$ ）做列并行： $X \to [Y_0 | Y_1]$ ，无通信
激活函数（GeLU）各 rank 独立执行
第二个线性层（ $W_{down}$ ）做行并行： $[Y_0 | Y_1] \to Z_0 + Z_1$ ，需要一次 all-reduce

整个 FFN 只需一次 all-reduce。如果只用列并行，每一层都需要 all-gather 还原完整激活值，通信代价更大。

通信特征

通信类型：主要是 all-reduce（聚合部分和）
通信量：每层约 $2S$ （ $S$ 为激活值总大小），与 batch size 和激活专家数无关
通信拓扑：通常限制在单机 NVLink 域内（~900 GB/s），因为 all-reduce 频率高、对延迟敏感
跨机 TP 不常见——跨节点带宽低一个数量级（RDMA 50-100 GB/s），且 TP 通信量恒定，无法通过减少激活来降低

对模型的侵入性

TP 要求重写模型层代码（ColumnParallelLinear、RowParallelLinear），指定每一层如何切分。这使得 TP 的工程侵入性远高于 DP/FSDP（后者只需用包装器把模型包起来）。

TP vs FSDP：一个关键对比

维度	TP	FSDP
核心思想	为计算而切分	为存储而切分
计算方式	先算后聚（直接在分片上计算）	先聚后算（聚合完整参数再计算）
通信内容	激活值的部分和	完整参数
代码侵入性	高（重写模型层）	低（包装器）

两者的最佳实践是组合使用：节点内用 TP 切分大层，节点间用 FSDP2 做数据并行。

在 MoE 模型中的局限

对于 MoE(Mixture of Experts) 模型，TP 会把本就不大的单个专家矩阵进一步切细，产生”瘦矩阵”（thin matrix），严重损害 GPU Tensor Core 的 GEMM 效率。尤其是 DeepSeek 式的细粒度专家设计（256+ experts），TP 的计算效率退化到难以接受的程度。此时 EP 是更优选择。