<aside>

本文数学推导由 GPT6-Astra 完成,文本部分由人类润色。

</aside>

<aside>

导读:在大模型推理部署中,工程师需要在给定数量的 GPU 上组合配置五种并行策略:DP(数据并行)、TP(张量并行)、EP(专家并行)、CP(上下文并行)和 PP(流水线并行)。

许多经验总结试图用一张简单的正负号表来概括各并行维度的作用,例如“增大 TP 降低显存”、“增大 CP 降低 KV 缓存”或“MoE 模型必须使用大 EP”。然而在实际生产环境中,各并行维度共享同一批物理 GPU,调大某一个并行度必然要求按比例缩小另一个并行度。脱离了“固定总资源”这一约束前提,单纯看某一项指标的变化往往会得出错误的结论。

本文旨在将大模型推理的并行策略选择转化为一个严密的带约束数学优化问题。全文不需要读者预先掌握复杂的分布式系统底层实现,只需具备基础的高等数学知识(一元函数求导、链式法则、二次不等式、基本概率期望与极限),即可跟随推导理解各并行维度的内在权衡,并算出特定硬件与模型下的最优配置。

</aside>


01. 推理系统的物理基础与耗时模型

在建立数学方程之前,我们首先明确大模型推理的计算阶段划分、核心资源指标以及五种并行策略的严格定义。

1.1 推理的两个计算阶段:Prefill 与 Decode

Transformer 模型的自回归生成过程包含两个物理特性截然不同的阶段:

  1. Prefill(预填充阶段):系统一次性接收用户输入的完整提示词序列(设长度为 $S$),并行完成所有输入词元(Token)的前向计算,生成第一个输出 Token,并将每一层计算出的键(Key)和值(Value)向量保存在显存中,形成 KV Cache(键值缓存)。由于 $S$ 通常较大,矩阵乘法规模饱满,该阶段主要受限于 GPU 的计算能力(Compute-bound,计算受限)。
  2. Decode(解码阶段):系统以自回归方式逐个生成后续 Token。每生成一个新 Token,模型只需输入当前最新的 1 个 Token,但必须从高带宽显存(HBM)中读取完整的模型权重以及此前累积的所有历史 KV Cache。由于单步参算的 Token 数量较少,算力利用率偏低,该阶段主要受限于显存带宽与跨卡网络通信(Memory/Communication-bound,访存与通信受限)。

(注:若开启投机采样(Speculative Decoding),每步 Decode 会由轻量的 Draft 模块提出若干候选 Token,主干网络在 1 次前向传播中并行验证一小批 Token,单步参算的 Token 数因此从 1 个增加为 $1 + k_{\text{spec}}$ 个。后文将看到,这会显著改变 Decode 阶段算子所处的 Roofline 区间。)

1.2 四个相互独立的资源度量

在分析任意算子 $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} $$

这四个量不能混为一谈:

1.3 五种并行维度的定义

设单个流水线阶段(Stage)内共有 $G$ 张 GPU,五种并行策略对计算与状态的切分方式如下: