<aside>
本文数学推导由 GPT6-Astra 完成,文本部分由人类润色。
</aside>
<aside>
导读:在大模型推理部署中,工程师需要在给定数量的 GPU 上组合配置五种并行策略:DP(数据并行)、TP(张量并行)、EP(专家并行)、CP(上下文并行)和 PP(流水线并行)。
许多经验总结试图用一张简单的正负号表来概括各并行维度的作用,例如“增大 TP 降低显存”、“增大 CP 降低 KV 缓存”或“MoE 模型必须使用大 EP”。然而在实际生产环境中,各并行维度共享同一批物理 GPU,调大某一个并行度必然要求按比例缩小另一个并行度。脱离了“固定总资源”这一约束前提,单纯看某一项指标的变化往往会得出错误的结论。
本文旨在将大模型推理的并行策略选择转化为一个严密的带约束数学优化问题。全文不需要读者预先掌握复杂的分布式系统底层实现,只需具备基础的高等数学知识(一元函数求导、链式法则、二次不等式、基本概率期望与极限),即可跟随推导理解各并行维度的内在权衡,并算出特定硬件与模型下的最优配置。
</aside>
在建立数学方程之前,我们首先明确大模型推理的计算阶段划分、核心资源指标以及五种并行策略的严格定义。
Transformer 模型的自回归生成过程包含两个物理特性截然不同的阶段:
(注:若开启投机采样(Speculative Decoding),每步 Decode 会由轻量的 Draft 模块提出若干候选 Token,主干网络在 1 次前向传播中并行验证一小批 Token,单步参算的 Token 数因此从 1 个增加为 $1 + k_{\text{spec}}$ 个。后文将看到,这会显著改变 Decode 阶段算子所处的 Roofline 区间。)
在分析任意算子 $v$ 在第 $j$ 张 GPU 上的开销时,必须严格区分以下四个物理量:
$$ M_{v,j}:\text{常驻显存容量 (Byte)},\quad F_{v,j}:\text{执行的浮点运算量 (FLOPs)},\quad Q_{v,j}:\text{HBM 显存读写流量 (Byte)},\quad V_{v,j}:\text{网络发送字节数 (Byte)}. \tag{1} $$
这四个量不能混为一谈:
设单个流水线阶段(Stage)内共有 $G$ 张 GPU,五种并行策略对计算与状态的切分方式如下: