数据、张量、流水线、专家与混合并行
并行策略将推理工作分布到多个设备上,以提高吞吐量、容纳更大的模型并更高效地利用硬件。随着模型变得更大、更复杂,新的并行策略也在不断涌现。
数据并行
数据并行(DP)通过将单个请求或微批次(microbatch)分布到同一模型的多个副本上,来提高总吞吐量。每个副本都有自己的模型权重副本,并独立处理不同的请求。一个副本可以在单块 GPU 上运行,也可以在内部使用张量并行或流水线并行横跨多块 GPU。
AI 团队通常使用路由器将独立请求分布到各副本。对于批处理工作负载,一个更大的批次也可以拆分为更小的微批次,每个副本同时处理一个不同的微批次。这种方法允许处理更多并发的请求。
因此,增加更多副本会提高总吞吐量和并发度。它并不会直接降低单个请求的计算延迟,因为该请求仍然由单个副本处理。
数据并行还提供了一个故障边界。当健康检查将某个副本标记为不可用时,路由器可以停止向它发送新请求,并使用其余副本。这提高了可用性,不过健康的副本需要有备用容量来吸收流量。故障副本上任何飞行中的请求和本地 KV 缓存也可能丢失。
这里的一个关键考虑是路由请求的方式。对于 LLM 推理,请求在提示词长度、输出长度和 KV 缓存使用上差异很大。GPU 容量也很昂贵,因此负载分布不均会让一些副本过载,而另一些则利用不足。因此,生产环境中的路由器在选择副本时通常会考虑多种信号。更多信息请参阅推理路由。
数据并行最主要的成本是内存重复。每个副本都需要一份模型权重,并维护独立的 KV 缓存。如果一份模型无法放进单个 worker,仅靠数据并行并不能解决问题。请将其与张量或流水线并行结合使用。
张量并行
张量并行(TP)将模型的权重分布到多块 GPU 上,以部署无法容纳进单块 GPU 内存的大型模型。它分割每个模型层内的张量,因此每台设备持有该层的分片,而不是副本。
这是必要的,因为许多现代 LLM 即使经过量化也无法放进单块 GPU,尤其是还要计入 KV 缓存和运行时开销。例如,FP8 权重大约每十亿参数需要 1 GB,因此像 Llama 3.1 405B 这样的模型仅权重就需要大约 405 GB。这已经超过了当前任何单块 GPU 的内存,所以单台设备通常无法加载该模型或高效地服务它。
在矩阵乘法等操作期间,每块 GPU 持有权重张量的一个分片并计算结果的一部分。然后,GPU 使用集体通信操作(如 all-reduce 或 all-gather,具体取决于张量的划分方式)交换或合并中间结果。
通过分布权重、缓存和计算,张量并行使超出单块 GPU 容量的模型变得可服务。当并行计算的收益大于新增的通信成本时,它还可以降低单个请求的延迟。
通信是核心的权衡。张量并行中的 GPU 必须在一个请求流经 Transformer 各层时反复交换中间结果。因此,互连带宽和延迟对性能影响很大。
张量并行在单个节点内由高带宽互连(如 NVLink)连接的 GPU 之间效果最好。将张量拆分到位于不同节点上的 GPU,可能会使网络通信成为主要瓶颈。你跨设备分片的 GPU 数量也受到限制;它通常必须整除注意力头的数量,因此你无法扩展到任意的 GPU 数量。
在对模型进行分片以执行张量并行之前,请检查权重、KV 缓存和其他推理内存需求是否能轻松容纳进单块 GPU,尤其是在量化之后。单 GPU 服务完全避免了跨设备同步,而通过张量并行增加 GPU 很少能带来线性加速。当内存占用或延迟目标证明通信开销合理时,再使用张量并行。
流水线并行
流水线并行(PP)将模型的层划分为按顺序排列的块,称为阶段(stage),每个阶段分配给单独的设备。数据像流水线一样流过这些阶段,一台设备的输出成为下一台的输入。例如,在四路流水线中,每台设备处理模型四分之一层。
与张量并行不同,流水线并行不要求设备在每一层内合并部分结果。一个阶段只有在完成自己所有层之后,才会将激活传递给下一阶段。这种较低的通信频率使流水线并行更适合低带宽连接,包括节点之间的互连。
代价是流水线气泡(pipeline bubble)。一个阶段可能在等待上一阶段工作时处于空闲。一个慢的或内存密集的阶段也会拖累其后的每一个阶段。因此,阶段边界应平衡执行时间和内存需求,而不仅仅是层数。
为了缩短空闲期,服务器可以让多个请求或微批次同时处于飞行中。当后面的阶段处理一个微批次时,前面的阶段可以开始处理下一个微批次。这种调度机制可以提高吞吐量,尽管它并不能完全消除流水线开始和结束时的空闲时间。更多信息请参阅 GPipe 论文。
请注意,流水线并行并不一定会降低单个请求的延迟。每个请求仍然按顺序通过所有阶段,并付出激活传输的成本。当有足够的并发工作来保持流水线填满时,或者当更低的通信频率比流水线气泡更重要时,流水线并行效果最好。
专家并行
专家并行(EP)是用于混合专家(MoE)模型的专用并行策略。在这些模型中,每个 token 只激活一部分专家。专家并行不是将全部专家复制到每台设备(例如 GPU)上,而是将专家本身拆分到不同的设备上。
每块 GPU 只持有部分专家的完整权重,而不是全部。这意味着每块 GPU 只处理分配给存储在其上的专家的那些 token。相比之下,如果你对 MoE 模型应用张量并行,它只是对所有专家的权重矩阵进行切片,并将这些切片分布到所有设备上。
通过使用专家并行,GPU 可以被更高效地利用。与在每台设备上复制整个模型相比,这减少了内存和计算开销。
混合并行
对于某些模型,仅仅依赖单一的并行策略往往是不够的。混合并行结合两种或多种并行技术,以实现更好的可扩展性、效率和硬件利用率。
一个典型的混合配置可能如下所示(组合数据并行和张量并行):
如果你有 8 块 GPU,你可以对前四块 GPU 应用张量并行(TP=4),然后使用数据并行(DP=2)将该配置复制到其余的 GPU 上。
请注意,这只是可能的组合之一,每种组合各有优缺点。在上面的例子中,张量并行在 GPU 之间引入了通信开销,尤其是在推理期间。因此,使用较高的 TP 度并不总是意味着更好的性能。
另一种配置是减少张量并行并增加数据并行。例如,你可以设置 TP=2 和 DP=4:
这减少了跨 GPU 通信,可能有助于降低推理期间的延迟。然而,有一个问题:模型权重占用 GPU 内存的很大一部分,尤其是对于大型模型。降低张量并行意味着更少的 GPU 共享模型,留给 KV 缓存的空间更少。这可能会削弱诸如前缀缓存之类的推理优化。
这些权衡并非张量和数据并行所独有。在设计混合并行方案时,务必根据你的具体模型大小、硬件配置和推理需求,对不同的配置进行基准测试。没有放之四海而皆准的配置。最优策略通常是通过调优和实验发现的。