张量并行 (TP)
列/行权重怎么切、GQA 下 KV 头怎么约束 TP 度、Decode 小消息怎么算、TP 度怎么选
核心要点：
- TP 的列并行 / 行并行交替切分原理
- GQA/MQA 下 TP 度必须整除 KV head 数的约束
- AllReduce 消息大小公式：Prefill 大消息 vs Decode 小消息
- Decode 阶段 alpha 主导与算法切换
- TP 度选择的工程决策框架（NVLink 域 / KV 头整除 / 显存 / 利用率）
张量并行 (Tensor Parallelism, TP) 把单层权重矩阵切到多设备，用 AllReduce 合并部分积，由 Megatron-LM 提出[1]，是推理场景的主力并行策略，通常部署在同节点高带宽互联 (NVLink / C2C) 上。
TP 怎么切权重才能最小化通信？
TP 用"列并行 → 行并行"交替，让每个 Transformer 层只产生 2 次 AllReduce，而不是每个 Linear 都产生一次。
- 列并行 Linear (W1)：各设备持有 W1 的列分片，独立计算，输出无需合并
- 行并行 Linear (W2)：各设备持有 W2 的行分片，独立计算部分积，末尾 AllReduce 合并
MLP 块结构：
Input X
→ [Column-parallel Linear W1] 各设备独立计算
→ GeLU
→ [Row-parallel Linear W2] 各设备计算部分积
→ AllReduce 合并部分积
→ Output
Attention 块结构相同，区别是 QKV projection 按头 (head) 切列，Output projection 切行：
Input X
→ [Column-parallel QKV projection] Q/K/V 各按头切分
→ Scaled Dot-Product Attention 各设备独立计算分配到的头
→ [Row-parallel Output projection] 按行切分计算部分积
→ AllReduce 合并部分积
→ Output
每个 Transformer 层 = 2 次 AllReduce （MLP 出口 + Attention 出口）。
GQA/MQA 下 TP 怎么切 Attention？
GQA 让 KV head 数远小于 Q head 数，TP 度必须整除 KV head 数，否则无法均匀切分 KV 权重。
标准 MHA 中 Q/K/V head 数相同，TP 度整除总 head 数即可。GQA (Grouped Query Attention) 让多个 Q head 共享同一组 KV head，KV head 成为更紧的约束：
$$\begin{equation} t \mid n_{\text{kv}} \quad (\text{TP 度 } t \text{ 必须整除 KV head 数 } n_{\text{kv}}) \label{eq:par-tp-gqa-constraint} \end{equation}$$当 $t > n_{\text{kv}}$ 时，KV head 必须被复制到多个 rank（每个 rank 持有完整 KV head 的副本），不再节省 KV 相关显存，且引入冗余计算。主流推理框架（vLLM / TensorRT-LLM）均要求 $t \le n_{\text{kv}}$，超过时报错或自动退化为复制模式[2]。
典型模型的 GQA 配置与可选 TP 度：
| 模型 | Q heads ($n_q$) | KV heads ($n_{\text{kv}}$) | GQA 比 | 可选 TP 度 |
|---|---|---|---|---|
| LLaMA-3 8B | 32 | 8 | 4:1 | 1, 2, 4, 8 |
| LLaMA-3 70B | 64 | 8 | 8:1 | 1, 2, 4, 8 |
| LLaMA-3 405B | 128 | 8 | 16:1 | 1, 2, 4, 8 |
| Qwen-2.5 72B | 64 | 8 | 8:1 | 1, 2, 4, 8 |
@tbl-tp-gqa 主流 GQA 模型的 KV head 配置与 TP 度约束
GQA 不改变 AllReduce 消息大小：AllReduce 发生在 row-parallel output projection 的输出端，其形状为 $[b, s, h]$（完整 hidden dimension），与 KV head 数无关。GQA 影响的是可选 TP 度的集合，不影响选定 TP 度后的通信量。
MLA（Multi-Latent Attention）的特殊情况：DeepSeek-V3 用 MLA 将 KV 压缩到低秩 latent（$d_c = 512$），不再有传统 KV head 的概念。MLA 的 TP 切分按 Q head 数（128）划分，KV latent 矩阵按列分片，$t \mid 128$ 即可，没有 GQA 那样的 KV head 数约束[3]。
AllReduce 在何时触发，频率是多少？
| 阶段 | 每层 AllReduce 次数 | 说明 |
|---|---|---|
| 推理前向 | 2 | Attention 出口 + MLP 出口 |
| 训练前向 | 2 | 同上 |
| 训练反向 | 2 | 梯度也需要 AllReduce 合并 |
| 训练合计 | 4 | 前向 + 反向 |
@tbl-tp-01 TP AllReduce 触发时机
单次 AllReduce 通信量有多大？
单次 AllReduce 消息大小 = 该层激活值张量大小：
$$\begin{equation} M_{\text{TP}} = b \cdot s \cdot h \cdot \text{dtype\_size} \label{eq:par-tp-comm-volume} \end{equation}$$- $b$: micro-batch size
- $s$：序列长度
- $h$: hidden dimension
- $\text{dtype\_size}$: BF16 = 2, FP32 = 4
典型推理配置 ($b=1$, $s=4096$, BF16) 的单次 AllReduce 消息：
| 模型 | $h$ | 单次消息 |
|---|---|---|
| LLaMA-70B | 8192 | 67.1 MB |
| DeepSeek-V3 | 7168 | 58.7 MB |
| Qwen-72B | 8192 | 67.1 MB |
@tbl-tp-02 典型大模型 TP 单次 AllReduce 消息
这些都属于大消息 (远超 $\alpha$-$\beta$ 模型的带宽-延迟交叉点)，应使用 Ring AllReduce 算法。
Decode 阶段 TP 消息退化到什么程度？
Decode 时序列长度 $s=1$，AllReduce 消息从 MB 级退化到 KB 级，alpha 项（启动延迟）主导，Ring AllReduce 不再是最优算法。
Prefill 处理整个 prompt（$s$ 可达数千到百万），消息大、beta 主导。Decode 每步只生成 1 个 token（$s=1$），代入 ：
$$\begin{equation} M_{\text{TP-decode}} = b \cdot 1 \cdot h \cdot \text{dtype\_size} = b \cdot h \cdot \text{dtype\_size} \label{eq:par-tp-decode-volume} \end{equation}$$典型 decode 场景的消息量级：
| 模型 | $h$ | 单请求 ($b=1$) | 批推理 ($b=32$) | 批推理 ($b=128$) |
|---|---|---|---|---|
| LLaMA-3 70B | 8192 | 16 KB | 512 KB | 2 MB |
| DeepSeek-V3 | 7168 | 14 KB | 448 KB | 1.75 MB |
| Qwen-2.5 72B | 8192 | 16 KB | 512 KB | 2 MB |
@tbl-tp-decode Decode 阶段 TP AllReduce 消息量级（BF16）
单请求 decode 的 16 KB 消息完全处于 alpha 主导区域。代入 Ring AllReduce 公式 ，以 H100 NVLink（$\beta \approx 450 \text{ GB/s 单向}$, $\alpha \approx 3 \mu s$）、$N=8$ 为例：
- 带宽项：$\frac{2 \times 7}{8} \times \frac{16\text{ KB}}{450\text{ GB/s}} \approx 0.06 \mu s$
- 延迟项：$2 \times 7 \times 3 \mu s = 42 \mu s$
- 延迟项是带宽项的 700 倍，通信时间几乎全部由启动开销决定
小消息下的算法切换
Ring AllReduce 的 $O(N)$ 延迟步在小消息下是纯损失。NCCL 在消息小于约 256 KB 时自动切换到 Tree AllReduce：
| 算法 | 延迟步数 | 带宽利用率 | 适用区间 |
|---|---|---|---|
| Ring AllReduce | $2(N-1)$ | $(N-1)/N$（渐近最优） | 大消息（> 256 KB） |
| Tree AllReduce | $2 \log_2 N$ | $1/\log_2 N$（较低） | 小消息（< 256 KB） |
| NVLink Multicast (NVLS) | $\approx 2$ | 接近线速 | 小消息 + NVSwitch V3+ |
@tbl-tp-algo Decode 小消息下的算法选择
- Tree AllReduce：$N=8$ 时延迟步从 14 降到 6（$2\log_2 8$），小消息下快 2× 以上
- NVLink Multicast (NVLS)：NVSwitch V3（H100/H200）支持 in-switch multicast + reduction，AllReduce 退化为近似 2 步（scatter + gather via switch），延迟接近单跳 NVLink 往返[4]
- 环境变量：
NCCL_ALGO=Tree可强制切换；NCCL_NVLS_ENABLE=1启用 NVLink multicast 路径
Decode 批大小对通信机制的影响
批推理（continuous batching）下 $b$ 动态变化。当 $b \ge 128$ 时消息回到 MB 级，Ring 重新占优；$b \le 32$ 时仍处于 alpha 主导区。推理框架的 TP 通信路径必须支持动态消息大小，不能为单一算法调优。
Speculative decoding 每步验证 $k$ 个候选（消息变为 $(1+k) \times h \times \text{dtype}$），部分缓解 alpha 敏感性但仍远小于 prefill 量级（详见 推理部署模式）。
为什么 TP 通信难以隐藏延迟？
TP AllReduce 在关键路径上：每层的计算必须等 AllReduce 完成才能进入下一层，推理前向无法与同层计算重叠。
TP 通信特征：
| 特征 | 值 |
|---|---|
| 通信原语 | AllReduce |
| 消息大小 | 10 ~ 100 MB （大消息） |
| 通信组大小 | 通常 2 ~ 8 （同 board / node 内） |
| 频率 | 推理每层 2 次，训练每层 4 次 |
| 延迟敏感性 | 高 （关键路径上，推理无 overlap） |
| 推荐算法 | Ring AllReduce （大消息带宽最优） |
| 推荐拓扑 | NVLink / C2C 等高带宽互联 |
@tbl-tp-03 TP 通信特征汇总
TP 度增大为什么开销上升？
Ring AllReduce 的延迟随通信组规模线性上升：
$$\begin{equation} T_{\text{Ring-AR}} = 2(N-1)\alpha + \frac{2(N-1)}{N} \cdot \frac{M}{\beta} \label{eq:par-tp-ring-allreduce} \end{equation}$$- 延迟项 $2(N-1)\alpha$ 线性增长：$N=2$ 时 $2\alpha$, $N=8$ 时 $14\alpha$
- 带宽项 $\frac{2(N-1)}{N}$ 趋向 2 （渐近最优），带宽开销增长放缓
大消息场景的真实瓶颈：带宽项主导，单次 AllReduce 时间约 $\frac{2M}{\beta}$，与 $N$ 关系不大。但更高的 TP 度意味着每设备权重分片更小，单设备算力利用率下降，通信/计算比上升。
推理与训练受 TP 影响的差异？
| 维度 | 推理 | 训练 |
|---|---|---|
| AllReduce 次数/层 | 2 （仅前向） | 4 （前向 + 反向） |
| batch size | 较小 | 较大 |
| 消息大小 | 中等 ~ 大 | 大 |
| 实时性 | 高 （影响 TTFT / TPOT） | 较低 |
| Overlap 机会 | 少 （前向无 overlap） | 有 （反向可与下一层梯度 overlap） |
@tbl-tp-04 推理 vs 训练的 TP 影响差异
TP 度怎么选？
TP 度不是越大越好，是四个约束的交集：NVLink 域大小、KV head 整除性、单卡显存、decode 算力利用率。
约束 1：不超过 NVLink 域
TP 通信在关键路径上，跨节点走 InfiniBand/RoCE 时延迟上升 5-10×（NVLink 单跳 ~1-3 μs vs 跨节点 ~5-15 μs）。因此 TP 度应 ≤ 单节点 NVLink 域大小：
| 平台 | NVLink 域 | 最大实用 TP |
|---|---|---|
| DGX A100 | 8 GPU | 8 |
| DGX H100 | 8 GPU | 8 |
| GB200 NVL72 | 72 GPU | 72（但通常不需要） |
@tbl-tp-nvlink NVLink 域与最大实用 TP 度
跨节点 TP 仅在单节点装不下模型时使用（如 405B BF16 需 ~810 GB，单个 8×80GB 节点不够），此时优先考虑量化降精度或 PP 切分。
约束 2：整除 KV head 数
由 约束，$t$ 必须整除 $n_{\text{kv}}$。主流 GQA 模型 $n_{\text{kv}} = 8$，所以可选 TP ∈ {1, 2, 4, 8}——恰好对齐 NVLink 域的 2 的幂次划分。MLA 模型（DeepSeek-V3）无此约束。
约束 3：模型参数必须放得下
$$\begin{equation} \frac{|\theta| \cdot \text{dtype\_size}}{t} + M_{\text{KV}} + M_{\text{overhead}} \le M_{\text{GPU}} \label{eq:par-tp-mem-constraint} \end{equation}$$- $|\theta| \cdot \text{dtype\_size}$：模型权重总字节（70B × 2B = 140 GB for BF16）
- $M_{\text{KV}}$：KV cache 显存（与序列长度和 batch 成正比）
- $M_{\text{overhead}}$：激活、CUDA context 等开销
约束 4：decode 算力利用率地板
高 TP 度让每张卡的 GEMM 变小，decode 单请求时矩阵形状退化为 $[1, h/t] \times [h/t, h]$，GPU SM 填不满。经验门槛：单卡 GEMM 的 M 维（$b \times s$）应 ≥ 64 才有合理利用率；decode 单请求 M=1 时，TP=8 的每卡利用率通常 < 5%。
批推理通过增大 $b$ 缓解此问题：$b=64$ 时 M=64，利用率回到合理区间。
业界典型部署配置
| 模型 | 精度 | GPU | 推荐 TP | 理由 |
|---|---|---|---|---|
| LLaMA-3 70B | BF16 | 4×A100-80G | 4 | 内存恰好放下，TP=4 利用率优于 TP=8 |
| LLaMA-3 70B | FP8 | 2×H100-80G | 2 | FP8 减半显存，TP=2 即可 |
| LLaMA-3 405B | BF16 | 8×H100-80G | 8 | 810 GB 需 8 卡，刚好 NVLink 域 |
| LLaMA-3 405B | FP8 | 8×H100-80G | 8 | FP8 减半后仍需 8 卡放 KV cache |
| DeepSeek-V3 | BF16 | 8×H100 + EP | 8 | MoE 用 EP 跨节点，Dense 层 TP=8 节点内 |
@tbl-tp-deploy 业界推理部署的 TP 度选择
TP 度选择决策流程
- 从显存约束出发：计算最小 TP 度使模型 + KV cache 放得下
- 对齐 KV head 约束：向上取到 $n_{\text{kv}}$ 的因子
- 检查 NVLink 域：超过域大小则考虑量化或 PP
- 评估利用率：目标 batch 下单卡 GEMM 是否够大
- 最小化原则：满足所有约束的最小 TP 度是最优选择——TP 越大通信开销越大、利用率越低
Takeaway
| 知识点 | 核心结论 |
|---|---|
| TP 切分原理 | 列并行 → 行并行交替，每层 2 次 AllReduce |
| GQA 约束 | TP 度必须整除 KV head 数；主流模型 $n_{\text{kv}}=8$ → TP ∈ {1,2,4,8} |
| MLA 无 KV head 约束 | DeepSeek-V3 按 Q head 数切分，$t \mid 128$ 即可 |
| Prefill 通信量 | $b \times s \times h \times \text{dtype}$，典型 10~100 MB，Ring AllReduce 最优 |
| Decode 通信量 | $b \times h \times \text{dtype}$，单请求 ~16 KB，alpha 主导 |
| Decode 算法切换 | 小消息用 Tree AllReduce（$2\log_2 N$ 步）或 NVLS（~2 步） |
| TP 度选择 | 取满足显存 + KV head 整除 + NVLink 域的最小值 |
| 跨节点 TP | 延迟上升 5-10×，仅在单节点装不下时使用 |
| 延迟位置 | 关键路径上，推理前向无法 overlap |
@tbl-tp-05 TP 核心知识点
参考资料
- Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, arXiv:1909.08053, 2019. https://arxiv.org/abs/1909.08053
- vLLM Project, Tensor Parallelism Implementation — GQA head partitioning logic. https://github.com/vllm-project/vllm
- DeepSeek-AI, DeepSeek-V3 Technical Report, arXiv:2412.19437, 2024. https://arxiv.org/abs/2412.19437
- NVIDIA, NCCL Documentation — NVLink SHARP (NVLS) for intra-node collectives. https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/nvls.html