微软新的异构分布式推理技术，让transformers的性能提升数倍！

DeepSpeed是微软开源的一个大规模深度学习优化库，它最大的目标是让分布式训练变得更加简单易用。DeepSpeed包含了很多微软自研的先进的分布式训练技术，最近6月份他们发布了一篇最新的论文，提出了ZeroQuant这个模型压缩的方法和技术，并且已经在最新的DeepSpeed版本中使用了。

关于DeepSpeed详情，可以参考：https://www.datalearner.com/ai-tools/deep-speed

时隔一个月不到，DeepSpeed再次发布新论文：DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale，介绍了他们最新的DeepSpeed推理系统，相比较传统的推理系统，其最大的特点是可以利用异构硬件进行推理加速，同时利用CPU与GPU，其结果表明比单纯使用GPU的效率提升很多。具体来说，这个论文的贡献是：

1. 多GPU推理解决方案，在适合聚合GPU内存时，可最大限度地减少延迟，同时最大限度地提高密集和稀疏变压器模型的吞吐量；
2. 异构推理解决方案，除GPU内存和计算外，还利用CPU和NVMe内存，以实现不适合聚合GPU内存的大型模型的高推理吞吐量。对于面向延迟的场景，深度速度推理将延迟降低了最多7.3倍，对于面向吞吐量的场景，吞吐量提高了1.5倍以上。
3. 它通过利用数百个GPU，在实时延迟约束下实现万亿参数规模推理，这是前所未有的推理规模。它可以推断出比仅使用GPU解决方案大25倍的型号，同时提供84 TFLOPS的高吞吐量（超过A6000峰值的50%）。

根据论文的描述，高性能的推理的主要挑战体现在：

1、不同batch size大小的推理挑战

小批量性能受到读取模型权重中内存带宽利用率的限制。在小批量推理中优化内存带宽有三个主要挑战。首先，由于不同内核使用小批量执行Transformers层操作的工作有限，推理性能会受到内核调用开销的影响。其次，每次内核调用都会将数据写入全局内存，在下一次内核调用期间由GPU内核读取，GPU内核和全局内存之间的数据传输会增加额外的开销。最后，无论是cBLAS还是CTLAS GeMM库都没有为极小的批处理大小进行很好的调整，也无法实现良好的内存带宽利用率。另一方面，批量推理性能受到计算利用率的限制，虽然在Transformers层内像GeMM这样的计算繁重操作可以使用CUBLAS和CUTLAS库实现非常好的计算利用率，总体利用率仍然会受到内核启动开销和GPU内核和除GeMM以外的不同内核之间的全局内存之间的数据传输的限制。为了解决这些挑战，我们引入了两种技术：i)深度融合，通过在元素操作之外融合多个内核来减少内核调用和数据移动开销；ii)定制GeMM内核，旨在在批处理大小相对较小时提高内存带宽利用率，同时允许使用深度融合进行融合。接下来，我们将详细讨论这些技术。

2、深度融合的挑战

虽然算子融合是深度学习中用于减少内核启动和数据移动开销的常见技术，但它主要限于元素运算符。相比之下，Transformers由数据布局转换、缩减和GeMM等运算符组成，这些运算符在线程块之间创建数据依赖关系，使它们难以融合。这是因为在GPU上，如果线程块产生的数据被另一个线程块消耗，则需要全局内存同步，该同步调用新的内核。为了避免全局同步的需要，深度融合沿着迭代空间的维度平铺计算空间，而迭代空间不会产生跨平铺数据依赖关系，并在不同的线程块上并行执行它们。包含数据依赖关系的计算空间的维度不平铺，而是由同一线程块处理。在此平铺之后，如果第二个运算符的每个平铺正好依赖于第一个运算符的一个输出平铺，则可以使用DeepFusion融合两个运算符。通过在粒度上执行融合，深度融合不仅可以融合元素操作，还可以融合缩减、数据转置和GeMM，只要不存在跨磁贴依赖关系。例如，层规范中的所有微操作都可以沿着令牌维度平铺，而缩减维度则在平铺中处理。这允许层规范中的所有微操作融合到单个内核中，尽管由多个约简操作组成。此外，每个磁贴产生的数据要么保存在寄存器中，要么在可能的情况下保存在共享内存中，以允许跨运算符重用数据，而不会产生全局内存数据传输开销。

3、SBI-GeMM：适用于小批量的自定义GeMM

我们的定制GeMM实施旨在与Deep-Fusion融合，同时实现最大的内存带宽利用率。它的设计可以分为三个部分：平铺策略、协作组减少和数据布局转换，以提高内存带宽利用率。1)平铺策略：图1(a)描述了我们针对瘦矩阵乘法的GeMM调度。我们首先沿着输出维度平铺计算。这允许我们通过将缩减保持在磁贴中，使用单个内核来实现GeMM。对于输出维度太小，无法创建足够的并行磁贴以实现良好的内存带宽的小型模型，我们也将输入维度平铺，并将GeMM作为两个内核实现，以允许跨磁贴减少。

2)协作组减少：使用上述平铺策略，线程块中的每个扭曲负责为输出平铺产生部分减少的结果，并且需要在线程块中的所有扭曲中进行最终减少。通常，这是作为基于二叉树的共享内存减少实现的，这需要多个扭曲级同步，从而造成性能瓶颈。为了避免这种情况，我们在共享内存中执行单个数据布局转置，使得同一输出元素的部分结果在内存中是连续的，并且可以通过直接在寄存器中使用协作组集合的单个扭曲来减少（见图1(a）)。

最后，每个扭曲的第一个线程保存最终结果，并将其写入共享内存。共享内存中的结果是连续的，允许合并写入全局内存。3)利用完整缓存线：在GPU体系结构中，每个L1缓存线为128字节，但是在扭曲中每个线程使用单个FP16或INT8元素的合并内存访问不能完全消耗完整缓存线。沿输出维度读取每个线程的多个元素以解决此问题，可以减少并行磁贴的数量，这也会损害内存带宽。因此，我们的解决方案是在初始化期间转置权重矩阵，使每列的M行在内存中连续，允许每个线程沿着输入维度读取M个元素（见图1(b）)。考虑到128字节缓存线，我们将M设置为2，对于INT8数据类型，将M设置为4。

4、把它放在一起

小批量Transformers内核：图1.c显示了Transformers层的不同组件，以及在小批量推理情况下深度融合考虑的操作。如图所示，我们在四个主要区域融合变压器层内的操作：1)QKV GeMM和输入层范数，2)转置加注意，3)后注意层范数和中间GeMM,4)偏置和剩余加法。为了支持GeMM与单个内核中的其余操作的融合，我们在SM之间广播输入批处理，并执行GeMM之前的相同操作，这样就不需要在SM之间通信数据来添加GeMM计划。我们观察到，尽管在SM之间复制工作，但与非常小的批处理大小的非复制、非融合内核实现相比，我们仍然获得了性能优势。批量变压器内核：我们遵循与上面讨论的相同的融合策略，不同的是，我们使用CUBLAS进行GeMM操作，并保持它们不融合。支持不同数据类型：我们的内核支持GeMM操作的FP32、FP16和INT8数据类型。为了支持INT-8，我们使用针对不同批处理大小进行调整的切割[34] INT8 GeMM实现。我们还在GeMM之前添加了量化操作，我们使用DeepFusion进行融合，在GeMM之后添加了反量化操作，我们使用剪切激光的结语功能进行融合。通过Cuda-Graph消除内核调用开销：对于批量大小较小的中小型模型，随着我们减少内核的实际执行时间，主要延迟瓶颈从内核执行转移到CPU侧的内核启动开销。为了解决这个问题，我们在推理管道中添加了CUDA-Graph [35]支持。更具体地说，我们在推理时的正向计算期间首次启动内核时存储内核的跟踪，并创建可用于以下请求的计算图，这在很大程度上消除了内核启动开销，并大幅提高了性能。