论文标题:Orca:一种用于基于 Transformer 的生成式模型的分布式服务系统
作者/机构:Gyeong-In Yu and Joo Seong Jeong, Seoul National University; Geon-Woo Kim, FriendliAI and Seoul National University; Soojeong Kim, FriendliAI; Byung-Gon Chun, FriendliAI and Seoul National University

A1 主要贡献

本文旨在解决现有推理服务系统在处理基于 Transformer 的生成式模型(如 GPT-3)时遇到的性能瓶颈。这类模型的自回归特性要求为单个推理请求多次运行模型,每次迭代生成一个 token。现有系统采用请求级调度,即一个批次(batch)中的所有请求必须全部完成后才能返回结果,这导致了两个核心问题:1) 提前完成的请求必须等待批次中其他请求,增加了不必要的延迟;2) 新到达的请求必须等待整个当前批次处理完毕,增加了排队时间。

为解决这些问题,本文提出了以下核心贡献:

  1. 迭代级调度(Iteration-level Scheduling):提出一种新的调度机制,将调度的粒度从请求级降低到迭代级。在该机制下,调度器每次只调用执行引擎运行模型的一个迭代。这使得新到达的请求只需等待当前迭代完成即可加入处理队列,而已完成的请求也可以在每次迭代后立即返回给客户端,从而显著降低了排队延迟和请求处理延迟。
  2. 选择性批处理(Selective Batching):为了在迭代级调度的基础上有效应用批处理,本文提出了选择性批处理技术。由于迭代级调度会导致批次内请求已处理的 token 数量不同,从而使得 Attention 操作的输入张量形状各异,无法直接批处理。选择性批处理通过仅对模型中的部分操作(如矩阵乘法和 LayerNorm)应用批处理,而对 Attention 操作则单独处理每个请求,从而解决了这一难题。该方法之所以有效,是因为 Attention 操作不含模型参数,批处理对其内存读取效率的提升有限。
  3. 分布式服务系统 ORCA:基于上述技术,本文设计并实现了一个名为 ORCA 的分布式服务系统。ORCA 结合了层内(intra-layer)和层间(inter-layer)模型并行策略,以支持数百亿参数规模的大型模型。系统还设计了新的调度算法,并考虑了内存管理和跨工作节点的流水线执行优化。
  4. 性能验证:在 GPT-3 175B 模型上的评估显示,与 NVIDIA FasterTransformer 相比,ORCA 在相同延迟水平下实现了 36.9 倍的吞吐量提升,证明了所提方法的显著优越性。

A3 背景知识/关键Observation/设计原则

GPT 的推理过程

GPT 模型推理的自回归特性。GPT 是一种基于 Transformer 架构变体的自回归语言模型,它接收一系列输入 token 并通过生成后续 token 来完成序列。图 1a 展示了一个三层 GPT 模型的简化计算图,其中节点代表 Transformer 层,边代表层间依赖。模型按节点上的数字顺序执行,相同颜色的节点共享同一组模型参数。生成的输出 token 会被反馈回模型用于生成下一个 token,形成一种顺序、逐个的推理过程。我们将运行模型所有层的过程定义为一次迭代(iteration)。图 1a 的示例包含三次迭代:第一次迭代(“iter 1”)处理所有输入 token(“I think this”)并生成下一个 token(“is”),这一步构成了初始化阶段(initiation phase);接下来的两次迭代(“iter 2” 和 “iter 3”)构成了增量阶段(increment phase),每次接收前一迭代的输出 token 并生成新 token,直到生成特殊的序列结束符 “”。增量阶段包含多次迭代,因为每次迭代只处理一个 token,而初始化阶段通常通过并行处理所有输入 token,在一次迭代中完成。

图1:GPT推理过程、Transformer层及内部状态使用图示。
图1:GPT推理过程、Transformer层及内部状态使用图示。

Transformer 层的内部状态。与原始 Transformer 包含编码器和解码器两个层堆栈不同,GPT 架构只包含一个解码器层堆栈。图 1b 展示了 GPT 中使用的 Transformer 层。Attention 操作是其核心,它计算一系列 token 的加权平均值,使得序列中的每个 token 都能感知到其他 token。Attention 操作接收三个输入:query、key 和 value,通过计算当前 token 的 query 与所有目标 token 的 key 的点积,应用 Softmax 获得权重,最后对所有 value 进行加权求和。由于 Attention 操作需要所有先前 token 的 key 和 value,我们将这些 key 和 value 视为应在多次迭代间保持的内部状态。一种朴素的无状态推理方式会在每次迭代时重新计算所有 token 的 key 和 value,效率低下。为了避免这种重复计算,fairseq 【索引43,fairseq: A Fast, Extensible Toolkit for Sequence Modeling,2019,NAACL】 提出了增量解码(incremental decoding),即保存 key 和 value 以便在后续迭代中重用,FasterTransformer 【索引4,NVIDIA FasterTransformer】 和 Megatron-LM 【索引3,Megatron-LM】 等系统也采用了这种方法。

Transformer 状态大小的可变性。图 1c 展示了 Transformer 和 LSTM 【索引25,Long Short-Term Memory,1997,Neural Computation】 的状态使用模式对比。主要区别在于,Transformer 的状态(Attention 的 key k 和 value v)大小随迭代次数增加而增长,而 LSTM 的状态(内部记忆 c 和层输入/输出 h)大小保持不变。在处理索引为 t 的 token 时,Attention 操作需要使用之前所有的 key $k_{l,1:t-1}$ 和 value $v_{l,1:t-1}$,以及当前的 key $k_{l,t}$ 和 value $v_{l,t}$。因此,Attention 操作需要处理的张量形状会根据已处理的 token 数量而变化。Transformer 层的其他操作按顺序包括:层归一化(LayerNorm)、QKV 线性变换(获取 query, key, value)、Attention 输出线性变换、残差连接(Add)、另一次层归一化、多层感知机(MLP)以及另一次残差连接。

机器学习推理服务系统

推理服务系统的目标与功能。随着机器学习驱动应用需求的增长,机器学习推理服务已成为现代数据中心的关键工作负载。用户向推理服务提交请求,服务利用预置资源(通常是 GPU、TPU 等加速器)运行 DNN 模型生成输出。一个优秀的推理服务应在合理的成本内提供低延迟和高吞吐量。为满足这些要求,运营商常使用 Triton Inference Server 【索引7,NVIDIA Triton Inference Server】 和 TensorFlow Serving 【索引42,TensorFlow-Serving: Flexible, High-Performance ML Serving,2017,NIPS Workshop on Machine Learning Systems】 等系统。这些系统位于底层模型执行引擎(如 TensorRT 【索引6,NVIDIA TensorRT】、TVM 【索引14,TVM: An Automated End-to-End Optimizing Compiler for Deep Learning,2018,OSDI】)之上,负责接收推理请求、调度引擎执行并发送响应。它们关注的重点包括执行批处理、多模型版本选择、多模型部署等。

图2:使用现有系统为生成式语言模型提供服务的整体工作流程。
图2:使用现有系统为生成式语言模型提供服务的整体工作流程。

批处理(Batching)的重要性。批处理是利用 GPU 等加速器实现高利用率的关键。通过批处理,来自多个请求的输入张量被合并成一个大的输入张量,这使得加速器能够更好地利用其大规模并行计算单元,从而提高引擎的吞吐量。此外,对于内存密集型操作,跨多个请求重用从片外内存加载的模型参数也是批处理执行的另一个优点。图 2 展示了现有服务系统和执行引擎为生成式语言模型服务的工作流程。服务系统(如 Triton)的调度器负责 ➀ 从队列中检索请求创建批次,并 ➁ 调度执行引擎(如 FasterTransformer)处理该批次。执行引擎 ➂ 通过运行模型多次迭代来处理接收到的批次,并 ➃ 将生成的文本返回给服务系统。

图3:一个请求具有相同输入长度,但部分请求比其他请求提前完成的示例。阴影部分表示输入token。“-”表示由调度产生的额外计算的输入和输出。
图3:一个请求具有相同输入长度,但部分请求比其他请求提前完成的示例。阴影部分表示输入token。“-”表示由调度产生的额外计算的输入和输出。

挑战与解决方案

挑战1:提前完成和延迟加入的请求。现有系统的一个主要限制是,服务系统与执行引擎仅在两种情况下交互:(1) 服务系统在空闲引擎上调度下一个批次;(2) 引擎完成处理当前批次。这种请求粒度的调度机制,要求引擎在批次中所有请求完成前保持批次固定。这对于生成式模型服务是有问题的,因为批次中的每个请求可能需要不同数量的迭代,导致某些请求比其他请求提前完成。如图 3 所示,尽管请求 x2 比 x1 提前完成,引擎仍需在所有迭代中为“活跃”和“非活跃”请求执行计算。为非活跃请求进行的额外计算限制了批处理执行的效率。更糟糕的是,这种行为阻止了已完成请求的提前返回,增加了大量延迟。同样,在当前批次执行过程中新到达的请求必须等待整个批次完成,这显著增加了排队时间。因此,我们认为当前的请求级调度机制无法高效处理具有多迭代特性的工作负载。值得注意的是,在语言模型训练中不存在此问题,因为训练过程使用 teacher forcing 技术【索引64,A Learning Algorithm for Continually Running Fully Recurrent Neural Networks,1989,Neural Computation】在单次迭代中完成整个批次的处理。

图4:ORCA的系统概览。组件间的虚线交互表示该交互在执行引擎的每次迭代中都发生。$x_i^j$是第i个请求的第j个token。阴影token表示从客户端接收的输入token,而未着色的token是由ORCA生成的。例如,请求x1最初带有两个输入token ($x_1^1, x_1^2$),目前已运行两次迭代,生成了$x_1^3$和$x_1^4$。另一方面,请求x3只包含输入token ($x_3^1, x_3^2$),因为它尚未运行任何迭代。
图4:ORCA的系统概览。组件间的虚线交互表示该交互在执行引擎的每次迭代中都发生。$x_i^j$是第i个请求的第j个token。阴影token表示从客户端接收的输入token,而未着色的token是由ORCA生成的。例如,请求x1最初带有两个输入token ($x_1^1, x_1^2$),目前已运行两次迭代,生成了$x_1^3$和$x_1^4$。另一方面,请求x3只包含输入token ($x_3^1, x_3^2$),因为它尚未运行任何迭代。

解决方案1:迭代级调度。为了解决上述限制,我们提议在迭代粒度上调度执行。调度器重复以下过程:(1) 选择要运行的请求;(2) 调用引擎对选定的请求执行一次迭代;(3) 接收该次迭代的执行结果。由于调度器在每次迭代后都会收到返回,它可以检测到请求的完成并立即将其生成的 token 返回给客户端。对于新到达的请求,它在当前调度迭代执行后就有机会开始处理,从而显著减少排队延迟。通过迭代级调度,调度器可以完全控制每次迭代中处理哪些请求以及处理多少个请求。图 4 描绘了使用迭代级调度的 ORCA 系统架构和整体工作流程。ORCA 暴露一个端点,新请求到达后被放入请求池(request pool)。调度器从池中选择一组请求,调度执行引擎运行模型的一次迭代,接收输出 token,并更新请求池。引擎是一个执行实际张量操作的抽象,可以跨多台机器的多个 GPU 并行化。与图 2 中必须等待批次中所有请求完成的方法不同,ORCA 的调度器可以在每次迭代时更改要处理的请求。

挑战2:批处理任意请求集。在实践中使用迭代级调度时,一个主要挑战是批处理。为实现高效率,执行引擎应能以批处理方式处理任何选定的请求集。不幸的是,即便是任意一对请求 $(x_i, x_j)$,也无法保证它们的下一次迭代可以被合并成批处理版本。存在三种无法批处理的情况:(1) 两个请求都处于初始化阶段且输入 token 数量不同(如图 4 中的 x3 和 x4);(2) 两个请求都处于增量阶段但正在处理不同索引的 token(x1 和 x2);(3) 两个请求处于不同阶段(x1 和 x3)。要对多个请求的执行进行批处理,每个请求的执行必须由相同的操作组成,且每个操作消耗相同形状的输入张量。上述三种情况都违反了这一要求,因为输入张量的“长度”维度或 Attention 状态张量的形状不同。这大大降低了在实际工作负载中进行批处理的可能性,使得使用大批次大小变得不切实际。

图5:ORCA执行引擎使用选择性批处理在一个批次请求上运行Transformer层的示意图。为简化起见,我们只描绘了QKV线性、Attention和Attention输出线性操作。
图5:ORCA执行引擎使用选择性批处理在一个批次请求上运行Transformer层的示意图。为简化起见,我们只描绘了QKV线性、Attention和Attention输出线性操作。

解决方案2:选择性批处理。我们提出了选择性批处理(selective batching),这是一种允许在组合批次时具有高度灵活性的批处理执行技术。该技术不是对模型的所有张量操作进行“批处理化”,而是有选择地仅对部分操作应用批处理。上述批处理问题根源在于不规则形状的输入(或状态)张量无法合并。然而,并非所有操作都与不规则形状的张量不兼容。像非 Attention 矩阵乘法和层归一化等操作可以通过展平张量来处理。例如,图 4 中请求 x3 和 x4 的输入张量可以构成一个形状为 $[\sum L, H] = [5, H]$ 的二维张量,并输入到所有非 Attention 操作中。而 Attention 操作需要请求的概念(即批次维度)来计算同一请求内 token 之间的注意力。选择性批处理利用了不同操作的特性:它在 Attention 操作前分割批次,对每个请求单独处理 Attention,然后将结果合并,而对模型的其他操作则应用 token 级的批处理。图 5 展示了选择性批处理机制。输入张量形状为 [7,H],在 Attention 操作前通过 Split 操作分割,分别执行 Attention,然后通过 Merge 操作合并回 [7,H] 的张量。ORCA 通过 Attention K/V 管理器来维护和提供增量阶段请求所需的历史 key 和 value。

A2 方法细节

4.1 分布式架构

并行化策略。近期的研究【索引12,Language Models are Few-Shot Learners,2020,NeurIPS;索引31,Scaling Laws for Neural Language Models,2020,arXiv】表明,扩展语言模型能显著提升模型质量,因此为大型语言模型提供服务支持变得愈发重要,尤其是在模型无法装入单个 GPU 时。此时,必须将模型参数及相应计算分割并分布到多个 GPU 和机器上。ORCA 采用了已知的 Transformer 模型并行化技术:层内并行(intra-layer parallelism)层间并行(inter-layer parallelism)。这两种模型并行策略,也被 FasterTransformer【索引4,NVIDIA FasterTransformer】使用,最初是为分布式训练开发的。层内并行【索引55,Mesh-TensorFlow: Deep Learning for Supercomputers,2018,NeurIPS;索引58,Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism,2019,arXiv】将矩阵乘法(即 Linear 和 Attention 操作)及其相关参数分割到多个 GPU 上。而层间并行则将 Transformer 层分割到多个 GPU 上,ORCA 为每个 GPU 分配相同数量的 Transformer 层。图 6 展示了将层内和层间并行应用于一个 4 层 GPT 模型的示例,其中 4 个层被分成 2 个层间分区,每个分区内的层又被细分为 3 个层内分区,共使用 6 个 GPU。

图6:层内和层间并行的一个示例。垂直虚线表示层间分区,水平线表示层内分区。
图6:层内和层间并行的一个示例。垂直虚线表示层间分区,水平线表示层内分区。

图7:使用图6中并行化配置的ORCA执行引擎分布式架构示意图。例如,图6中的第一个层间分区(Layer1和Layer2)被分配给Worker1,而第二个分区被分配给Worker2。
图7:使用图6中并行化配置的ORCA执行引擎分布式架构示意图。例如,图6中的第一个层间分区(Layer1和Layer2)被分配给Worker1,而第二个分区被分配给Worker2。

ORCA 执行引擎架构与流程。ORCA 执行引擎支持上述分布式执行技术。图 7 描绘了 ORCA 引擎的架构。每个工作进程(worker process)负责模型的一个层间分区,并且可以位于不同的机器上。每个 worker 管理一个或多个专用于控制 GPU 的 CPU 线程,其数量取决于层内并行的程度。引擎的执行流程如下:当引擎被调度为一批请求运行一次迭代时,引擎主节点(engine master)将批次信息转发给第一个工作进程(Worker1)。这些信息包括当前迭代的 token 和一个控制消息(包含批次内请求的 ID、当前 token 索引、输入 token 数量等)。Worker1 的控制器(controller)将信息传递给 GPU 控制线程,后者解析信息并向其关联的 GPU 发出适当的 GPU 核函数。同时,Worker1 的控制器将控制消息转发给下一个 worker(Worker2)的控制器,而无需等待 Worker1 的 GPU 核函数完成。最后一个 worker(Worker2)的控制器会等待其 GPU 核函数完成,以便获取每个请求的输出 token 并将其发送回引擎主节点。

控制与数据平面分离。为了使 GPU 尽可能保持繁忙,ORCA 引擎的设计旨在最小化 CPU 和 GPU 之间的同步。我们观察到,当前的分布式推理系统(如 FasterTransformer【索引4,NVIDIA FasterTransformer】和 Megatron-LM【索引3,Megatron-LM】)在每个进程接收控制消息时都会有 CPU-GPU 同步,因为它们通过 GPU-to-GPU 通信通道(NCCL【索引5,NVIDIA NCCL】)交换消息,这在每次迭代中都会发生,带来了不可忽视的性能开销。与此不同,ORCA 将控制消息(以及 token)和张量数据传输的通信通道分开。如图 7 所示,ORCA 引擎仅使用 NCCL【索引5,NVIDIA NCCL】交换由 GPU 生产和消费的中间张量数据(虚线箭头表示),而 CPU 线程用于发出 GPU 核函数的控制消息则通过一个不涉及 GPU 的独立通信通道(如 gRPC【索引2,gRPC】)在引擎主节点和 worker 控制器之间发送。

4.2 调度算法

迭代级 FCFS 调度。ORCA 调度器在每次迭代时决定选择和处理哪些请求。由于选择性批处理技术允许引擎以批处理方式运行任何请求集,调度器在组合批次时具有很高的灵活性。ORCA 调度器采用一种简单的算法,不改变客户端请求的处理顺序,即先到先处理。我们将其定义为迭代级先到先服务(FCFS)属性:对于请求池中的任意一对请求 $(x_i, x_j)$,如果 $x_i$ 比 $x_j$ 先到达,那么 $x_i$ 已运行的迭代次数应该等于或多于 $x_j$。这意味着某些后到达的请求如果需要较少的迭代次数,可能会比先到的请求更早返回给客户端。

调度考虑因素:批次大小与内存。调度器还需要考虑两个额外因素:批次大小增加带来的收益递减和 GPU 内存限制。增加批次大小会用延迟换取吞吐量,但当批次大小变得很大时,吞吐量的增益会减小。因此,ORCA 也设置了最大批次大小(max batch size),系统操作员可以调整此参数以在满足延迟预算的同时最大化吞吞吐量。另一个因素是 GPU 内存。虽然中间结果的缓冲区可以被重用,但 Attention K/V 管理器用于存储 key 和 value 的缓冲区在请求完成前无法回收。一个朴素的实现可能会导致死锁,即调度器因为没有足够空间存储新的 Attention key 和 value 而无法为任何请求调度迭代。因此,ORCA 调度器需要知道为管理器预分配的内存区域的剩余大小。

调度算法细节。ORCA 调度器综合考虑了所有这些因素:它根据到达时间选择最多 "max batch size" 个请求,并在首次调度请求时为其预留足够的 key 和 value 存储空间。算法 1 描述了调度过程。Select 函数(第 17 行)根据请求的到达时间从请求池中选择最多 max_bs 个请求(第 20-22 行)。当调度器考虑一个处于初始化阶段的请求时,它会使用请求的 max_tokens 属性提前预留 max_tokens 个 slot 的 GPU 内存用于存储 key 和 value(第 23-26 行)。调度器根据当前已预留的 slot 数量 n_rsrv 和总 slot 数量 n_slots 来判断是否可以进行预留(第 25 行)。一个 slot 定义为存储单个 token 的 Attention key 和 value 所需的内存量。n_slots 是一个由操作员配置的参数。一旦预留成功,就保证了管理器在请求完成前都能为新生成的 key 和 value 分配缓冲区。与 max_bs 的调整需要权衡延迟和吞吐量不同,n_slots 的配置很简单:操作员只需在内存约束下使用尽可能大的 n_slots 值即可。

Algorithm 1 ORCA scheduling algorithm.

Params: n_workers: number of workers, max_bs:
max batch size, n_slots: number of K/V slots
1 n_scheduled ← 0
2 n_rsrv ← 0
3 while true do
4   batch, n_rsrv ← Select(request_pool, n_rsrv)
5   schedule engine to run one iteration of
    the model for the batch
6   foreach req in batch do
7     req.state ← RUNNING
8   n_scheduled ← n_scheduled + 1
9   if n_scheduled = n_workers then
10    wait for return of a scheduled batch
11    foreach req in the returned batch do
12      req.state ← INCREMENT
13      if finished(req) then
14        n_rsrv ← n_rsrv − req.max_tokens
15    n_scheduled ← n_scheduled − 1
16
17 def Select(pool, n_rsrv):
18   batch ← {}
19   pool ← {req | req ∈ pool, req.state != RUNNING}
20   SortByArrivalTime(pool)
21   foreach req in pool do
22     if batch.size() = max_bs then break
23     if req.state = INITIATION then
24       new_n_rsrv ← n_rsrv + req.max_tokens
25       if new_n_rsrv > n_slots then break
26       n_rsrv ← new_n_rsrv
27   batch ← batch ∪ {req}
28   return batch, n_rsrv

流水线并行。ORCA 的调度器使得引擎中的 worker 执行可以在多个批次之间形成流水线。调度器在当前已调度批次数 n_scheduled 达到 worker 数量 n_workers 之前,不会等待已调度批次的返回(算法 1 的第 9-10 行)。通过这种方式,调度器使引擎中并发运行的批次数保持为 n_workers,这意味着引擎中的每个 worker 都在处理一个批次而不会空闲。图 8a 描绘了 3 个 ORCA worker 的执行流水线。调度器首先选择请求 A 和 B 组成批次 AB 并调度,然后在 AB 返回之前继续注入批次 CD 和 EF。相比之下,由于请求级调度,Triton【索引7,NVIDIA Triton Inference Server】和 FasterTransformer【索引4,NVIDIA FasterTransformer】的组合在当前批次 AB 完成前无法注入下一个请求 C。为了实现流水线执行,FasterTransformer 将一个批次分割成多个微批次(microbatches)【索引28,GPipe: Efficient Training of Giant Neural Networks Using Pipeline Parallelism,2019,NeurIPS】并在这些微批次上流水线化分区。如图 8b 所示,这种方法可能因微批次较小而降低批处理的性能增益,或者因微批次过大而导致流水线气泡。ORCA 得益于迭代级调度,无需将批次分割成微批次即可轻松实现流水线,避免了这种权衡。

图 8: ORCA 和 FasterTransformer 中流水线并行使用的比较,其中 $X_i$ 是请求 X 的第 i 次迭代。(a) ORCA 执行流水线。(b) FasterTransformer 执行流水线。

5. 实现

技术栈与模型支持。我们使用 13000 行 C++ 代码实现了 ORCA,基于 CUDA 生态系统。我们使用 gRPC【索引2,gRPC】进行 ORCA 引擎控制平面的通信,而数据平面则使用 NCCL【索引5,NVIDIA NCCL】进行层间和层内通信。由于 ORCA 专注于基于 Transformer 的生成模型,它提供了流行的 Transformer 层作为模型构建块,包括原始的编码器-解码器 Transformer【索引60,Attention is All you Need,2017,NeurIPS】、GPT【索引47,Language Models are Unsupervised Multask Learners,2019】以及 Raffel 等人【索引49,Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer,2020,JMLR】讨论的其他变体。

核函数融合。我们还为 LayerNorm、Attention 和 GeLU 算子实现了融合核函数,这与其他用于 Transformer 模型训练或推理的系统【索引1,DeepSpeed;索引4,NVIDIA FasterTransformer;索引58,Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism,2019,arXiv】类似。例如,计算 Attention query 和 key 之间的点积、对点积应用 Softmax 以及对 Attention value 进行加权平均的过程被融合成一个单一的 CUDA 核函数。此外,我们更进一步,通过简单地连接用于不同请求的核函数的所有线程块,来融合分割后的 Attention 算子的核函数。尽管这种融合使得一个核函数内的线程块具有不同的特性和生命周期(这通常不被 CUDA 编程实践所鼓励),因为它们处理不同形状的张量,但我们发现这种融合通过提高 GPU 利用率和减少核函数启动开销【索引34,Nimble: Lightweight and Parallel GPU Task Scheduling for Deep Learning,2020,NeurIPS;索引39,Rammer: Enabling Holistic Deep Learning Compiler Optimizations with rTasks,2020】是有益的。

A4 实验环境

  • 硬件配置

    • 计算平台:Azure ND96asr A100 v4 虚拟机。
    • GPU:每台虚拟机配备 8 块 NVIDIA 40-GB A100 GPU,通过 NVLink 连接。实验最多使用 4 台虚拟机。
    • 网络:每台虚拟机配备 8 个 Mellanox 200Gbps HDR Infiniband 适配器,提供 1.6Tb/s 的虚拟机间互联带宽。
  • 模型与数据集

    • 模型架构:使用 GPT 模型作为代表,涵盖了多种参数规模:13B、101B、175B 和 341B。具体配置见表 1。13B 和 175B 的配置源自 GPT-3 论文【索引12,Language Models are Few-Shot Learners,2020,NeurIPS】,101B 和 341B 的配置是基于前两者调整得出的。
    • 模型参数:所有模型最大序列长度为 2048,使用 fp16 格式的模型参数和中间激活值。
    • 数据集:由于没有公开的生成式语言模型请求追踪数据,实验使用合成的客户端请求追踪。每个请求的输入 token 数量从 U(32, 512) 的均匀分布中采样,最大生成 token 数 max_gen_tokens 从 U(1, 128) 中采样。请求到达时间遵循泊松过程,通过调整泊松参数来模拟不同负载。
  • 软件配置

    • 实现:ORCA 使用 C++ 实现,基于 CUDA 生态系统。
    • 依赖库:控制平面通信使用 gRPC,数据平面通信(层内和层间)使用 NCCL。
    • 基线系统:与 NVIDIA 的 FasterTransformer【索引4,NVIDIA FasterTransformer】进行比较,这是一个支持大规模 Transformer 模型分布式执行的推理引擎。


表1:实验中使用的模型配置。

A4 实验结果

实验分为引擎微基准测试和端到端性能测试两部分。

引擎微基准测试

  • 实验内容:此测试旨在单独评估 ORCA 执行引擎的性能,不涉及迭代级调度。测试脚本模拟请求级调度的行为,对一批具有相同输入长度(32 或 128)和相同生成长度(32)的请求进行处理,并测量处理整个批次的时间。
  • 实验结果
    • 在单机(13B 和 101B 模型)上,ORCA 引擎的性能与 FasterTransformer 相当或略差(图 9a, 9b)。这是因为 ORCA 对 Attention 操作不进行批处理,但由于 Attention 没有模型参数,不批处理对效率影响较小。
    • ORCA 通过基于 max_tokens 的动态内存分配,可以支持比 FasterTransformer 更大的批次大小(图 9a, 9b)。
    • 在跨机(175B 模型)且禁用流水线的场景下,ORCA 引擎性能比 FasterTransformer 高出 47%(图 9c)。这主要归功于 ORCA 的控制-数据平面分离设计,避免了因控制消息交换引起的 CPU-GPU 同步开销。
    • 图9:使用FasterTransformer和ORCA引擎(无调度组件)处理一批请求的执行时间。标签“ft(n)”代表FasterTransformer处理具有n个输入token的请求的结果。导致内存不足错误的配置显示为缺失条目。
      图9:使用FasterTransformer和ORCA引擎(无调度组件)处理一批请求的执行时间。标签“ft(n)”代表FasterTransformer处理具有n个输入token的请求的结果。导致内存不足错误的配置显示为缺失条目。

端到端性能

  • 实验内容:使用合成的请求追踪来评估整个 ORCA 系统的端到端延迟和吞吐量,请求的输入和输出长度是异构的。对比 ORCA(不同 max_bs)和 FasterTransformer(不同 max_bsmbs 组合)。
  • 实验结果
    • 总体性能:ORCA 在延迟和吞吐量上显著优于 FasterTransformer。随着负载增加,ORCA 的吞吐量大幅提升,而延迟增长相对平缓。相比之下,FasterTransformer 无法有效处理异构请求,导致吞吐量受限且延迟高(图 10)。
    • 性能提升量化:在 175B 模型上,为了达到约 190ms 的中位数归一化延迟,FasterTransformer 的吞吐量为 0.185 req/s,而 ORCA 的吞吐量为 6.81 req/s,实现了 36.9 倍的提升(图 10b)。
    • 批次大小影响:对于 ORCA,增加 max_bs 能有效提高吞吐量而不显著影响延迟,这得益于迭代级调度解决了提前完成和延迟加入的问题。对于 FasterTransformer,增加 max_bs 并不总能提升吞吐量,因为它难以处理请求长度差异大的批次。
    • 图10:中位数端到端延迟(按生成token数归一化)和吞吐量。标签“orca(max_bs)”表示ORCA在最大批处理大小为max_bs时的结果。标签“ft(max_bs, mbs)”表示FasterTransformer在最大批处理大小为max_bs和微批处理大小为mbs时的结果。
      图10:中位数端到端延迟(按生成token数归一化)和吞吐量。标签“orca(max_bs)”表示ORCA在最大批处理大小为max_bs时的结果。标签“ft(max_bs, mbs)”表示FasterTransformer在最大批处理大小为max_bs和微批处理大小为mbs时的结果。
    • 同构请求追踪:当所有请求具有相同的输入和输出长度时,FasterTransformer 的性能随 max_bs 增加有明显改善,因为提前完成的问题不存在了。尽管如此,ORCA 仍然优于 FasterTransformer(max_bs=8 的情况),除非 ORCA 的 max_bs 设为 1,此时其退化为无批处理的简单流水线(图 11)。
    • 图11:使用175B模型和由同构请求组成的追踪数据得到的中位数端到端延迟和吞吐量。由于所有请求特性相同,延迟未进行归一化。
      图11:使用175B模型和由同构请求组成的追踪数据得到的中位数端到端延迟和吞吐量。由于所有请求特性相同,延迟未进行归一化。

A5 结论

本文提出了迭代级调度选择性批处理两种创新技术,以解决服务基于 Transformer 的生成式模型时面临的低延迟和高吞吐量挑战。迭代级调度将调度粒度从请求级细化到迭代级,允许请求灵活地加入和离开批次;选择性批处理则使得在迭代级调度的基础上,对任意组合的请求进行有效批处理成为可能。基于这些技术,我们设计并实现了一个名为 ORCA 的分布式服务系统。实验结果表明,ORCA 的方法非常有效,与当前最先进的系统相比,在同等延迟水平下,其吞吐量提升了一个数量级。未来的工作可以探索设计一个通用的接口,以支持这两种技术,同时保持服务层和执行层抽象的分离。

引用文献汇总

本文在缩写过程中引用了以下文献:

  • [2] gRPC:

    • 引用位置: 4.1 节,4.2 节,5 节
    • 原文描述: 用于 ORCA 引擎主节点和工作节点控制器之间的通信,因为它不涉及 GPU,避免了不必要的同步。
  • [3] Megatron-LM:

    • 引用位置: 2 节,4.1 节
    • 原文描述: 作为一种采用增量解码的 Transformer 系统,以及一种现有的分布式推理系统,但它在每次迭代中通过 GPU-to-GPU 通信交换控制消息,会引入同步开销。
  • [4] NVIDIA FasterTransformer:

    • 引用位置: 1 节,2 节,4.1 节,4.2 节,5 节,6 节
    • 原文描述: 作为当前主流的 Transformer 模型推理优化执行引擎,是本文的主要比较基线。它也采用增量解码和模型并行策略,但其请求级调度和基于微批次的流水线存在效率瓶颈。
  • [5] NVIDIA NCCL:

    • 引用位置: 4.1 节,5 节
    • 原文描述: 用于 ORCA 引擎中 GPU 之间中间张量数据的传输,即数据平面的通信。
  • [6] NVIDIA TensorRT:

    • 引用位置: 2 节
    • 原文描述: 作为底层模型执行引擎的一个例子。
  • [7] NVIDIA Triton Inference Server:

    • 引用位置: 1 节,2 节,4.2 节
    • 原文描述: 作为现有的主流 ML 推理服务系统,它采用请求级调度,是导致生成式模型服务效率低下的原因之一。
  • [12] Language Models are Few-Shot Learners (GPT-3):

    • 引用位置: 1 节,4.1 节,6 节
    • 原文描述: 作为大型生成式语言模型的代表,其成功凸显了为这类模型提供高效服务系统的必要性。实验中使用的 175B 模型配置来源于此。
  • [14] TVM: An Automated End-to-End Optimizing Compiler for Deep Learning:

    • 引用位置: 2 节
    • 原文描述: 作为底层模型执行引擎的一个例子。
  • [23] Low Latency RNN Inference with Cellular Batching (BatchMaker):

    • 引用位置: 7 节
    • 原文描述: 作为相关工作,BatchMaker 提出了针对 RNN 的单元格级批处理,但该方法不适用于 Transformer,因为 Transformer 中每个 token 位置的计算单元都是不同的。
  • [25] Long Short-Term Memory (LSTM):

    • 引用位置: 2 节
    • 原文描述: 与 Transformer 对比,说明 LSTM 的内部状态大小是固定的,而 Transformer 的状态大小随序列长度增长。
  • [28] GPipe: Efficient Training of Giant Neural Networks Using Pipeline Parallelism:

    • 引用位置: 4.2 节
    • 原文描述: FasterTransformer 使用其提出的微批次(microbatch)技术来实现流水线并行。
  • [31] Scaling Laws for Neural Language Models:

    • 引用位置: 4.1 节
    • 原文描述: 证明了扩展语言模型可以提升模型质量,强调了支持大模型服务的重要性。
  • [34] Nimble: Lightweight and Parallel GPU Task Scheduling for Deep Learning:

    • 引用位置: 5 节
    • 原文描述: 引用其关于减少 GPU 核函数启动开销的观点,来支持 ORCA 的核函数融合策略。
  • [39] Rammer: Enabling Holistic Deep Learning Compiler Optimizations with rTasks:

    • 引用位置: 5 节
    • 原文描述: 引用其关于减少 GPU 核函数启动开销的观点,来支持 ORCA 的核函数融合策略。
  • [42] TensorFlow-Serving: Flexible, High-Performance ML Serving:

    • 引用位置: 1 节,2 节
    • 原文描述: 作为现有的主流 ML 推理服务系统的一个例子。
  • [43] fairseq: A Fast, Extensible Toolkit for Sequence Modeling:

    • 引用位置: 2 节
    • 原文描述: 提出了增量解码(incremental decoding)技术,通过缓存 Attention 的 key 和 value 来避免重复计算,是 Transformer 推理优化的关键技术。
  • [47] Language Models are Unsupervised Multitask Learners (GPT-2):

    • 引用位置: 2 节,5 节
    • 原文描述: 作为 Transformer 生成模型的代表,ORCA 支持其架构。
  • [49] Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer:

    • 引用位置: 5 节
    • 原文描述: ORCA 支持其论文中讨论的 Transformer 变体。
  • [55] Mesh-TensorFlow: Deep Learning for Supercomputers:

    • 引用位置: 1 节,4.1 节
    • 原文描述: 提出了层内模型并行(intra-layer model parallelism)技术,被 ORCA 采用以支持大模型。
  • [58] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism:

    • 引用位置: 1 节,4.1 节,5 节
    • 原文描述: 提出了层内模型并行技术,被 ORCA 采用。ORCA 的融合核函数也借鉴了此类系统的做法。
  • [60] Attention is All you Need (Transformer):

    • 引用位置: 1 节,2 节,5 节
    • 原文描述: 提出了 Transformer 架构及其核心的 Attention 机制,是本文所研究模型的基础。
  • [64] A Learning Algorithm for Continually Running Fully Recurrent Neural Networks:

    • 引用位置: 3 节
    • 原文描述: 引用其提出的 teacher forcing 技术,来说明为什么模型训练时不存在请求级调度导致的问题。