视觉大模型训练和推理加速
百度Geek说约 19 分钟大家好，我是来自 NVIDIA GPU 计算专家团队的陶砺，很高兴今天有机会在这里跟大家分享一下我和我的同事陈庾，在 Swin Transformer 这个视觉大模的型训练和推理优化上的一些工作。其中一些的方法与策略，在其他的模型训练、推理的优化上都可以使用，来提高模型的吞吐、提升 GPU 的使用效率、加快模型的迭代。
我会介绍 Swin Transformer 模型的训练部分的优化，在推理优化部分的工作，将由我的同事来做详细的介绍。
这里是我们今天分享的目录，主要分为四个部分，既然是针对特定模型进行的优化，那么我们首先会简单介绍一下 Swin Transformer 模型。然后，我会结合 profiling 的工具，也就是 nsight system 对训练的流程进行分析和优化。在推理部分，我的同事会给出推理优化的策略和方法，包含较为细节的 cuda 层面的优化。最后，是今天优化内容的一个总结。
首先是第一部分，也就是 Swin Transformer 的介绍。
一、Swin Transformer简介
GEEK TALK
从模型的名称我们可以看出，这是一个基于 transformer 的模型，我们先对 transformer 进行一下简单的回顾。
Transformer 模型从 attention is all you need 这篇文章中被提出后，在自然语言处理领域的很多任务上大放异彩。
Transformer 模型的核心就是所谓的注意力机制，也就是 attention mechanism。对于注意力模块，通常的输入是 query，key 和 value 三个张量。通过 query 和 key 的作用，加上 softmax 的计算，可以得到通常被称为 attention map 的注意力结果，根据 attention map 中的数值的高低，模型就可以学习到需要更加注意 value 中的哪些区域，或者说模型可以学习到，value 中的哪些数值对我们的任务有很大的帮助。这就是最基础的单头注意力模型。
我们通过增加这样单头注意力的模块的数量，也就可以构成常见的多头注意力模块。常见的 encoder、decoder 都是基于这样的多头注意力模块搭建的。
很多模型通常包含了 self-attention，cross-attention 这两种注意力模块，或者是一个或多个模块的堆叠。如著名的 BERT 就是由多个 encoder 模块组成，现在大热的 diffusion 模型通常同时包含了 self-attention 和 cross-attention。
在 Swin Transformer 之前， Vision Transformer (ViT) 首先将 transformer 应用到了计算机视觉领域。ViT 的模型结构，如下图左侧所示，ViT 会将一个图像分割成一系列的 patch，每一个 patch 类比于自然语言处理中的 token，然后通过一个 Transformer-based 的 encoder 对这一系列 patch 进行 encode，最后得到可用于分类等任务的 feature。
而来到 Swin Transformer，它引入了 window attention 的概念，不同于 ViT 对整个图像进行 attention，Swin Transformer 会先将图像划分成若干个 window，然后仅对 window 内部的 patch 进行 attention，从而减少计算量。
为了弥补 window 带来的边界问题，Swin Transformer 进一步引入 window shift 的操作。同时为了使得模型有更丰富的位置信息，还在 attention 时引入了 relative position bias。其实这里的 window attention 和 window shift，就是 Swin Transformer 中的 Swin 名称的由来。
这里给出的是 Swin Transformer 的网络结构，大致的一个网络结构和传统的 CNN 如 ResNet 十分相近。
可以看到整个网络结构被划分为多个 stage，在不同 stage 中间，会有对应的降采样的过程。每个 stage 的分辨率是不一样的，从而形成了一个分辨率金字塔，这样也使得每个 stage 的计算复杂程度也逐渐降低。
然后每个 stage 中会有若干个 transformer block。每一个 transformer block 中，就会用到上面提到的 window attention 模块。
接下来，我们从具体操作的角度来对 Swin Transformer 进行解构。
可以看到，一个 transformer block 中涉及到三大部分，第一部分是 window shift/partition/reverse 的 window 相关的操作，第二部分是 attention 计算，第三部分是 FFN 计算；而 attention 和 FFN 部分又可以进一步细分为若个 op，最终我们可以将整个模型细分为几十个 op 的组合。
这样的算子划分对于我们进行性能分析，定位性能瓶颈以及开展加速优化而言，都是非常重要的。
以上就是第一部分的介绍。接下来，我们来介绍一下在训练上我们进行的一些优化工作，特别的，我们结合 profiling 工具，也就是 nsight system，对整体的训练流程做一个分析和优化。
二、Swin Transformer优化训练
GEEK TALK
对于大模型的训练而言，通常会用到多卡、多节点的计算资源。针对 Swin Transformer，我们发现卡间通讯的开销占比会相对较少，随着卡数的增长，整体速度的提升几乎呈现线性的增长，所以在这里，我们优先对单 GPU 上的计算瓶颈进行分析和优化。
nsight system 是一个系统层面的性能分析工具，通过这个工具，我们可以很方便的看到模型的各个模块的 GPU 的使用情况，是否存在数据等待等可能存在的性能瓶颈和优化空间，可以便于我们合理的规划 CPU、GPU 之间的负载。
nsight system 可以捕捉到 CUDA，以及一些 gpu 计算库如 cublas，cudnn，tensorRT 等调用的核（kernel）函数的调用和运行情况，以及可以方便用户添加一些标记，来统计标记范围内对应 gpu 的运行情况。
一个标准的模型优化流程如下图所示，我们对模型进行 profiling，拿到性能分析报告，发现性能优化点，然后有针对性的去做性能调优。
这里是一个 nsight system 的界面，我们可以很清晰地看到核函数的发射，也就是 kernel launch；核函数的运行，也就是这里的 runtime 部分。对于具体的核函数，我们可以看到在整个流程里的时间占比，以及 gpu 是否存在空闲等信息。在添加完 nvtx 标记之后，我们可以看到模型前向，反向所需要的时间。
在前向部分，如果放大，我们也可以清晰地看到具体每个 SwinTransformer Block 的计算需要的时间。
我们首先通过 nsight system 性能分析工具来看一下整个 baseline 的性能表现，下图中展示的就是 FP32 的 baseline，可以看到它的 GPU 利用率是很高的，而其中占比最高的是矩阵乘的 kernel。
那么对于矩阵乘法而言，我们的一个优化手段，就是充分利用 tensor core 进行加速。
我们知道 NVIDIA 的 GPU 内有 cuda core 和 tensor core 这样的硬件资源，tensor core 是专门为了矩阵乘法的加速的模块。我们可以考虑直接采用 tf32 tensor core 或者混合精度下，采用 fp16 tensor core。要知道，使用 fp16 的 tensor core 在矩阵乘法上的吞吐，会比 tf32 要高，对比纯 fp32 的矩阵乘也会有很高的加速效果。
在此，我们采用了混合精度的方案。通过采用 torch.cuda.amp 的混合精度的模式，我们可以取得了 1. 63 倍的吞吐提升。
在 profiling 的结果里也能够很清晰地看到，原本占最高的矩阵乘，经过优化后，在整个 timeline 中的占比降到了 11.9%。至此，占比较高的 kernel 都是 elementwise kernel。
对于 elementwise kernel，我们首先要了解哪里会用到 elementwise 的 kernel。
Elementwise kernel 里，比较常见的 unrolled elementwise kernel 和 vectorized elementwise kernel。其中 unrolled elementwise kernel 广泛存在于一些有偏置的卷积，或者线性层中，以及一些保证数据在内存连续性的 op 中。
vectorized elementwise kernel 则经常出现在一些激活函数，如 ReLU 的计算中。如果想要减少这里大量的 elementwise kernel，一个常见的做法是做算子融合，比如矩阵乘法中，我们可以通过将 elementwise的操作与矩阵乘法的算子融合在一起，来降低这部分的时间开销。
对于算子融合，一般而言可以为我们带来两个好处：
一个是减少 kernel launch 的开销，如下图所示，两个 cuda kernel 的执行需要两次 launch，那样可能会导致 kernel 之间存在 gap，使得 GPU 空闲，那么如果我们将两个 cuda kernel 融合成一个 cuda kernel，一方面节省了一次 launch，同时也可以避免 gap 的产生。
另外一个好处是减少了 global memory 的访问，因为 global memory 的访问是非常耗时的，而两个独立的 cuda kernel 之间要进行结果传递，都需要通过 global memory，将两个 cuda kernel 融合成一个 kernel，我们可以在寄存器或者 share memory 上进行结果传递，从而避免了一次 global memory 写和读，提升性能。
对于算子融合，我们第一步是采用现成的 apex 库来进行 Layernorm 和 Adam 中操作的融合，可以看通过简单的指令替换，我们可以使能 apex 的 fused layernorm 和 fused Adam，从而使得加速从 1.63 倍提升至 2.11 倍。
从 profling 的日志我们也可以看到，经过算子融合之后，elementwise kernel 在这个 timeline 的占比大幅降低，矩阵乘法重新成为时间占比最大的 kernel。
除了利用现有的 apex 库，我们也进行了手工的融合算子开发。
通过观察 timeline，以及对模型的理解，我们发现 Swin Transformer 中有特有的 window 相关操作，如 window partition/shift/merge 等，这里的一次 window shift，需要调用两个 kernel，并在 shift 完成之后调用 elementwise 的 kernel。并且，attention 模块前如果需要做一次这样的操作，那么之后会有对应的 reverse 操作。这里单单 window shift 调用的 roll_cuda_kernel 就在整个 timeline 中占比 4.6%。
刚才提到的这些操作，其实只是对数据进行了划分，即对应的数据会被划分到一个 window 中去，对应的原始代码如下图所示。
我们发现，这部分的操作其实本质上只是 index mapping，因此，我们对这一部分进行的融合算子开发。开发的过程，我们需要掌握 CUDA 编程的相关知识，并且编写算子的前向计算和反向计算的相关代码。
如何向 pytorch 中引入自定义算子，官方给出了教程，我们可以按照教程编写 CUDA 代码，编译好后就可以作为一个模块引入原始的模型。可以看到，通过引入我们的定制化融合算子，我们可以将加速比进一步提升至 2.19 倍。
接下来展示的是，我们对 mha 部分的融合工作。
Mha 部分是 transformer 模型中一个占比很大的模块，因此对它的优化往往可以带来较大的加速效果。从图中可以看到，在没有进行算子融合之前，mha 部分的操作占比为 37.69%，其中包括了不少 elementwise 的 kernel。如果我们能够将相关操作融合成一个独立的 kernel，并具有更快的速度，加速比可以得到进一步提升。
三、Swin Transformer推理优化
GEEK TALK
推理上的算子融合不需要考虑反向，所以 kernel 开发过程中不需要考虑保存计算梯度所需要的中间结果； 推理过程允许预处理，我们可以对一些只需要一次计算便可重复使用的操作，提前算好，保留结果，每次推理时直接调用从而避免重复计算。
首先我们需要创建句柄和乘法描述符； 接下来我们为每个矩阵创建一个矩阵描述符； 因为一般我们的输入都是常规 layout 的，所以我们需要对常规布局的矩阵进行布局转换，使其变为 IMMA-specific 的布局； 然后再进行 int8 矩阵乘法，得到结果之后，我们可以考虑继续用这个结果进行下游的乘法计算，这样可以避免转变会常规布局的开销； 只有最后一个矩阵乘法的结果，我们需要转换常规布局以便输出。
首先 IMMA-specific 布局对矩阵是有特定的尺寸要求，为了避免推理过程中需要额外分配空间的操作，我们会提前分配好符合 IMMA-specific 布局尺寸的 buffer； 然后，由于 weight 可以一次处理重复使用，所以我们会提前对 weight（相当于乘法中的 B 矩阵）进行布局变换，避免在推理过程中反复变换 weight； 第三个技巧是，对于不得不进行特殊布局变换的 A 和 C，我们会把变换和上游或下游 op 进行算子融合，以便隐藏这部分的开销； 最后一点，是与布局无关，而是 int8 矩阵乘法必有的量化和反量化的操作，我们同样会采用算子融合的方式，把它的 latency 隐藏起来。
四、Swin Transformer优化总结
GEEK TALK
END