TP（Tensor Parallelism）

Tensor Parallelism（TP）在每个 Transformer Layer 内切分权重、激活和计算，由多个 TP Rank 共同完成同一批 token 的前向计算。Dense Transformer 的 Attention 和 MLP 通常都采用“Column Parallel 切分中间特征，Row Parallel 归并输出”的 Megatron-LM 风格。

下文只讨论 Dense Decoder Layer 的前向计算，设 $T=B\times S$ 为 token 数，$H$ 为隐藏维度，$p$ 为 TP Degree。

1. Column Parallel 与 Row Parallel

对线性层 $Y=XW$，Column Parallel 沿权重的输出维度切分：

$$W=[W_0,\dots ,W_{p-1}],Y_r=X{W}_r$$

每个 Rank 只生成一部分输出特征，完整结果在逻辑上为 $Y=\mathrm{Concat}(Y_0,\dots ,Y_{p-1})$。如果后续算子能直接消费同样的分片，不需要立即 All-Gather。

Row Parallel 沿矩阵乘的归约维度 $K$ 切分输入和权重：

$$X=[X_0,\dots ,X_{p-1}],W=\left[\begin{array}{c}{W}_0\\ ⋮\\ W_{p-1}\end{array}\right]$$$$Y=XW=\sum _{r=0}^{p-1}{X}_r{W}_r$$

每个 Rank 计算形状与 $Y$ 相同的局部部分和，再通过 All-Reduce Sum 得到完整输出。

2. Attention

设输入 $X\in {\mathbb{R}}^{T\times H}$，Query Head 数为 $h_q$，KV Head 数为 $h_{kv}$，Q/K Head Dimension 为 $d_k$，Value Head Dimension 为 $d_v$。多数 LLM 中 $d_k=d_v=d_h$。

2.1 Q/K/V/O 的维度关系

四个投影矩阵的形状为：

权重	形状	投影或输入形状
$W_Q$	$H\times (h_q{d}_k)$	$Q:[T,h_q,d_k]$
$W_K$	$H\times (h_{kv}{d}_k)$	$K:[T,h_{kv},d_k]$
$W_V$	$H\times (h_{kv}{d}_v)$	$V:[T,h_{kv},d_v]$
$W_O$	$(h_q{d}_v)\times H$	$Z:[T,h_q{d}_v]$

推导的核心是：Attention 按 Query Head 独立计算。令 GQA 的共享比例为 $g=h_q/h_{kv}$，第 $i$ 个 Query Head 使用第 $\lfloor i/g\rfloor$ 个 KV Head：

$$Z_i=\mathrm{Softmax}\left(\frac{Q_i{K}_{\lfloor i/g\rfloor }^T}{\sqrt{d_k}}\right)V_{\lfloor i/g\rfloor },Z_i\in {\mathbb{R}}^{T\times d_v}$$

KV Head 可以被多个 Query Heads 共享，但每个 Query Head 都产生一个独立的 $d_v$ 维输出。因此完整 Attention 输出为：

$$Z=\mathrm{Concat}(Z_0,\dots ,Z_{h_q-1})\in {\mathbb{R}}^{T\times (h_q{d}_v)}$$

所以 $W_O$ 的输入维度由 $h_q$ 决定，而不是由 $h_{kv}$ 决定。

2.2 Q/K/V 投影按 Head 切分

Q/K/V 投影使用 Column Parallel：

$$W_Q=[W_Q^{(0)},\dots ,W_Q^{(p-1)}],Q_r=X{W}_Q^{(r)}$$

$W_K$ 和 $W_V$ 使用相同切分。当 Head 数可整除 $p$ 时，每个 Rank 持有 $h_q/p$ 个 Query Heads 和 $h_{kv}/p$ 个 KV Heads：

Q_r: [T, h_q  / p, d_k]
K_r: [T, h_kv / p, d_k]
V_r: [T, h_kv / p, d_v]

不同 Query Heads 的 Attention 计算之间没有数据依赖。只要共享同一 KV Head 的 Query Heads 与该 KV Head 位于同一 Rank，各 Rank 就可独立执行 RoPE、Attention Score、Softmax 和 Value 加权。因此 Q/K/V 投影和 Attention 核心计算之间不需要 TP 集合通信；KV Heads 被切分时，KV Cache 也按本地 KV Heads 分片。

2.3 本地 Head Concat 与 Output Projection

Attention 为每个本地 Query Head 产生输出 $Z_{r,j}\in {\mathbb{R}}^{T\times d_v}$。在进入本地 $W_O$ 前，需要沿 Head 维连接这些输出：

$$Z_r=\mathrm{Concat}(Z_{r,0},\dots ,Z_{r,h_q/p-1})\in {\mathbb{R}}^{T\times (h_q{d}_v/p)}$$

这是 Rank 内部的语义操作，不涉及 TP 集合通信。实现通常已经将本地 Heads 存在同一个张量中，因此不需要逐 Head 调用 concat，而是通过布局转换和 reshape 将 Head 维展平到 GEMM 的归约维：

[B, h_q / p, S, d_v]
  → transpose
[B, S, h_q / p, d_v]
  → reshape
[T, h_q d_v / p]

输出布局连续时可能只需要 reshape；否则需要 transpose + contiguous，Fused Attention Kernel 也可直接产生适合 Output Projection 的布局。

局部 Concat 完成后，不需要收集其他 Rank 的 Heads。将 $W_O\in {\mathbb{R}}^{(h_q{d}_v)\times H}$ 按对应 Query Heads 切成 Row Parallel 分片：

$$Z=[Z_0,\dots ,Z_{p-1}],W_O=\left[\begin{array}{c}{W}_O^{(0)}\\ ⋮\\ W_O^{(p-1)}\end{array}\right]$$

则：

$$Z{W}_O=\sum _{r=0}^{p-1}{Z}_r{W}_O^{(r)}$$

每个 Rank 先计算 $P_r=Z_r{W}_O^{(r)}\in {\mathbb{R}}^{T\times H}$，再对 $P_r$ 执行 All-Reduce Sum。这与“Gather 所有 Heads → 乘完整 $W_O$”数学等价；输出 Bias 应在归并后只加一次。

GQA/MQA 的 KV Head 数可能小于 TP Degree。此时框架需要复制 KV Heads，或限制 $p$ 必须整除 $h_{kv}$；具体策略取决于实现。

3. MLP

以 SwiGLU MLP 为例，中间维度为 $F$：

$$G=X{W}_{gate},U=X{W}_{up},Z=\mathrm{SiLU}(G)\odot U,Y=Z{W}_{down}$$

$W_{gate},W_{up}\in {\mathbb{R}}^{H\times F}$ 使用 Column Parallel，且必须沿相同的 $F$ 区间切分：

$$G_r=X{W}_{gate}^{(r)},U_r=X{W}_{up}^{(r)},Z_r=\mathrm{SiLU}(G_r)\odot U_r$$

这样 SwiGLU 的激活和逐元素乘法都可在本地完成，$Z_r\in {\mathbb{R}}^{T\times (F/p)}$ 无需 All-Gather。

$W_{down}\in {\mathbb{R}}^{F\times H}$ 使用 Row Parallel，每个 Rank 持有 $W_{down}^{(r)}\in {\mathbb{R}}^{(F/p)\times H}$：

$$P_r=Z_r{W}_{down}^{(r)},Y=\mathrm{AllReduceSum}(P_r)$$

4. 通信与约束

经典 TP 布局中，每个 Rank 在 Decoder Layer 输入处持有完整 $X$。前向通信边界为：

模块	Column Parallel	本地计算	Row Parallel	归并
Attention	Q/K/V Projection	本地 Heads Attention	Output Projection	1 次 All-Reduce
MLP	Gate/Up Projection	SwiGLU	Down Projection	1 次 All-Reduce

因此一个标准 Decoder Layer 通常有两次 TP All-Reduce。归并后的激活在各 Rank 上一致，Residual Add 和 RMSNorm 可以本地执行。

与 Sequence Parallel 组合时，Row Parallel 后的 All-Reduce 可拆成 Reduce-Scatter，使各 Rank 只保留一段 token 激活；在下一个需要完整 token 输入的 Column Parallel 前再执行 All-Gather。

TP 还有以下实现约束：

• $h_q$、$h_{kv}$ 和 $F$ 通常需要与 TP Degree 兼容。
• 量化 Group Size、Fused QKV 布局和 SwiGLU Gate/Up 排列必须与分片边界对齐。
• TP Degree 越大，每个 Rank 的 GEMM 越小，集合通信占比越高；TP Group 通常优先放在 NVLink/NVSwitch 等高速互联域内。

5. 参考

1. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
2. NVIDIA Megatron Core：Parallelism Strategies Guide
3. Hugging Face Transformers：Llama Tensor Parallel Plan