DeepSpeed 推理：使用自定义推理内核和量化支持的多GPU推理
虽然DeepSpeed支持训练先进的大规模模型，但由于现有推理解决方案中的三大限制，在所需的应用场景中使用这些训练好的模型仍然具有挑战性：1）缺乏对多GPU推理的支持，以适应大型模型并满足延迟要求，2）在使用小批量大小运行推理时，GPU内核性能有限，3）在利用量化方面存在困难，包括量化模型以减少模型大小和延迟，以及在没有专用硬件的情况下支持量化模型的高性能推理。
为了应对这些挑战，我们引入了DeepSpeed Inference，它无缝地为在DeepSpeed中训练的大型模型添加了高性能推理支持，具有三个关键特性：适用于多GPU推理的推理适应并行性、针对小批量大小优化的推理优化内核，以及对量化模型的量化感知训练和推理内核的灵活支持。
使用自适应并行性进行多GPU推理
并行化是适应大型模型并减少训练和推理过程中每个设备内存消耗的有效方法。然而，简单地将训练并行化选择和程度应用于推理效果不佳。MP和PP配置通常在模型训练期间设置，除了数据并行（DP）之外，基于内存占用和计算风格以及资源预算。一方面，推理计算本质上需要较少的内存，因此每个设备可以承担更大的分区。这有助于减少模型部署所需的并行化程度。另一方面，优化延迟或满足延迟要求通常是推理中的首要任务，而训练则优化吞吐量。
为了获得所需的延迟，DeepSpeed Inference 自动采用 MP 作为一种有效的方法来减少模型延迟，并且其并行度通常首先确定。通过 MP，我们可以分割模型并在多个设备（GPU）上并行化计算操作以减少延迟，但这会降低计算粒度并增加可能影响吞吐量的通信。一旦达到延迟目标，DeepSpeed 可以应用管道并行性以最大化吞吐量。总体而言，DeepSpeed Inference 支持从训练到推理的并行方法和度选择的灵活适应，在最小化延迟的同时节省部署成本。
定制推理内核以提高Transformer模块的计算效率
为了实现高计算效率，DeepSpeed-inference 提供了专门为 Transformer 块设计的推理内核，通过操作符融合，并考虑多 GPU 的模型并行性。我们的内核融合方案与类似方法的主要区别在于，我们不仅融合了元素级操作（如偏置加法、残差和激活函数），还将通用矩阵乘法（GeMM）操作与其他操作合并。为此，我们设计了一种高效的向量-矩阵或瘦矩阵-矩阵乘法实现，使我们能够在 GeMM 操作的归约边界处融合更多操作。
内核融合
我们采取两种主要策略来融合操作：1）在整个融合操作序列中保持输入和输出的访问模式不变；2）在每个全归约边界处融合操作。第一种策略确保不同的线程块不会在流式多处理器（SMs）之间传输数据。这是因为除了使用主内存外，SMs之间没有直接的通信方式，而主内存由于内存访问的非确定性行为会增加块同步的开销。第二种策略背后的原因是我们无法继续执行，除非在模型并行的GPU之间对部分结果进行归约。
图1：带有Megatron风格模型并行全归约组件的Transformer层。图中用虚线（线的宽度表示融合深度）展示了层中融合在一起的部分。
图1展示了Transformer层的不同组件，以及我们在推理优化中考虑进行融合的操作组。我们还考虑了NVIDIA Megatron-LM风格的并行性，该并行性将注意力（Attn）和前馈（FF）块分布在多个GPU上。因此，我们包括了在Attn和FF块之后在并行GPU之间减少结果的两个all-reduce操作。如图1所示，我们在Transformer层内的四个主要区域融合了操作：
- 输入层归一化加上查询、键和值的GeMMs及其偏置添加。
- 变换加注意力。
- 中间FF、Layer-Norm、Bias-add、残差和高斯误差线性单元（GELU）。
- 偏置加残差。
为了融合这些操作，我们利用共享内存作为中间缓存，用于在层归一化和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量化工具包，支持灵活的量化感知训练和用于量化推理的高性能内核。
在训练过程中，我们引入了一种名为量化混合（Mixture of Quantization, MoQ）的新方法，该方法受到混合精度训练的启发，同时无缝应用量化。通过MoQ，我们可以在训练过程中更新参数时模拟量化的影响，从而控制模型的精度。此外，它支持灵活的量化策略和调度——我们发现，通过在训练过程中动态调整量化位数，最终量化模型在相同压缩比下提供了更高的精度。为了适应不同的任务，MoQ还可以利用模型的二阶信息来检测其对精度的敏感性，并相应地调整量化调度和目标。
为了最大化量化模型的性能提升，我们提供了专为量化模型设计的推理内核，通过优化数据移动来减少延迟，但不需要专门的硬件。最后，我们的工具包不需要在客户端进行任何代码更改，使其易于使用。
性能结果
提升吞吐量并降低推理成本。图3显示了对应于三个Transformer网络（GPT-2、Turing-NLG和GPT-3）的三种模型大小的每GPU推理吞吐量。当使用与基线相同的FP16精度时，DeepSpeed Inference的每GPU吞吐量提高了2到4倍。通过启用量化，我们进一步提高了吞吐量。对于GPT-2，我们实现了3倍的吞吐量提升；对于Turing-NLG，实现了5倍的提升；对于与GPT-3在特性和大小上相似的模型，实现了3倍的提升，这直接转化为在服务这些大型模型时3-5倍的推理成本降低。此外，我们在不牺牲延迟的情况下实现了这些吞吐量和成本的改进，如图5所示。
图3：不同模型大小的推理吞吐量。DeepSpeed推理比基线实现了3倍到5倍的吞吐量提升。
降低推理成本的一个来源是通过减少用于托管大型模型的GPU数量，如图4所示。优化的GPU资源来自1）使用适应推理的并行性，允许用户从训练好的模型检查点调整模型和管道并行度，以及2）通过INT8量化将模型内存占用减少一半。如图所示，我们通过调整并行性，使用2倍更少的GPU来运行17B模型大小的推理。通过DeepSpeed MoQ的INT8量化，我们分别使用4倍和2倍更少的GPU来处理17B和175B大小的模型。
图4：用于在不同模型大小上运行推理的GPU数量如图4所示。
减少推理延迟。对于推理延迟至关重要的应用场景，我们可以增加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倍的加速。
图5. 使用不同并行配置优化延迟的17B模型的推理延迟。