在昨天举行的“2022百度云智能峰会·智能计算峰会”上，百度智能云发布了中国首个全栈自研AI基础设施“百度AI大基地”，引发受到业界的广泛关注。其中，百度白鸽·AI异构计算平台是AI IaaS层的重要组成部分，是AI基础的坚实基础设施。百度百格·AI异构计算平台集成了百度自研AI芯片“昆仑芯”，在AI计算、存储、加速、容器等方面对系统进行优化，最大限度地发挥算力的性能和性能。效率。

本次峰会上，NVIDIAGPU计算专家陈宇分享了题为“视觉大模型训练与推理优化”的演讲，基于Swin Transf以为基础以ormer模型为例，介绍了NVIDIA在视觉大模型训练和推理优化方面的探索和实践经验，包括Swin Transformer在训练和推理中使用的加速技术以及相应的加速效果。这些加速方法现已应用于百度白鸽·AI异构计算平台的AIAK加速功能，帮助各行业客户加速AI业务落地。

以下为演讲摘要。

大型视觉模型介绍：游泳变形金刚

在介绍Swin Transformer之前，我们可以先回顾一下Vision Transformer（ViT）的模型结构。如左图所示，ViT会将一张图像划分为一系列图像块（patch），然后通过基于Transformer的encoder对这一系列图像块进行编码，最终得到结果可用于分类特征等任务。

至于Swin Transformer，它在tenti上引入了窗口的概念，这与ViT对整个图像进行注意力计算（attention）不同。 Swin Transformer首先会将图像划分为若干个窗口，然后只对窗口内部的图像块进行相应的注意力计算，从而减少计算量。为了弥补窗口带来的边界问题，Swin Transformer进一步引入了窗口移位操作。同时，为了让模型拥有更丰富的位置信息，Swin Transformer还在注意力计算中引入了相对位置偏差。

在Swin Transformer的网络结构中，我们还可以看到Swin Transformer被分为多个阶段（stages），其中每个阶段的分辨率不同，从而形成了一个分辨率金字塔，这也逐渐降低了每个阶段的计算复杂度。每级对应几个变压器块。

接下来我们从具体操作的角度来划分Swin Transformer的操作符。可以看出，变压器块包括三个部分。第一部分是窗口平移、分区、反转等一些与窗口相关的操作，第二部分是注意力计算，第三部分是FFN计算。 Attention和FFN部分可以进一步细分为几个算子。最后如图所示，我们可以将整个模型细分为几十个算子的组合。这样的算子划分对于我们进行性能分析、定位性能瓶颈、进行加速优化非常重要。

大型视觉模型Swim Transformer的训练加速

接下来我们将介绍NVIDIA针对Swim Transformer模型训练加速的优化工作。对于Swin Transformer，我们一般使用数据并行进行多GPU训练。我们发现卡间通信的开销会比较小，所以这里我们优先优化单GPU上的计算瓶颈。

首先我们通过nsight系统性能分析工具看一下整个基线的性能。左图为FP32的基线。可以看到它的GPU利用率非常高，占比最高的是与矩阵乘法相关的内核。

所以对于矩阵乘法，我们的优化方法之一就是充分利用张量核来进行加速。可以考虑直接使用tf32张量核心或使用混合精度的fp16张量核心。 这里，我们采用混合精度解决方案。正如右图所示，矩阵乘法的内核由于使用了fp16张量内核，显着降低了其耗时比例。另外，利用混合精度，我们还可以在内核io上取得一定的增益。 总体而言，在 Swin-Large 模型上，我们可以通过混合精度实现 1.63 倍的吞吐量提升。

下一个优化解决方案是算子融合。 算子融合一般可以给我们带来两个好处：一是降低内核启动的成本。如左图所示，两个cuda内核的执行需要两次启动，这可能会导致内核之间存在一定的间隙，从而使GPU空闲。

如果我们将两个 cuda 内核合并为一个 cuda 内核，我们就可以节省一次启动并避免出现间隙。另一个好处是减少了全局内存访问，因为全局内存访问非常耗时。为了在两个独立的 cuda 内核之间传输结果，我们需要使用全局内存将两个 cuda 内核合并为一个内核。我们可以将结果转移到寄存器或者共享内存中，这样就避免了全局内存的读写，从而提高性能。

对于算子融合，我们的第一步是使用现成的 apex 库来融合 Layernorm 和 Adam 中的操作。可以看到，通过简单的指令替换，我们可以启用apex的fused Layernorm和fused Adam，从而将加速效果从1.63倍提升到2.11倍。

除了使用现有的apex库之外，NVIDIA还开发了手动融合算子。这里展示的是我们针对Swin Transformer特有的窗口相关操作开发的融合算子，比如窗口分割、移位、合并等。可以看出，通过引入我们定制的融合算子，加速比可以进一步提升至2.19倍。

接下来展示的是NVIDIA 在 mha 部分的融合工作。 Mha部分是Transformer模型中的一个大模块，因此对其进行优化往往可以带来更大的加速效果。从图中可以看出，在进行算子融合之前，mha部分的运算比例为37.69%，其中包含了很多elementwise kernel。通过将相关操作融合到单独的 fmha 内核中，我们进一步将加速比提高了 2.85 倍。

以上是我们在单卡上实现的训练加速效果。我们看一下单机8卡的训练情况。我们可以看到，通过上述优化，我们可以将训练吞吐量从 1,612 提高到 3,733，实现了 2.32 倍的加速。

大型视觉模型Swim Transformer的推理加速

接下来我们将介绍NVIDIA对Swim Transformer模型推理加速的优化工作。

和训练一样，推理的加速离不开算子融合的解决。不过，与训练相比，推理中的算子融合有两个优点：一是推理中的算子融合不需要考虑反向，因此内核开发过程中不需要考虑保存计算梯度所需的中间结果；其次，推理过程允许进行预处理。我们可以把一些只需要一次计算、可以重复使用的运算，提前计算好，保留结果，每次推理时直接调用，避免重复计算。

在推理方面，NVIDIA进行了很多算子融合，这里我们选取了两个加速效果明显的算子进行介绍。

第一个是算子融合 mha 部分。我们将位置偏差查找操作提前到预处理部分，以避免每次推理时都执行查找，然后将批处理 gemm、softmax 和批处理 gemm 合并到独立的 fmha 内核中。可以看到，集成后，这部分实现了10x的加速，端到端的加速也实现了1.58x。

另一种算子融合是QKV gemm + 偏差的融合。 gemm和bias的融合是一种非常常见的融合方法。这里，为了配合我们前面提到的fmha内核，我们需要提前变换一下weight和bias的格式。这种提前变换也体现了上述算子融合在推理中的灵活性。最后，通过QKV gemm+bias的融合，我们可以进一步实现1.1倍的端到端加速。

除了算子融合之外，我们还使用了矩阵乘法填充的优化方法。 在Swin Transformer的计算中，有时我们会遇到奇数主维度的矩阵乘法。此时，不利于我们的矩阵乘法内核进行向量化读写，从而使得内核的运行效率较低。这时我们可以考虑对参与运算的矩阵的主要维度进行padding，使其成为8的倍数。这样就可以将矩阵乘法内核向量化，alignment=8，一次读写8个元素。读写，提高性能。

如下表所示，将 n 从 49 填充到 56 后，矩阵乘法的延迟从 60.54us（微秒）下降到 40.38us，实现了 1.5 倍的加速。

下一个优化方法是使用half2或char4等数据类型。

以half2为例，它可以带来以下三个好处：

首先，矢量化读写可以提高显存的带宽利用率，减少内存访问指令的数量。如右图所示，通过half2的使用，内存访问指令减少了一半，内存的SOL也明显提升。

其次，结合 half2 专有的高吞吐量数学指令，可以减少内核延迟。

第三，在进行归约时，使用half2数据类型意味着一个cuda线程同时处理两个元素，可以有效减少空闲线程的数量。最后，通过使用half2，我们可以观察到addBiasResidual内核的延迟从20.96us下降到10.78us，加速了1.94倍。

接下来的优化方法就是巧用寄存器数组了。 当我们优化Transformer模型的常用算子（例如layernorm或softmax）时，我们经常需要在内核中多次使用相同的输入数据。因此，我们可以不用每次都从全局内存中读取数据，而是使用寄存器数组来缓存数据，以避免重复读取全局内存。

由于寄存器是每个CUDA线程独享的，因此在设计内核时，我们需要提前设置每个CUDA线程需要缓存的元素数量，从而开辟相应大小的寄存器数组并分配给每个CUDA线程。当涉及到我们负责的元素时，我们需要确保我们可以合并访问。

如右图所示，当我们有8个线程时，线程0可以处理元素0，当我们有4个线程时，线程0可以处理元素0和4，以此类推。我们一般建议使用模板函数，通过模板参数来控制每个cuda线程的寄存器数组大小。

另外，在使用寄存器数组时，我们需要保证我们的下标是常量。如果使用循环变量作为下标，我们应该尽量确保可以进行循环展开，以避免编译器将数据高延迟地放置到本地内存中。 ，如图所示，我们在循环条件中添加了限制，通过ncu报告可以看到，避免了本地内存的使用。

最后要介绍的优化方法是INT8量化。 INT8量化是推理加速非常重要的加速方法。对于基于Transformer的模型，INT8量化往往可以带来良好的加速效果。

对于Swin Transformer来说，通过结合合适的PTQ或QAT量化方案，可以在保证量化精度的同时达到良好的加速效果。通过QAT，我们可以保证精度损失在千分之五以内。

在PTQ一栏，我们可以看到swinlarge的积分下降严重。我们可以通过禁用FC2和PatchMerge中的一些量化节点来进一步提高量化精度。这实际上是性能和准确性之间的一个平衡。

接下来介绍一下NVIDIA在推理侧实现的加速效果。

左上图为优化后 FP16 与 pytorch 的延迟对比，右下图为优化后 FP16 与 pytorch 的加速比，以及 INT8 优化与 FP16 优化的加速比。可以看到，通过优化，我们在FP16精度方面相对于pytorch可以实现2.82x~7.34x的加速。结合INT8量化，我们可以在此基础上进一步实现1.2x~1.5x的加速。

总结：可视化大模型Swim Transformer训练与推理加速技巧

在本次分享中，我们介绍了一系列训练和推理加速技术，包括：

1.混合精度训练，低精度推理，使用更高性能的Tensor Core进行矩阵乘法或卷积计算；

2、算子融合；

3.cuda内核优化技巧：如矩阵补零、向量化读写、寄存器数组的巧妙运用等

结合以上优化技巧，加速效果为：

训练方面，以Swin-Large模型为例，实现了单卡~2.85x的加速比，8卡2.32x的加速比；

推理方面，以Swin-tiny模型为例，在FP16精度下实现了2.82x - 7.34x的加速比。结合INT8量化，进一步实现了1.2x - 1.5x的加速比。

上述介绍的视觉大模型训练和推理的加速方法已在百度白鸽·AI异构计算平台的AIAK加速功能中实现。