yuzheng
← 全部笔记

翻译与阅读笔记

图解大模型推理并行策略:从数据、张量到混合并行

全面解析大模型推理中的多卡并行切分策略:数据并行、张量并行、流水线并行、上下文并行与专家并行,权衡通信开销并给出实战决策路径。

修改文字

如何将大模型推理负载切分到多张 GPU 上,同时不让通信开销侵蚀性能?本文系统梳理数据并行(DP)、张量并行(TP)、流水线并行(PP)、上下文并行(CP)与专家并行(EP),阐述它们的混合编排方式,并给出根据模型、硬件和工作负载选择合适部署拓扑的工程检查清单与决策框架。


为什么需要分布式推理

4 张 GPU 拥有 4 个独立的显存空间和 4 组计算单元。它们并不是“容量自动乘以 4 倍”的单块 GPU。

在单张 GPU 上运行的算子内核,若没有显式的数据搬运,根本无法直接访问存放在另一张 GPU 上的权重参数或 KV Cache 显存块。

4 张独立 GPU 与互连总线(中文版)
                                                **图 1,   4 张独立 GPU 拥有各自独立的显存与计算单元,跨卡访问必须经由互连总线**

因此,分布式推理的第一步始终是数据放置(Data Placement)。推理服务引擎必须明确把模型权重、激活值、用户请求以及缓存的 Key/Value 分配给具体的设备。

任何跨设备的依赖,都会在请求的处理关键路径(Critical Path)上引入点对点传输或集合通信算子。

工程上的核心挑战在于:如何选择一种数据放置方案,在消除当前硬件资源瓶颈的同时,不引入超出互连带宽所能承受的通信开销。

在评估每种并行策略时,我们需要带着以下五个关键问题进行对比:

  • 切分了什么状态?(What state is partitioned?)
  • 保留/复制了什么状态?(What remains replicated?)
  • 互连总线上跨卡传输了什么?(What crosses the interconnect?)
  • 传输发生的频次是多少?(How often does that transfer occur?)
  • 该拓扑改善的是显存容量、单请求时延,还是系统吞吐量?(Does the layout improve memory capacity, request latency, or throughput?)

在本文的后续讨论中,我们统一设定一个基准硬件环境:一台搭载 4 张 GPU 的服务器。待部署的模型至少需要 2 张 GPU 的显存才能装下,且服务器同时承接简短对话与长上下文请求。我们保持模型、请求负载与底层硬件不变,仅改变推理任务在 4 张 GPU 之间的切分方式。

关于集合通信(Collective Communication)

本文会频繁使用“集合通信”这一概念。集合通信指多个 GPU 作为一个通信组共同交换或聚合数据的操作。例如,两张 GPU 分别计算某一层的部分输出,All-Reduce 会将各卡的部分和累加,并将完整结果同步回每张 GPU。通信组内的每张 GPU 都必须完成各自的任务,集合通信才算结束。因此,任何一张较慢的 GPU 或较慢的通信链路都会拖慢整个通信组(即长尾掉队效应,Straggler Effect)。这与两张卡之间的直接点对点传输不同。批注 2

作者批注 2展开阅读

集合通信(Collective Communication)是同一通信组中的多个进程或 GPU,按约定共同完成的数据交换操作。 在分布式训练中,一个参与者通常称为一个 rank。

理解它主要看两点:数据怎么分发,以及是否做求和等计算。

操作做什么结果放在哪里常见用途
Broadcast(广播)一个 rank 的数据复制给所有 rank所有 rank同步模型参数
Scatter(分散)一个 rank 把数据切成多份,每个 rank 收一份各 rank 持有不同分片分发数据
Gather(收集)把各 rank 的数据收集、拼接起来指定的一个 rank集中输出结果
AllGather(全收集)每个 rank 的分片都发给其他 rank所有 rank 都拿到完整数据收集分片参数或激活
Reduce(归约)对各 rank 的对应元素求和、取最大值等指定的一个 rank汇总统计量
AllReduce(全归约)归约后,让所有 rank 都拿到结果所有 rank数据并行梯度同步
ReduceScatter(归约后分散)先按对应元素归约,再把结果切分各 rank 拿到结果的一部分分片梯度同步
AllToAll(全互换)每个 rank 都给每个其他 rank 发送各自指定的分片各 rank 收到面向自己的数据MoE token 分发
Barrier(屏障)等所有 rank 都到达此处后再继续不产生数据结果阶段同步

用两个 GPU 看清几个关键操作

假设初始数据为:

GPU 0: [1, 2]
GPU 1: [3, 4]

AllGather:拼接,不计算。

GPU 0: [1, 2, 3, 4]
GPU 1: [1, 2, 3, 4]

AllReduce(SUM):对应位置相加,所有人拿完整结果。

GPU 0: [4, 6]
GPU 1: [4, 6]

ReduceScatter(SUM):对应位置相加,每人拿一片结果。

GPU 0: [4]
GPU 1: [6]

AllToAll:每人把输入按目标 rank 切分,然后交换。

GPU 0: [1, 3]   ← 收到两个 GPU 发给 GPU 0 的部分
GPU 1: [2, 4]   ← 收到两个 GPU 发给 GPU 1 的部分

最值得记住的关系是:

AllReduce≡ReduceScatter+AllGather\boxed{\text{AllReduce} \equiv \text{ReduceScatter} + \text{AllGather}}

也就是:先归约并分片,再把各分片收集完整。 这里表示结果语义上的等价,实际通信库可以采用不同算法。

放到大模型训练里,可以这样记:

  • 数据并行 DDP:常用 AllReduce,同步各 GPU 的梯度。
  • 参数分片 FSDP / ZeRO:常用 AllGather 收集参数、ReduceScatter 归约并分片梯度,具体取决于分片策略。
  • 张量并行 TP:根据矩阵如何切分,使用 AllReduce、AllGather 或 ReduceScatter。
  • 专家并行 MoE:常用 AllToAll,把 token 发给对应专家,再把输出发回。

一个容易忽略的点:集合通信要求各 rank 的调用顺序和数据约定匹配,但不代表每个操作都是 Barrier。 调用不匹配可能造成挂起或报错。


大模型推理基础与性能指标

在深入优化之前,我们需要先回顾大模型推理的基本运行方式。

大模型生成过程分为两个阶段:

大模型推理的 Prefill 与 Decode 阶段(中文版)

图 2 大模型推理包含 Prefill(预填充)与 Decode(解码)两个阶段

  • Prefill(预填充)阶段:处理输入的 Prompt 并构建对应的 Key-Value Cache。它同时处理大量 Prompt Token,计算以大规模矩阵乘法为主。Prefill 阶段直接决定了首字时延(Time to First Token,TTFT)。
  • Decode(解码)阶段:自回归逐步生成后续 Token,每步只生成一个 Token。每一步都需要重新读取模型权重并访问持续增长的 KV Cache。Decode 阶段直接决定了每输出 Token 耗时(Time Per Output Token,TPOT),在某些系统中也称为 Token 间时延(Inter-Token Latency,ITL)。
预填充与解码时间线(中文版)

图 3 Prefill 批量处理输入,Decode 逐步生成后续 Token

这两个阶段对硬件资源的利用特征截然不同:

  • Prefill 通常在大量 Prompt Token 上执行大尺寸矩阵乘法,算术密度高,通常处于**计算受限(Compute-bound)**状态。
  • Decode 执行小尺寸算子,每步反复读取权重与 KV Cache,通常处于**访存受限(Memory-bound)**状态。一种能够缩短 Prefill 时延的切分方案,如果给每个解码生成的 Token 都增加了跨卡同步开销,反而可能会拖慢 Decode 阶段。

关于推理内部执行路径更详细的剖析,可参阅此前的精读笔记:图解大模型推理:从 Prefill 到 Decode。

四大硬件瓶颈与服务指标

在生产部署中,系统通常受制于四类硬件资源:

  1. 参数显存(Parameter Memory):决定模型权重能否完整装入显存。
  2. KV Cache 显存(KV-cache Memory):限制驻留的上下文 Token 总量与系统最大并发请求数。
  3. GPU 计算耗时(Compute Time):单卡本地矩阵运算和内核执行耗时。
  4. 设备间通信耗时(Communication Time):卡间或节点间的数据传输与集合通信耗时。

服务端的观测指标必须与业务目标直接对应:

  • TTFT:衡量初始响应速度。
  • TPOT:衡量文本流式生成的步频与流畅度。
  • 请求吞吐量(Request Throughput):单位时间内完成的完整请求数(Req/s)。
  • Token 吞吐量(Token Throughput):区分统计输入 Token 与输出 Token 的处理速率(Tokens/s)。
  • 有效吞吐量(Goodput):仅统计满足服务等级目标(SLO/SLA)的请求吞吐。

对任何指标取平均值都会掩盖长尾延迟,因此必须同时监控中位数(P50)与高分位值(如 P95、P99)。

硬件资源与性能指标(中文版)

图 4 硬件资源与服务性能指标的对应关系

P95 TTFT 表示 95% 的请求都能在该阈值内生成首字。在相同的请求分布下,必须用它与服务 SLO 进行对齐评估。

注意:脱离延迟约束去谈“每秒生成多少 Token”毫无意义。一个过载饱和的集群可能报告极高的 Token 吞吐,但其长尾时延已严重超时违约。

这些服务指标揭示了系统未达标的表象,而底层显存、算力或通信的测量数据则指明了根因。接下来,我们需要针对具体的瓶颈来选择对应的并行策略。

各并行策略的核心对比如下:

大模型推理并行策略全景对比矩阵(中文版)

图 5 大模型推理常用并行策略在状态切分、复制状态、通信开销与频次上的全景对比

上表中的通信频次至关重要。除了单次传输的数据量大小,传输发生的频率往往对性能起决定性作用。例如,在每个 Transformer Decoder 层内部都重复执行的小尺寸集合通信,其累计开销通常远高于在流水线各阶段边界之间发生的一次大尺寸激活值点对点传输。

我们首先从数据并行讲起,因为数据并行将单次前向传播完全保留在单个副本内部;而其他并行方式则是将单次前向传播切分到多张 GPU 上,直接把通信开销置于执行路径之中。


1. 数据并行(Data Parallelism, DP)

数据并行在请求维度进行切分。

每个副本(Replica)都保存一份完整的模型,拥有独立的调度器与 KV Cache 池。网关路由负责将各个到来的请求分发给具体的副本。

推理数据并行架构(中文版)

图 6 推理数据并行:每个副本独立持有完整模型,调度器独立管理各自的 KV Cache

这与模型训练时的数据并行有本质区别:

在推理阶段,前向传播结束后不需要同步梯度(无需任何 All-Reduce 梯度通信),因此一个请求在被分配到特定副本后,整个生成生命周期都完全在本地执行。

假设一个由 2 张 GPU 组成的模型副本能够满足单请求延迟目标,但无法承受高峰期的请求到达速率。在我们的 4 卡服务器上,数据并行允许我们启动第 2 个独立的双卡副本。两个副本并发处理不同的请求,因此该拓扑的目标是提升系统整体吞吐量,而非降低单个请求的时延。

数据并行的 3 项工程检查

  1. 显存容量检查:每个副本都必须能完整容纳模型。这意味着 4 卡服务器上会存放两份完整的模型权重。因此,只有在单个副本已经能够完整塞入显存的前提下,数据并行才可行。它无法解决单个副本放不下的显存超限问题。
  2. 请求路由策略:当存在两个活跃副本时,路由层需要决定将新请求发往哪个队列和 KV Cache 池。朴素的轮询(Round-Robin)假设两边负载完全一致,但实际场景中一个副本可能排队更长、可用 KV 显存更少,或者另一副本可能已经缓存了该请求可复用的前缀(Prefix Cache)。
请求路由调度对比(中文版)

图 7 轮询调度对队列、KV 显存与前缀命中无感知;感知状态的路由可显著提升命中率与吞吐

  1. 准入容量控制(Admission Capacity):请求刚到达时占用的 KV 显存很小,但随着自回归 Decode 的推进,显存占用会持续线性增长。路由层必须为后续的生成长度预留显存预算;否则,副本可能会在 Prefill 阶段接纳请求,但在 Decode 途中发生 KV 显存耗尽(OOM 或被迫抢占调度)。

工程准则:在增加副本后,结合队列延迟、剩余 KV 显存余量以及前缀缓存(Prefix Cache)命中率来进行智能路由。在接纳请求前,为预期的解码长度预留出 KV 显存配额。

适用与不适用场景

当流量处于低谷时,第 2 个副本处于空闲状态,完全无法降低单个请求的时延。只有当单副本能够装下模型,且并发请求量足够高、能够让多个副本充分运转时,才应当采用数据并行。


2. 张量并行(Tensor Parallelism, TP)

假设模型无法放入单张 GPU,但若将其最大的权重矩阵切分到两张卡上即可装下。

数据并行无法解决这个问题,因为每个 DP 副本都必须持有完整模型。而张量并行能够解决,因为它让多张 GPU 在同一个 Transformer 层内部、协同处理同一个请求。

这使得张量并行首先是一个模型放置工具。它也能缩减大尺寸矩阵乘法的单步计算耗时,但前提是节省下来的计算时间必须大于引入的跨卡通信开销。

Megatron-LM 经典论文中给出了双层投影的切分构造。我们可以观察一个前馈网络(FFN)模块中的数据依赖关系:

第一层投影(Gate/Up Projection)产生中间激活值,第二层投影(Down Projection)将其投射回隐藏层维度。

MLP 矩阵切分结构(中文版)

图 8 MLP 模块中各张量的维度与依赖关系:输入 XX、权重 AA 与 BB、中间激活以及输出

在 2 个 Rank 的配置下:

  • 第一层权重矩阵 AA 按列切分(Column Parallel),每个 Rank 计算不同的中间特征分量。
  • 第二层权重矩阵 BB 按行切分(Row Parallel),每个 Rank 消费与其对应的中间特征切片。两张卡各自计算出部分输出张量(Partial Sum)。批注 1
作者批注 1展开阅读

因为这样配对,第一层算出的中间特征可以直接交给本卡的第二层使用,两层之间不用通信,最后再统一求和。

按文章的乘法约定,忽略偏置,整个 FFN 是:

Y=σ(XA)BY=\sigma(XA)B

其中 XX 是输入,σ\sigma 是逐元素激活函数。

① 第一层按列切:每张卡算一半中间特征

把 AA 左右切开:

A=[A0    A1]A=[A_0\;\;A_1]

于是:

XA=[XA0    XA1]XA=[XA_0\;\;XA_1]

逐元素激活可以分别计算,所以:

H=σ(XA)=[σ(XA0)⏟H0    σ(XA1)⏟H1]H=\sigma(XA)=[\underbrace{\sigma(XA_0)}_{H_0}\;\;\underbrace{\sigma(XA_1)}_{H_1}]

两张卡各持有同一份完整输入 XX,分别得到 H0H_0、H1H_1。每个分片里的特征值已经完整,可以直接做激活,无需跨卡求和。

② 第二层按行切:正好接住这一半特征

把 BB 上下切开,与 HH 的特征分片对应:

Y=HB=[H0    H1][B0B1]=H0B0+H1B1\begin{aligned} Y=HB &=[H_0\;\;H_1]\begin{bmatrix}B_0\\B_1\end{bmatrix}\\ &=H_0B_0+H_1B_1 \end{aligned}

因此,每张卡只需要本地数据:

第一层及激活第二层
GPU 0H0=σ(XA0)H_0=\sigma(XA_0)P0=H0B0P_0=H_0B_0
GPU 1H1=σ(XA1)H_1=\sigma(XA_1)P1=H1B1P_1=H_1B_1

最后用 All-Reduce 得到 Y=P0+P1Y=P_0+P_1。注意:P0P_0、P1P_1 都与完整输出同形状,但每个元素只算了部分贡献,所以要相加,不能拼接。

如果反过来呢?

第一层按行切,会得到:

XA=X0A0+X1A1XA=X_0A_0+X_1A_1

这时必须先跨卡求和,才能做激活,因为一般有:

σ(U+V)≠σ(U)+σ(V)\sigma(U+V)\ne\sigma(U)+\sigma(V)

而第一层按列切后,若第二层也按列切,每张卡通常又需要完整的 HH,得先收集中间特征。

所以这里的“先列后行”,本质上是沿同一个中间特征维度,切开 AA 的列和 BB 的行,让两次矩阵乘法在本卡连续完成,减少通信。

双 Rank 计算与 All-Reduce 规约(中文版)

图 9 各 GPU 利用本地权重分片执行矩阵乘法,随后通过 All-Reduce 累加合并出完整输出

每张 GPU 使用本地权重分片计算一次局部矩阵乘法,随后通过 All-Reduce 算子将这两个部分张量相加。在该集合通信完成之前,任何一张卡都不具备正确的层输出。

以下 Python 脚本直观展示了双 Rank 的数据依赖与计算过程:

import numpy as np

def tensor_parallel_output(x, a_shards, b_shards, activation=np.tanh):
    partials = []
    for a_shard, b_shard in zip(a_shards, b_shards, strict=True):
        hidden_shard = activation(x @ a_shard)
        partials.append(hidden_shard @ b_shard)
    return np.sum(partials, axis=0)

该函数严格遵循双 Rank 的依赖关系:每张卡计算局部乘积,最后的规约累加必不可少,因为每个局部结果都只包含了输出的一部分。

运行该测试产生如下输出:

tensor_parallel_output = [[-3.294438, 1.406896]]

未切分的单卡标准 MLP 运算产生完全相同的数值,各元素在浮点精度容差范围内严格一致。

自注意力(Self-Attention)机制采用相同的核心思路:

  • Query、Key、Value 投影矩阵按注意力头维度进行列切分。
  • 输出投影(Out Projection)矩阵按行切分,并通过 All-Reduce 聚合各头部的局部结果。
  • 具体的切分布局取决于多头注意力(MHA)、分组查询注意力(GQA)的配置以及模型专用的算子实现。
双 Rank 的张量并行 MLP(中文版)

图 10 双 Rank 的张量并行 MLP:第一层投影切分计算,第二层投影的部分结果通过 All-Reduce 聚合

在选择张量并行度(TP Degree)前,需要明确两个基本概念:

  • 张量并行通信组(TP Group):共同处理同一个请求的一组 GPU。
  • Rank:通信组内的具体 GPU 进程。TP=2 包含 2 个 Rank,TP=4 包含 4 个 Rank。
TP2 与 TP4 的切分粒度(中文版)

图 11 提高 TP 度数会使每个分片矩阵变小,但会增加参与集合通信的 Rank 数量

提高 TP 度数会将每个分片矩阵切得更细。以输出维度为 8,192 列的矩阵为例:

  • TP=2 时,每个 Rank 分配 4,096 列。
  • TP=4 时,每个 Rank 分配 2,048 列。

各 GPU 存储的权重更少,单次执行的矩阵乘法规模更小。这看似是直接的性能提升,但层计算无法使用孤立的局部结果继续向下传播,各个 Rank 必须互相交换或聚合数据。

因此,TP=4 虽然减少了每个 Rank 的本地计算量,但涉及的通信节点增加了一倍。TP=4 是否更快,完全取决于节省的计算时间与增加的通信开销之间的权衡。

张量并行的 3 项工程检查

  1. 显存是否真的能装下? 矩阵切分只减少了被切分的权重矩阵显存。其他显存分配可能在每个 Rank 上完全复制,包括部分不可分参数、临时中间激活缓冲区、CUDA Graph 预分配显存以及单卡固定的上下文管理开销。必须实测每张 GPU 的峰值显存占用,绝不能简单假设 TP=2 会让所有显存占用减半。
切分状态与复制状态的显存占用(中文版)

图 12 权重矩阵被均分,但 CUDA Graph、输入缓冲区与部分状态在各卡上依然是复制的

  1. Rank 之间的通信耗时占比多大? 张量并行的集合通信深入每个 Transformer 层内部(通常每层 2 次 All-Reduce)。较慢的互连带宽会导致 GPU 在每次局部计算完成后陷入等待。这种等待在 Prefill 和每一次 Decode 迭代中都会发生。
  2. 每个 Rank 的本地计算负载是否依然饱满? GPU 之所以能高效执行大尺寸矩阵乘法,是因为大规模并行线程能充分打满计算单元(SM 占用率高)。提高 TP 度数会减小局部矩阵尺寸;一旦尺寸过小,局部计算的执行效率就会大幅下降,而集合通信开销却并未减少。此外,注意力头数(必须被 TP 整除)和量化分组也会对合法的 TP 度数施加硬性约束。
张量并行的计算与通信权衡(中文版)

图 13 张量并行在本地计算与跨卡通信之间的权衡

在长 Prompt 的 Prefill 阶段,由于每个 Rank 分配到的矩阵仍然很大,TP=2 通常明显快于单卡。但在 Decode 阶段,单步只有一个 Token,计算量极小,往往更早触碰拐点。

工程准则:使用满足显存容纳与延迟要求的最小张量并行度。每次增加并行度后都必须重新基准测试。卡间链路决定了通信代价,NVLink / NVSwitch 与跨节点网络在延迟和带宽上有数量级的差距。

对于我们的 4 卡服务器:由于模型需要两张卡才能装下,基准方案应选择 TP=2。若 TP=2 满足延迟目标,剩余两张卡应搭建第二个 TP=2 副本以增加吞吐;只有在 TP=2 依然超显存或无法满足单请求时延 SLO 时,才升级到 TP=4。


3. 流水线并行(Pipeline Parallelism, PP)

当多张 GPU 之间的互连链路较慢(如缺乏 NVLink 的跨节点环境),张量并行每层高频的 All-Reduce 通信代价会极其高昂。

此时,流水线并行提供了另一种模型放置方案:沿模型深度切分。将连续的一组 Transformer 层放置在同一个阶段(Stage),仅当执行跨越阶段边界时才传递中间激活值(Hidden States)。

  • Stage 0 运行前一部分网络层,将输出激活值传递给 Stage 1。
  • 最后一个 Stage 运行末尾网络层,并计算最终的词表 Logits。

请求依然需要严格按顺序流经所有阶段,但任何一个阶段的权重都持久保留在其本地 GPU(或本地 TP 组)中。

流水线预热、稳态与排空(中文版)

图 14 流水线执行中的预热、稳态与排空阶段

这与传统的训练流水线同样不同:推理没有反向传播(Backward Pass)。利用持续批处理(Continuous Batching),多个请求组可以同时在不同阶段上流动。

因此,流水线并行以远低于张量并行的通信频次解决了跨设备的模型放置问题。这使它成为跨节点或缺乏高速对等互连硬件上的优秀候选方案。但它不会自动降低单个请求的端到端时延,因为该请求必须完整穿越所有阶段。

流水线依赖微批次(Microbatching)来为各个阶段提供并发任务:

微批次执行与流水线气泡(中文版)

图 15 在线推理中利用持续批处理填充流水线,减少预热与排空阶段的气泡占比

在在线服务中,这些微批次可能包含完全不同的请求。调度器必须在维持流水线顺序的同时填补新释放的槽位。

流水线并行的 3 项工程检查

  1. 工作负载是否具备足够的独立并发请求? 在预热(Warm-up)和排空(Drain)期间,阶段会处于空闲状态(流水线气泡 Bubble)。并发微批次过少会拉大空闲比例,在低负载下系统甚至退化为同一时间只有一个阶段在工作。
  2. 阶段之间的负载是否均衡? 80 层的模型在 PP=4 下看似每阶段 20 层,但相等的层数并不保证相等的耗时或显存。输入 Embedding、输出投影头、异构架构都会打破平衡。最慢的阶段决定了整体吞吐,较快的阶段只能等待。
  3. 激活值缓冲与运维容错成本:当请求在流水线中流动时,中间激活值必须在显存中保留缓冲;此外,任何一个阶段发生故障都会导致整个模型实例崩溃。
流水线阶段的耗时均衡(中文版)

图 16 不同阶段的计算耗时与负载均衡

对于我们的 4 卡服务器:只有在模型必须占满全部 4 张卡、且在线流量并发足以填满 4 个阶段时,PP=4 才有意义。如果硬件是“两张卡之间有高速互连、两个卡对之间链路较慢”,则阶段内 TP=2 + 阶段间 PP=2 的混合配置更为合理:高频的张量聚合被限制在高速链路内,低频的激活传递走较慢的跨组链路。


4. 上下文并行(Context Parallelism, CP)

假设模型权重已经可以装入显存,但面对超长上下文请求时出现以下两类瓶颈之一:

  • 长 Prompt 在 Prefill 阶段无法达到 TTFT 延迟目标;
  • 长文本的 KV Cache 消耗了过多显存,导致 Decode 阶段允许并发的请求数过低。

上下文并行通过**切分序列位置维度(Sequence Positions)**而非模型权重来解决这一问题。针对上述两类瓶颈,上下文并行演化出了两套截然不同的机制:

  • Prefill 上下文并行(PCP):切分输入 Prompt 的计算负载。
  • Decode 上下文并行(DCP):切分存储在显存中的历史 KV Cache。

Prefill 上下文并行(PCP)

PCP 将 Prompt Token 序列均分给各个 Rank。每个 Rank 为其本地负责的 Token 分块独立计算 Query、Key 和 Value 张量。但注意力运算机制要求每个 Query 必须能与全部所需的 Key/Value 进行交互。

常见实现方式有两种:

  1. 全收集方式(All-Gather):在每张卡上将完整的 K、V 张量收集齐备,各卡用本地 Query 计算全局上下文。这只在收集后的张量不超出显存时可行。
  2. 分块循环传递(Ring Attention 风格):各卡循环轮转部分 KV 分块,逐步增量聚合局部的注意力输出。这避免了在任何单卡上物化完整的全局序列,但需要多轮通信迭代。
Prefill 上下文并行(PCP)实现架构(中文版)

图 17 PCP 将 Prompt 切片分配给各卡,通过 Ring 风格通信或 All-Gather 实现全序列注意力

译图说明:右下角原文“Owns Tans”含义不清,保留原词并标注,未作猜译。

PCP 的工程关键在于找到开销与收益的临界点。必须在不同的 Prompt 长度区间分别压测 TTFT。超长 Prompt 包含充足的注意力计算量来摊销通信成本,而短 Prompt 则得不偿失。因此,混合流量网关通常只将命中长文本规则的请求路由到启用了 PCP 的工作节点。

Decode 上下文并行(DCP)

DCP 按照序列位置切分已经生成的 KV Cache 历史。

Decode 步只产生单 Token 的 Query,每个 Rank 让该 Query 与本地持有的历史 KV 分片计算局部注意力,随后各 Rank 聚合局部注意力输出与 Softmax 统计量:

import numpy as np

def merge_attention_partials(partials):
    global_max = max(part.maximum for part in partials)
    scales = [np.exp(part.maximum - global_max) for part in partials]
    denominator = sum(
        scale * part.exponential_sum
        for scale, part in zip(scales, partials, strict=True)
    )
    numerator = sum(
        (
            scale * part.weighted_value
            for scale, part in zip(scales, partials, strict=True)
        ),
        start=np.zeros_like(partials[0].weighted_value),
    )
    return numerator / denominator

每个 Rank 上报其局部的最大值(Maximum)、指数和(Exponential Sum)以及加权 Value 和(Weighted Value Sum)。

全局最大值(global_max)用于在最终归约前对各 Rank 的局部统计量进行动态重缩放(Rescaling)。这保证了在跨 KV 分片的情况下 Softmax 运算的数值稳定性,防止指数上溢。测试脚本计算得出的结果与未切分的全局注意力在浮点精度下完全吻合:

merged_attention_output = [0.769627, 0.779901]
PCP 与 DCP 的数据流对比(中文版)

图 18 DCP 将历史 KV Cache 按位置切分,解码时通过全局最大值重缩放安全合并注意力输出

译图说明:DCP 在本文中指解码上下文并行。图中的 Σ 仅为合并示意,正确合并还需对 Softmax 统计量重标定,不能直接相加局部注意力输出。

DCP 的 2 项工程检查

  1. 显存容量增益:对比切分前后单卡驻留 KV Token 量及系统最大并发数。
  2. 解码延迟代价:每个 Decoder 层都必须合并局部注意力结果,这会不可避免地拉长 TPOT。只有当获得的 KV 显存容量或所支持的上下文长度收益高于每步合并的延迟代价时,DCP 才有实用价值。

序列并行(Sequence Parallelism, SP)辨析:

Megatron 风格的序列并行通常是指在张量并行中,沿着序列维度对 LayerNorm 和 Dropout 的激活值进行切分以节省显存。而上下文并行(CP)是指将 Attention 机制本身的上下文序列跨卡切分。各开源框架在命名习惯上略有差异,切勿混淆。


5. 专家并行(Expert Parallelism, EP)

专家并行仅适用于混合专家(Mixture-of-Experts, MoE)架构。

在 MoE 模型中,如果在每张 GPU 上都复制存放全部专家权重,会占据极其庞大的参数显存,哪怕每个 Token 在实际计算时只激活其中的少数几个(Top-KK)专家。

专家并行将不同的专家权重分散存放在不同的 Rank 上,从而在单模型实例内分散专家参数的显存压力。

专家并行架构与动态路由(中文版)

图 19 专家并行将不同专家权重分配至不同 GPU,运行时由门控路由分发 Token

在运行时,门控路由(Router)为每个 Token 动态指定专家。系统将 Token 表征通过 All-to-All 集合通信分发到持有对应专家的 Rank,执行本地专家网络的前向计算,最后再通过反向 All-to-All 将输出汇总。

该机制节省了显存,但在核心请求路径上引入了依赖输入数据的 All-to-All 动态通信。关于 MoE 推理架构更深入的技术探讨,可参阅作者的专题笔记:MoE 推理工程深度解析。

专家并行的 3 项工程检查

  1. 权重放置检查:核算每个 Rank 分配到的专家显存,并确认共享的自注意力层、Embedding、缓冲区和 KV Cache 依然能够装入。专家并行不负责稠密共享部分的切分。
  2. 路由通信开销:与稠密张量并行不同,All-to-All 传输的数据量完全取决于路由器的分配结果。必须使用真实的生产 Prompt 负载分布来实测每张卡发送的数据量和 All-to-All 耗时。
  3. 负载倾斜与掉队者问题(Straggler):输入 Token 经常向少数热门专家倾斜,而 All-to-All 集合通信的结束时间取决于最繁忙的那个 Rank。平均负载会掩盖严重的掉队现象。必须监控单专家的最大 Token 数、本地计算耗时以及计算后的等待时间。
专家负载不均衡导致的掉队效应(中文版)

图 20 专家路由倾斜会导致各卡计算量严重不均,集合通信耗时由最慢的 Rank 决定

注意:注意力和专家层不需要采用相同的切分策略。稠密注意力层可以使用 DP 或 TP,而 MoE 前馈层则可以采用 EP。


6. 混合并行(Hybrid Parallelism)

单一的并行方式往往只能解决系统部署的一部分挑战。在实际的大规模推理集群中,通常需要将多种策略复合编排:

例如,利用张量并行让模型能够在单节点内装下;利用流水线并行将模型跨节点拉伸;再利用数据并行复制出足够多的副本以应对高峰请求吞吐。

对于独立的稠密模型,总的工作节点数等于各并行度的乘积:

Total Workers=DP×PP×TP\text{Total Workers} = \text{DP} \times \text{PP} \times \text{TP}

混合并行组合维度(中文版)

图 21 混合并行的多维组合公式:总节点数由 DP、PP、TP 共同决定

节点物理拓扑与混合映射

以部署在 2 台 4 卡服务器上的 8 张 GPU 为例。节点内 GPU 互连带宽极高且延迟极低,而跨节点网络带宽相对较低:

  • 方案 A:节点内采用 TP=4,节点间采用 PP=2。高频的 Transformer 层内 All-Reduce 完全限制在节点内部的高速总线上,仅在跨节点时传输阶段间的激活值。8 张卡共同构成 1 个大模型实例。
  • 方案 B:如果该模型在 4 张卡上已经能够装下,则另一种方案更加优越:节点内采用 TP=4,节点间采用 DP=2 构成两个独立副本。请求完全无需穿过慢速的跨节点网络。
混合并行与物理拓扑映射(中文版)

图 22 将高频的 TP 绑定在节点内高速链路,低频的 PP 或无通信的 DP 跨越节点网络

读图提示:推理 DP 的副本独立处理请求,图中的双向箭头不表示副本之间需要同步梯度。

构建混合拓扑时应遵循以下固定顺序:

  1. 首先选择能让单个 Stage 放下的最小 TP 组,并将其绑定在最快的互连链路上。
  2. 只有当模型尺寸必须跨越慢速网络边界时,才引入 PP。
  3. 在单个完整模型实例能够放下后,引入 DP 扩充并发吞吐。
  4. 针对实测的长上下文瓶颈引入 CP;针对 MoE 专家权重显存瓶颈引入 EP。

这一顺序能避免常见错误:在尚未明确具体要解决什么资源瓶颈之前就盲目叠加并行维度。每一个新增的并行轴都必须对应一个清晰的工程目标和可量化的测量依据。

预填充与解码解耦(PD 分离)

当 Prefill 和 Decode 两个阶段对硬件的诉求存在显著结构性冲突时,**PD 分离(Prefill and Decode Disaggregation)**成为极具价值的系统架构。

系统设立独立的 Prefill 工作池与 Decode 工作池:

  • Prefill 池配置高算力硬件,优先追求大规模矩阵计算吞吐;
  • Decode 池配置高显存带宽与大容量显存,优先追求高 KV 缓存容量与稳定的步进生成。
Prefill 与 Decode 分离架构(PD 分离)(中文版)

图 23 PD 分离架构:解耦计算密集与访存密集阶段,通过高速网络传输 KV Cache 状态

两池之间通过高速网络传输完成 Prefill 后的 KV Cache 状态。需要强调的是:PD 分离并不能替代 TP、PP 或 CP;每一个独立的工作池内部依然需要规划各自的模型放置与并行切分方案。


落地决策框架(Decision Framework)

在面对具体的业务场景与硬件条件时,建议自顶向下遵循以下决策流程。每当选择出足以解决当前瓶颈的最小改动方案后,立即停下来进行基准测试;只有当仍然存在未解决的瓶颈时,才继续向下推进。

原图中的并行配置决策流程(中文版)

图 24 原图中的并行配置决策流程(原图存在疑点,见下方说明)

读图提示:原图“规模 ≥ 互连带宽”的量纲与含义不明确,且混入“训练时延”。此处保留原图信息供对照,部署决策请结合正文。

  1. 第一步:判断完整模型是否能装入单卡?
    • 如果能装下:从单副本开始。当并发流量上升时,通过**数据并行(DP)**增加副本数量,并在网关路由器中启用感知前缀缓存的调度策略。
    • 如果装不下:寻找能装下模型的最小张量并行度(TP),并在该 TP 度数下压测 Decode 延迟。如果网络对等连接较慢、或者各网络层组执行时间均衡,可考虑引入流水线并行(PP)。

原图待修复:指标与性能剖析矩阵

图 25 原始指标与性能剖析矩阵(文字不清,保留原图)

原图说明:图内多处文字不是可辨认的正常英文,暂时保留原图,未制作中文版。

  1. 第二步:模型是否需要跨节点?
    • 将 TP 严格限制在节点内高速互连链路上;跨节点使用 PP。在加大负载前,必须实测各 Stage 的时间平衡性。
  2. 第三步:长文本请求是否超标?
    • 分开评估:若超长 Prompt 的 TTFT 无法达标,评估 PCP;若长文本的 KV Cache 显存耗尽导致并发崩溃,评估 DCP。切勿在没有 Prompt 分段压测数据的情况下盲目开启。
并行策略决策矩阵(中文版)

图 26 并行策略决策矩阵

  1. 第四步:MoE 模型是否需要专家并行?
    • 只有在专家权重导致显存无法装下时才测试 EP,同时必须独立追踪 Token 分布倾斜度与 All-to-All 耗时。切勿将 EP 引入稠密模型,也切勿仅通过整体 GPU 利用率来推断 EP 的收益。

核心原则:只有当性能测量明确指出了某种资源限制时,才开启针对该限制的并行模式。首先定位受限的硬件资源,然后选择通信开销与当前物理互连拓扑相匹配的切分方案。


总结

在大模型推理工程中,并行策略本质上是数据放置(Placement)决策,而不是调节 GPU 数量的简单参数。

每一次切分都会引入通信、同步与调度上的额外开销。只有当切分节省出来的显存或计算耗时在实际硬件和真实请求分布上能够显著超越这些开销时,优化才具有正向收益。

在满足延迟与吞吐目标的方案中,最简单的那一个永远是最好的方案。 其他一切推断,都必须以严格的端到端压测数据为准。


← 返回笔记列表回到顶部 ↑