DeepSpeed 推理:具有定制推理内核和量化支持的多 GPU 推理

虽然 DeepSpeed 支持训练高级的大规模模型,但在现有推理解决方案中,在所需应用场景中使用这些训练模型仍然具有挑战性,主要原因是存在以下三个局限性:1) 缺乏对多 GPU 推理的支持,无法适应大型模型并满足延迟要求;2) 使用小批量大小运行推理时,GPU 内核性能有限;3) 量化利用的困难,包括量化模型以减小模型大小和延迟,以及支持量化模型的高性能推理,而无需专门的硬件。

为了应对这些挑战,我们引入了 DeepSpeed 推理,它以三种关键特性无缝地为在 DeepSpeed 中训练的大型模型添加了高性能推理支持:推理自适应并行用于多 GPU 推理、针对小批量大小进行优化的推理优化内核,以及对量化感知训练和量化模型的推理内核的灵活支持。

使用自适应并行的多 GPU 推理

并行是一种有效的方案,可以适应大型模型并减少训练和推理的每个设备内存消耗。但是,简单地将训练并行选择和程度应用于推理效果不佳。MP 和 PP 配置通常在模型训练期间设置,除了数据并行 (DP) 之外,还根据内存占用、计算方式和资源预算进行设置。一方面,推理计算本质上需要较少的内存,因此可以负担每个设备更大的分区,这有助于减少模型部署所需的并行程度。另一方面,优化延迟或满足延迟要求通常是推理中的首要任务,而训练优化的是吞吐量。

为了获得所需的延迟,DeepSpeed 推理自动调整 MP 作为一种有效的方案来减少模型延迟,并且其并行程度通常首先确定。使用 MP,我们可以拆分模式并在多个设备(GPU)上并行化计算操作以减少延迟,但这会降低计算粒度并增加可能损害吞吐量的通信。一旦满足延迟目标,DeepSpeed 就可以应用管道并行来最大化吞吐量。总的来说,DeepSpeed 推理支持从训练到推理的并行方法和程度选择的灵活适应,最大程度地减少延迟并节省部署成本。

定制推理内核,用于提升 Transformer 模块的计算效率

为了实现高计算效率,DeepSpeed 推理提供针对 Transformer 模块量身定制的推理内核,通过操作融合考虑了模型并行用于多 GPU。我们的内核融合方案与类似方法的主要区别在于,我们不仅融合了逐元素操作(如偏差加、残差和激活函数),还将通用矩阵乘法 (GeMM) 操作与其他操作合并。为此,我们为向量矩阵或瘦矩阵矩阵乘法设计了一种高效的实现,使我们能够在 GeMM 操作的缩减边界处融合更多操作。

内核融合

我们采取两种主要策略来融合操作:1) 在一起融合的操作序列中保持输入和输出的访问模式完整;2) 在每个全部减少边界处融合操作。第一个策略确保不同的线程块不会遇到在流多处理器 (SM) 之间传输数据的情况。这是由于 SM 之间没有直接的通信,除了使用主内存外,这会增加块同步开销,因为内存访问的行为是不确定的。第二个策略背后的原因是,除非部分结果在模型并行 GPU 之间减少,否则我们无法继续执行。

Inference-Kernel-Fusion

图 1:具有 Megatron 样式模型并行全部减少组件的 Transformer 层。该图说明了使用虚线融合在一起的层的部分(线的宽度显示融合深度)。

图 1 显示了 Transformer 层的不同组件,以及在我们的推理优化中考虑的融合操作组。我们还考虑了 NVIDIA Megatron-LM 样式的并行性,它将注意力 (Attn) 和前馈 (FF) 块分隔在多个 GPU 上。因此,我们包括了两个全部减少操作,这两个操作在 Attn 和 FF 块之后在并行 GPU 之间减少结果。如图 1 所示,我们在四个主要区域融合 Transformer 层内的操作

  1. 输入层归一化加上查询、键和值的 GeMM 以及它们的偏差添加。
  2. 变换加上注意力。
  3. 中间 FF、层归一化、偏差添加、残差和高斯误差线性单元 (GELU)。
  4. 偏差添加加上残差。

为了融合这些操作,我们利用共享内存作为中间缓存,用于在层归一化和 GeMM 中使用的缩减操作以及逐元素操作之间传输数据。此外,我们在缩减部分计算时使用 warp 级指令在线程之间通信数据。此外,我们使用了一个新的 GeMM 操作调度,它允许融合第三个内核融合所需的尽可能多的操作。我们还通过使用隐式矩阵转换来减少内存压力,将第二个内核融合中注意计算的 GeMM 组合在一起。与使用 cuBLAS GeMM 的未融合计算方式相比,我们分别将所有这些内核融合的性能提高了 1.5 倍、2.9 倍、3 倍和 1.2 倍。

从训练到推理的无缝管道,支持自动内核注入

要以推理模式运行模型,DeepSpeed 只需要模型检查点的位置和所需的并行配置,即 MP/PP 程度。DeepSpeed 推理内核也可以为许多知名模型架构启用,例如 HuggingFace(Bert 和 GPT-2)或 Megatron GPT 基于的模型,使用预定义的策略映射,该映射将原始参数映射到推理内核中的参数。对于其他基于 Transformer 的模型,用户可以指定他们自己的策略映射。请注意,DS-Inference 可以独立于训练管道运行,只要它接收所有模型检查点,并且如果定义了正确的映射策略,DeepSpeed Transformer 推理内核可以注入到任何 Transformer 模型中。有关如何启用 Transformer 推理内核以及指定并行的更多信息,请参阅我们的 推理教程

灵活的量化支持

为了进一步降低大型模型的推理成本,我们创建了 DeepSpeed 量化工具包,支持灵活的量化感知训练和用于量化推理的高性能内核。

对于训练,我们引入了一种名为混合量化 (MoQ) 的新方法,该方法的灵感来自于混合精度训练,同时无缝地应用量化。使用 MoQ,我们可以通过模拟量化对训练的每个步骤中更新参数的影响来控制模型的精度。此外,它支持灵活的量化策略和时间表——我们发现,通过在训练期间动态调整量化位数,最终的量化模型在相同的压缩比下提供了更高的精度。为了适应不同的任务,MoQ 还可以利用模型的二阶信息来检测它们对精度的敏感性,并相应地调整量化时间表和目标。

为了最大程度地提高量化模型的性能提升,我们提供了专为量化模型量身定制的推理内核,这些内核通过优化数据移动来减少延迟,但不需要专门的硬件。最后,我们的工具包不需要在客户端进行任何代码更改,使其易于使用。

性能结果

提高吞吐量并降低推理成本。图 3 显示了三个 Transformer 网络(GPT-2、Turing-NLG 和 GPT-3)的三个模型大小的每个 GPU 的推理吞吐量。DeepSpeed 推理在使用与基线相同的 FP16 精度时,每个 GPU 的吞吐量提高了 2 到 4 倍。通过启用量化,我们进一步提高了吞吐量。对于 GPT-2,我们实现了 3 倍的吞吐量提升,对于 Turing-NLG 实现了 5 倍的吞吐量提升,对于一个在特性和大小上类似于 GPT-3 的模型,我们实现了 3 倍的吞吐量提升,这直接转化为在服务这些大型模型时 3-5 倍的推理成本降低。此外,我们实现了这些吞吐量和成本改进,而不会影响延迟,如图 5 所示。

Inference-Throughput

图 3:不同模型大小的推理吞吐量。DeepSpeed 推理比基线实现了 3 倍到 5 倍更高的吞吐量。

推理成本降低的一个来源是通过减少托管大型模型的 GPU 数量,如图 4 所示。优化的 GPU 资源来自 1) 使用推理自适应并行,允许用户从训练的模型检查点调整模型和管道并行程度,以及 2) 使用 INT8 量化将模型内存占用减少一半。如图所示,我们使用少 2 倍的 GPU 来为 17B 模型大小运行推理,方法是调整并行性。结合通过 DeepSpeed MoQ 进行的 INT8 量化,我们分别为 17B 和 175B 大小使用了少 4 倍和少 2 倍的 GPU。

Inference-Throughput

图 4:用于在图 4 中所示的不同模型大小上运行推理的 GPU 数量。

减少推理延迟。对于推理延迟至关重要的应用场景,我们可以在 DeepSpeed 推理中提高模型并行程度,以进一步减少推理延迟。如图 5 所示,当我们将模型并行大小增加到 4 时,与 PyTorch 相比,我们可以将延迟降低 2.3 倍。此外,通过在推理时调整并行性并使用 MoQ 量化模型,我们仍然可以使用更少的 GPU 来实现高延迟改进。我们在使用基线资源的 4 倍和 2 倍的情况下,分别获得了 1.3 倍和 1.9 倍的加速。

对于推理延迟至关重要的应用场景,我们可以在 DeepSpeed 推理中提高模型并行程度,以进一步减少推理延迟。如图 5 所示,当我们将模型并行大小增加到 4 时,与 PyTorch 相比,我们可以将延迟降低 2.3 倍。此外,通过在推理时调整并行性并使用 MoQ 量化模型,我们仍然可以使用更少的 GPU 来实现高延迟改进。我们在使用基线资源的 4 倍和 2 倍的情况下,分别获得了 1.3 倍和 1.9 倍的加速。

Inference-Throughput

图 5. 使用不同的并行配置优化延迟的 17B 模型的推理延迟。

更新: