Jellow 编辑发布,AI 辅助生成并完成技术核对。 计算示例和实测结果会在文中分别说明。
模型刚好能放进一张 GPU,但 Decode 的每 token 延迟没有达标。把 tensor_parallel_size 从 1 改成 4,会接近四倍快吗?另一种情况更直接:模型单卡放不下,究竟该用几路张量并行(Tensor Parallelism,TP),还是沿层切成流水线并行?
答案不能只看 GPU 数量。张量并行把矩阵乘法分给多张卡,同时把原来留在卡内的数据依赖变成集合通信。本文从一次调度步骤中的激活形状 [T,H] 出发,算出标准 Megatron 式稠密 Transformer 的通信账,再给出可测量的加速条件。
核对日期为 2026-09-12。讨论范围是采用一维列并行/行并行的稠密 Decoder 推理;MoE、Sequence Parallel、Context Parallel 和特殊融合实现放在失效边界中处理。
先把“放得下”和“跑得快”拆开
对每个候选 TP 度数 p,先检查容量约束:
[
M{rank}(p)=M{weight}(p)+M{KV}(p)+M{runtime}(p)\le M_{usable}
]
M_usable 是扣除安全余量后的单卡可用显存。权重通常可按 p 分片,但 KV Cache、词表、临时工作区和框架保留内存是否等比例缩小取决于实现,不能直接把总显存除以 p。
容量检查只产生“可运行的候选集合”,不证明更多卡更快。如果模型能在单卡运行,以吞吐为目标时,合理基线不是一组 TP=p,而是相同 p 张卡上的 p 个 TP=1 副本。vLLM 的官方部署指南也把“模型能放入单卡”作为优先采用单卡推理的起点;节点内模型过大时才优先 TP,缺少 NVLink 的节点则建议考虑流水线并行以降低通信压力。vLLM Parallelism and Scaling
如果模型单卡放不下,最小可行 p 是容量基线。此时“成功”可能只是让模型能够运行,不能把它表述成相对单卡的实测加速。
为什么一个 MLP 只在末尾归并
设当前调度步骤有 T 个活跃 token,隐藏维度为 H,SwiGLU 中间维度为 I:
[
X\in\mathbb{R}^{T\times H},\quad
W_g,W_u\in\mathbb{R}^{H\times I},\quad
W_d\in\mathbb{R}^{I\times H}
]
Decode 每条活跃序列通常贡献一个新 token,因此此时 T 近似活跃批大小;Prefill 或 Chunked Prefill 中,T 是本步骤实际调度的 prompt token 总数,不等于请求数。
把前两个权重沿输出维切成 p 份,把下降投影沿输入维切成对应的 p 份:
[
Wg=[W{g,1},\ldots,W_{g,p}],\quad
Wu=[W{u,1},\ldots,W_{u,p}],\quad
Wd=\begin{bmatrix}W{d,1}\ \vdots \ W_{d,p}\end{bmatrix}
]
第 r 张卡计算:
[
Zr=\operatorname{SiLU}(XW{g,r})\odot(XW_{u,r})
\in\mathbb{R}^{T\times I/p}
]
[
P_r=ZrW{d,r}\in\mathbb{R}^{T\times H},\qquad
Y=\sum_{r=1}^{p}P_r
]
激活函数和逐元素乘法都留在本地分片内,只有最后的 [T,H] 部分和需要 All-Reduce。Megatron Core 的官方 API 仍以 ColumnParallelLinear 和 RowParallelLinear 表达这两种切分。Megatron Core 张量并行层文档
自注意力采用相同配对:Q/K/V 投影按头或输出列切分,注意力在各卡本地计算,输出投影按行切分后归并。原始 Megatron-LM 论文由此得到一个标准 Transformer 层前向两次 All-Reduce:注意力一次,MLP 一次。Megatron-LM 论文
这里的“两次”是本文成本模型的关键假设,不是所有推理引擎和模型结构的常数。
从 [T,H] 得到通信账
设激活通信数据类型每元素占 e_a 字节。一次 All-Reduce 的逻辑载荷为:
[
S=THe_a\quad\text{bytes}
]
对 p 个 rank,NVIDIA nccl-tests 给出的 All-Reduce 总线带宽修正因子是 2(p-1)/p。按每个 rank 的出口带宽折算,一次集合通信的数据项为:
[
V_{AR}=\frac{2(p-1)}{p}S
]
一个有 L 层的步骤包含 2L 次这样的归并,所以总线等价通信量为:
[
V_{step}=4L\frac{p-1}{p}THe_a
]
NVIDIA 对该因子的推导同时提醒:小消息应直接观察操作时间;消息足够大之后,时间才近似为固定开销加字节数除以带宽。因此可用下式做初筛:
[
\tau_{AR}(S,p)\approx \alpha_p+
\frac{2(p-1)S}{pB_p}
]
α_p 是一次集合通信的固定延迟,B_p 是指定 rank 布局下的有效总线带宽。NCCL 可能选择树、环、NVLS 或其他路径,这个式子不是替代实测的硬件承诺;生产判断应直接代入相同拓扑、相同载荷下测得的 τ_AR。
看一个纯计算示例。假定 H=8192、L=80、T=32,激活以 BF16 通信,p=4:
- 每次载荷
S=32×8192×2=524288字节,即 512 KiB; - 每步有
2L=160次集合通信; - 每个 rank 的总线等价通信量为
4×80×3/4×512 KiB=120 MiB。
120 MiB 看起来不大,但它被拆成 160 次有依赖关系的调用,固定延迟可能比峰值带宽更重要。这也是小 Decode 批量经常无法按卡数线性加速的原因。
加速成立需要满足什么不等式
令 GQA 的 KV 投影总宽度为 H_KV。忽略 bias 后,注意力投影 Q、K、V、O 的主要 FLOPs 为:
[
F{attn}=4T(H^2+HH{KV})
]
SwiGLU 三个线性层的 FLOPs 为:
[
F_{mlp}=6THI
]
因此每层主要投影计算量近似为:
[
F{proj}=4T(H^2+HH{KV})+6THI
]
这些式子按一次乘加计 2 FLOPs。它们没有包括注意力分数、归一化、采样和调度开销;长上下文会增加注意力计算,从而改变计算与通信的比例。
先作一个乐观假设:分片后的 GEMM 效率不下降,投影计算可按 p 等分。则:
[
tp\approx \frac{t{proj,1}}{p}+2L\tau_{AR}(THea,p)+t{other,p}
]
TP 比单卡快的必要条件是:
[
\left(1-\frac{1}{p}\right)t_{proj,1}
2L\tau{AR}+(t{other,p}-t_{other,1})
]
左侧是理想情况下省下的计算时间,右侧是新增通信及其他开销。代价并没有消失,而是从单卡 GEMM 转移到了集合通信、更小的本地 GEMM、同步等待和多卡资源占用。
还能得到一个活跃 token 数阈值。若单卡全部投影时间写成 t_proj,1=aT,并暂时忽略 t_other 的差值,则:
[
T>
\frac{2L\alpha_p}
{a(1-1/p)-4L\frac{p-1}{p}\frac{He_a}{B_p}}
]
分母必须为正;否则仅带宽项就已经吃掉理想计算收益,在这个简化模型中不存在能获益的 T。a 不应从 GPU 标称峰值推测,最好由 TP=1 的应用 trace 按 T 分桶拟合。
用纯 Python 核对分片恒等式和阈值直觉
下面程序只用 Python 标准语法。前半段以 ReLU 代替 SiLU,使整数结果可以精确比较;任何逐元素激活都不改变分片恒等式。后半段枚举一个明确标注的成本模型,结果不是 GPU 实测。
def mm(a, b):
return [[sum(x*y for x, y in zip(row, col))
for col in zip(*b)] for row in a]
def elem(a, b, fn):
return [[fn(x, y) for x, y in zip(ar, br)]
for ar, br in zip(a, b)]
def add(a, b):
return elem(a, b, lambda x, y: x + y)
def col_shard(a, rank, p):
n = len(a[0])
assert n % p == 0
q = n // p
return [row[rank*q:(rank+1)*q] for row in a]
def row_shard(a, rank, p):
n = len(a)
assert n % p == 0
q = n // p
return a[rank*q:(rank+1)*q]
X = [[1, -1, 2, 0], [0, 2, -1, 1]] # [T=2, H=4]
Wg = [[(r + 2*c) % 5 - 2 for c in range(6)] for r in range(4)]
Wu = [[(2*r + c) % 7 - 3 for c in range(6)] for r in range(4)]
Wd = [[(r - c) % 5 - 2 for c in range(4)] for r in range(6)]
gate = mm(X, Wg)
up = mm(X, Wu)
z = elem(gate, up, lambda g, u: max(0, g) * u)
dense = mm(z, Wd)
p = 2
partial = []
for rank in range(p):
g = mm(X, col_shard(Wg, rank, p))
u = mm(X, col_shard(Wu, rank, p))
z_rank = elem(g, u, lambda x, y: max(0, x) * y)
partial.append(mm(z_rank, row_shard(Wd, rank, p)))
reduced = add(partial[0], partial[1])
assert reduced == dense
print("algebra: OK", dense)
# 以下全是示例假设,不是硬件实测。
H, I, H_KV, L = 8192, 28672, 1024, 80
ACT_BYTES, P_EFF = 2, 120e12 # BF16;单卡投影有效 120 TFLOP/s
ALPHA, BUS_BW = 8e-6, 150e9 # 每次 8 us;总线带宽 150 GB/s
payload = 32*H*ACT_BYTES
traffic = 4*L*(4-1)*payload//4
assert payload == 512*1024
assert traffic == 120*1024**2
print("TP=4,T=32:", payload//1024, "KiB/call,",
traffic//1024**2, "MiB/step/rank")
def estimate(T, p):
f_layer = 4*T*(H*H + H*H_KV) + 6*T*H*I
t1 = L*f_layer/P_EFF
if p == 1:
return t1, 1.0
payload = T*H*ACT_BYTES
t_ar = ALPHA + payload*2*(p-1)/p/BUS_BW
tp = t1/p + 2*L*t_ar
return tp, t1/tp
print("T " + " ".join(f"TP={p}" for p in (2, 4, 8)))
for T in (1, 4, 8, 16, 32):
print(f"{T:<2} " + " ".join(
f"{estimate(T, p)[1]:.2f}x" for p in (2, 4, 8)))
执行输出:
algebra: OK [[-14, -20, 4, 38], [86, 40, -6, -52]]
TP=4,T=32: 512 KiB/call, 120 MiB/step/rank
T TP=2 TP=4 TP=8
1 0.61x 0.72x 0.79x
4 1.26x 1.81x 2.31x
8 1.53x 2.42x 3.42x
16 1.71x 2.91x 4.51x
32 1.82x 3.25x 5.35x
这个枚举故意给分片 GEMM 理想的 1/p 加速,T=1 仍然变慢。真实小矩阵还可能因每卡的 I/p 过窄而降低计算效率,所以模型通常偏乐观。改变 ALPHA、BUS_BW 和 P_EFF,可以先判断哪些 TP 度数值得进入 GPU 压测,而不能据此宣布性能结果。
公平基线不是只有 TP=1
| 方案 | 回答的问题 | 必须观察的代价 |
|---|---|---|
| 单个 TP=1 副本 | 单请求或单批的本地基线 | 是否满足容量与延迟 SLO |
p 个 TP=1 副本 |
相同 GPU 数下的吞吐基线 | 路由、各副本批量与负载均衡 |
一个 TP=p 副本 |
TP 是否降低单步延迟或解决显存不足 | All-Reduce、GEMM 变窄、同步等待 |
PP=p 或节点内 TP+跨节点 PP |
弱互连、切分不整齐或跨节点时的替代方案 | 流水线气泡、阶段不均衡、激活点对点传输 |
线上比较应报告卡数归一化 goodput:
[
G_{card}=\frac{\text{满足延迟 SLO 的输出 token 数}}
{\text{时间}\times\text{GPU 数}}
]
TP 可能让一个请求更快,却让每卡 goodput 下降;如果模型原本能单卡运行,这正是多副本基线不可省略的原因。
权重量化也不会自动等比例降低这里的通信量。公式中的 e_a 是归并激活的数据类型,不是权重位宽。4 bit 权重配 BF16 激活时,权重读取减少了,而 [T,H] 的 All-Reduce 仍可能按每元素 2 字节进行;必须从 trace 核对实际通信 dtype。
可复现的压测设计
先固定模型、权重与 KV 精度、框架提交版本、CUDA/NCCL 版本、GPU 型号与 rank 布局。对 TP=1、2、4、8 中满足容量和可整除条件的候选执行两阶段测试。
第一阶段隔离组件。记录真实 T 分布,并在目标拓扑上扫描覆盖 THe_a 的消息大小。例如单节点 4 卡可使用:
./build/all_reduce_perf -b 16K -e 4M -f 2 -g 4 \
-d half -n 200 -w 50 -I 1 -K 20
nccl-tests 官方仓库说明了消息范围、预热、逐迭代时延和正确性检查参数;跨节点测试需改用与部署相同的 MPI rank 布局。NVIDIA nccl-tests
同时在完整引擎 trace 中记录:
- 每步活跃 token 数
T,并区分 Prefill 与 Decode; - 每次 collective 的载荷、dtype、耗时和每层调用数;
- collective 位于关键路径上的时间,而不只是所有 kernel 时间之和;
- 分片 GEMM 的形状、耗时、SM 利用率与显存带宽;
- TPOT/ITL、TTFT、端到端延迟的 p50、p95、p99;
- 输出 token/s、满足 SLO 的 goodput,以及每卡归一化结果。
第二阶段回放相同到达过程与输入、输出长度分布。先保持调度参数一致以隔离 TP,再允许每个候选独立调优批量上限,比较各自能达到的最佳 SLO—goodput 前沿。只报离线大批吞吐会掩盖小 T 时的集合通信固定开销。
什么时候公式会失效
以下任一条件成立,都应从 trace 重新建立通信账:
p不能整除注意力头数、KV 头数或 MLP 中间维度,框架发生 padding、复制 KV 头或拒绝加载;- GQA/MQA 的 KV Cache 分片策略与假设不同,尤其
p大于 KV 头数; - 开启 Sequence Parallel 后,部分 All-Reduce 被 All-Gather 与 Reduce-Scatter 替代;
- MoE 引入 token dispatch 和 All-to-All,负载不均衡也进入关键路径;
- 引擎融合了归并、残差或归一化,或者通过自定义 collective 改变调用次数;
- 跨节点链路、并发通信或 rank 放置使
α_p、B_p与微基准不同; - 长上下文的注意力计算占主导,此时只计算投影 FLOPs 会低估可并行计算;
- 本地 GEMM 过窄,实际计算时间明显高于
t_proj,1/p。
最直接的校验是:若 trace 中每步 collective 数不接近 2L,或载荷不是约 THe_a,就停止套用本文的简化公式。
排障要按因果链走
- 先核对正确性与形状。 确认实际
p、rank 数、dtype、头数和中间维度切分;检查是否有 padding、KV 复制或意外降级。 - 再核对物理路径。 运行
NCCL_RUN_DIAGNOSTICS=1,查看nvidia-smi topo -m和 P2P 状态。NCCL 2.31 的官方诊断可在通信器初始化时验证 GPU 间实际路径。NCCL Diagnostics - 用相同载荷测 collective。 如果
nccl-tests本身就慢,先处理 NVLink、PCIe、NIC、NUMA 或 GPUDirect RDMA;不要先怪模型。 - 比较微基准与应用 trace。 微基准正常而应用中的 All-Reduce 慢,检查 rank 放置、并发 stream、CPU 亲和性、框架同步及其他通信争用。
- 通信正常再看 GEMM。 若 collective 关键路径占比不高但加速仍差,检查
I/p、每卡头数和实际 kernel 效率,修正理想1/p假设。 - 最后回到调度与排队。 离线步骤变快而线上 p99 或 goodput 没改善,通常要检查
T分布、排队、批量形成速度和副本负载,而不是继续增加 TP。
NCCL 的官方性能排障文档也建议先区分硬件路径、NCCL 和应用问题,并提醒不要因单一微基准收益就把环境变量强制带入生产。NCCL 2.31.2 Performance and Tuning
把卡数选择落成一条规则
先选出能放下模型且形状可分的 p。对每个候选,用真实 T 分桶测量 τ_AR(THe_a,p),再检查:
[
\text{实测计算节省} > 2L\tau_{AR}+\text{其他新增开销}
]
模型能单卡运行时,只有 TP 在目标延迟下的每卡 goodput 优于多副本基线,或多副本无法满足单请求延迟 SLO,才值得占用多张卡。模型不能单卡运行时,默认选择满足容量和 SLO 的最小 TP 度数;跨越弱链路之前,先比较节点内 TP 加跨节点 PP。
卡数不是越多越好。真正的分界由活跃 token 数、每层 [T,H] 的归并延迟、本地 GEMM 退化和服务 SLO 一起决定。