视觉大模型训练和推理减速

大家好，我是来自 NVIDIA GPU 计算专家团队的陶砺，很快乐当天无时机在这里跟大家分享一下我和我的共事陈庾，在 Swin Transformer 这个视觉大模的型训练和推理优化上的一些上班。其中一些的方法与战略，在其余的模型训练、推理的优化上都可以经常使用，来提高模型的吞吐、优化 GPU 的经常使用效率、放慢模型的迭代。

我会引见 Swin Transformer 模型的训练局部的优化，在推理优化局部的上班，将由我的共事来做详细的引见

这里是我们当蠢才享的目录，关键分为四个局部，既然是针对特定模型启动的优化，那么我们首先会繁难引见一下 Swin Transformer 模型。而后，我会结合 profiling 的工具，也就是 nsight system 对训练的流程启动剖析和优化。在推理局部，我的共事会给出推理优化的战略和方法，蕴含较为细节的 cuda 层面的优化。最后，是当天优化内容的一个总结。

首先是第一局部，也就是 Swin Transformer 的引见。

从模型的称号我们可以看出，这是一个基于 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，对全体的训练流程做一个剖析和优化。

2. Swin Transformer 训练优化

关于大模型的训练而言，通常会用到多卡、多节点的计算资源。针对 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，这局部的模块除了 query，key 和 value 外，mask 和 bias 都是以 tensor 的方式传入的，我们开发了 fMHA 这样的一个模块，可以将原本的若干 kernel 融合起来。从 fMHA 这个模块触及到的计算来看，针对 Swin Transformer 中遇到的一些 shape，该模块都有比拟清楚的优化。

模型用上 fMHA 模块后，我们可以将减速比进一步优化 2. 85 倍。上述是我们在单卡上取得的训练减速成果，那么我们来看一下单机 8 卡的训练状况，可以看到，经过上述优化，我们可以将训练吞吐从 1612 优化至 3733，取得 2.32 倍的减速。

关于训练优化而言，减速比我们宿愿越高越好，对应的，我们也宿愿减速后的性能能够与减速前坚持分歧。

叠加上上述若干减速打算后，可以看到，模型的收敛性与原始的 baseline 坚持分歧，优化前后的模型的收敛、精度的分歧性，在 Swin-Tiny，Swin-Base 以及 Swin-Large 上都获取了验证。

关于训练局部，一些其余的减速战略包括 CUDA graph、multi-stream 等，都能对 Swin Transformer 的性能有进一步优化；其余方面，目前我们引见的是经常使用混合精度的打算，也就是 Swin Transformer 官网 repo 驳回的战略；经常使用纯 fp16 的打算（即 apex O2 形式）可以到达更快的减速成果。

只管 Swin 对通讯的要求不高，但是关于多节点大模型的训练，相比于原始的散布式训练，经常使用正当的战略去暗藏通讯的开支，能够在多卡训练上取得进一步的收益。

接上去，有请我的共事来引见一下我们在推理上的减速打算和成果。

3. Swin Transformer 推理优化

大家好，我是来自英伟达 GPU 计算专家团队的陈庾，十分感谢陶砺在训练减速上的引见，接下来由我来引见一下推理上的减速。

跟训练一样，推理的减速离不开算子融合这一打算。不过相关于训练而言，在推理上启动算子融合有更好的灵敏性，关键表现有两点：

在推理侧，我们可以启动不少的算子融合，这里给出的是我们在 Transformer 模型中经常出现的一些算子融合的 pattern 以及成功关系 pattern 所须要用到的工具。

首先，我们独自列出矩阵乘法和卷积，是由于有一大类算子融合是围绕他们启动的，关于矩阵乘法关系的融合，我们可以思考驳回 cublas，cutlass，cudnn 这三个库；关于卷积，我们可以驳回 cudnn 或许 cutlass。那么关于矩阵乘法的算子融合而言，在 Transformer 模型中，我们演绎为 gemm + elementwise 的操作，比如 gemm + bias, gemm + bias + 激活函数等，这一类的算子融合，我们可以思考间接调用 cublas 或 cutlass 来成功。

此外，假设我们 gemm 之后的 op 操作比拟复杂，比如 layernorm，transpose 等，我们可以思考将 gemm 和 bias 分开，而后把 bias 融合到下一个 op 中，这样可以更为容易地调用 cublas 来成功繁难的矩阵乘法，当然这种 bias 和下一个 op 启动融合的 pattern 普通是须要我们手写 cuda kernel 来成功。

最后，有一些特定 op，雷同须要我们以手写 cuda kernel 的方式启动融合，比如 layernorm + shift + window partition。

由于算子融合须要我们比拟奇妙地设计 cuda kernel，所以我们普通倡导先经过 nsight system 性能剖析工具对全体 pipeline 启动剖析，优先针对热点模块启动算子融合优化，以到达性能和上班量的平衡。

那么在泛滥的算子融合优化中，我们筛选了两个减速成果比拟清楚的算子启动引见。

首先是 mha 局部的算子融合，我们将 position bias lookup 这一操作提早到预处置局部，从而防止每次推理时都启动 lookup。

而后将 batch gemm，softmax，batch gemm 融分解一个独立的 fMHA kernel，同时我们把 transpose 关系的操作融合到了 fMHA kernel I/O 操作中，经过必定的数据读写的 pattern 来防止显式的 transpose 操作。

可以看到，融合后该局部取得了 10 倍的减速，而端到端也取得了 1.58 倍的减速。

另一个我想引见一下的算子融合是 QKV gemm + bias 的融合。

gemm 和 bias 的融合是一个十分经常出现的融合手腕，在这里为了配合我们前面提到的 fMHAkernel，我们须要对 weight 和 bias 提早启动格局上的变换。

我之所以在这里选用引见这个算子融合，也正是由于这种提早变换表现了我们前面提到的，推理上启动算子融合的灵敏性，我们可以对模型的推理流程做一些不影响其精度的变动，从而成功更好算子融合 pattern，取得更好的减速成果。

最后，经过 QKV gemm+bias 的融合，我们可以进一步取得 1.1 倍的端到端减速。

下一个优化手腕是矩阵乘法 padding。

在 Swin Transformer 的计算中，有时刻我们会遇到主维为奇数的矩阵乘法，这时刻并不利于我们的矩阵乘法 kernel 启意向量化读写，从而使得 kernel 的运转效率变低，此时我们可以思考对介入运算的矩阵主维启动 padding 操作，使其变为 8 的倍数，这样一来，矩阵乘 kernel 就可以以 alignment=8，一次性读写 8 个元素的方式来启意向量化读写，优化性能。

如下表所示，我们将 n 从 49 padding 到 56 后，矩阵乘法的 latency 从 60.54us 降低为 40.38us，取得了1.5 倍的减速比。

下一个优化手腕是巧用 half2 或许 char4 这样的数据类型。

以下的代码是一个 half2 优化的示例，它成功的是一个繁难的加 bias 再加残差这样的算子融合操作，可以看到经过经常使用 half2 数据类型，相关于 half 数据类，我们可以将 latency 从 20.96us 降低到 10.78us，减速 1.94 倍。

那么驳回 half2 数据类型普通有什么好处呢？关键有三点：

第一个好处是向量化读写可以优化 memory 的带宽应用效率并降低访存指令数；如下图右侧所示，经过 half2 的经常使用，访存指令缩小了一半，同时 memory 的 SOL 也有清楚优化；

第二个好处是结合 half2 专有的高吞吐的数学指令，可以减低 kernel 的 latency。这两点都曾经体如今了这个示例程序中；

第三个好处是在启动 reduction 关系 kernel 开发时，驳回 half2 数据类型象征着一个 cuda 线程同时处置两个元素，可以有效缩小闲暇的线程数，也可以缩小线程同步的 latency。

下一个优化手腕是巧用寄存器数组。

在我们启动 layernorm 或许 softmax 等 Transformer 模型经常出现的算子优化时，我们经常须要在一个 kernel 中屡次经常使用同一个输入数据，那么相关于每次都从 global memory 读取，我们可以驳回寄存器数组来缓存数据，从而防止重复读取 global memory。

由于寄存器是每个 cuda 线程独占的，所以在启动 kernel 设计时，我们须要提早设定好每个 cuda 线程所须要缓存的元素个数，从而开拓对应大小的寄存器数组，并且在调配每个 cuda 线程所担任元素时，须要确保我们可以做到兼并访问，如下图右上侧所示，当我们有 8 个线程时，0 号线程可以处置 0 号元素，当我们有 4 个线程是，0 号线程则处置 0 号和 4 号元素，如此类推。

我们普通倡导可以驳回模板函数的方式，经过模板参数来控制每 个cuda 线程的寄存器数组大小。

此外，在经常使用寄存器数组时，须要保障我们的下标是常量，假设是循环变量作为下标，我们应该尽量保障可以启动循环展开，这样可以防止编译器将数据放到了 latency 很高的 local memory 中，如下图所示，我们在循环条件中参与限度，经过 ncu report 可以看到，防止了 local memory 的经常使用。

最后一个我想引见优化手腕是 INT8 量化。

INT8 量化是推理减速十分关键的减速手腕，关于 Transformer based 的模型而言，INT8 量化可以在缩小显存消耗的同时带来更好的性能。

而关于 Swin 来说，经过结适宜合的 PTQ 或 QAT 量化打算，可以在取得良好减速的同时，保障量化精度。普通我们启动 int8 量化，关键是对矩阵乘法或许卷积启动量化，比如 int8 矩阵乘法中，我们会先将原始的 FP32 或 FP16 的 input 和 weight 量化为 INT8 而后再启动 INT8 矩阵乘法，累加到 INT32 数据类型上，这是我们会启动反量化操作，获取 FP32 或 FP16 的结果。

比拟经常出现调用 INT8 矩阵乘法的工具是 cublasLt，为了可以取得更好的性能，我们有必要深化地了解一下 cublasLt api 的一些个性。

cublasLt 关于 int8 矩阵乘法，提供了两种输入类型，区分是下图左侧所示，以 INT32 输入，或许下图右侧所示，以 INT8 输入，图中蓝框所示的 cublasLt 的计算操作。

可以看到相关于 INT32 输入而言， INT8 输入会多了一对反量化和量化操作，这样一来普通会带来更多的精度损失，但是由于 INT8 输入，在写出到 global memory 时相对 INT32 输入少了 3/4 的数据量，性能会更好，所以这外面存在着精度和性能 tradeoff。

那么关于 Swin Transformer 而言，我们发现配合 QAT，以 INT8 输入会在取好的减速比的前提下，保障精度，由于我们驳回了 INT8 输入的打算。

另外，关于 cublasLt 中 INT8 矩阵乘法，还须要思考数据的规划疑问，cublasLt 支持两种规划，一种 IMMA-specific 的规划，其中触及到一些比拟复杂的格局，而且在这种规划只支持 NT-gemm，另外一种是惯例的列优先的规划，在此规划下支持 TN-gemm。

普通来说，驳回列优先的规划，会更无利于整个 pipeline 代码的开发，由于假设我们用 IMMA-specific 规划的话，我们为了兼容这种规划或许须要很多额外的操作，以及上下游 kernel 也须要为这种不凡规划做兼容。但是在一些尺寸的矩阵乘法上，IMMA-specific 规划或许会有更好的性能，所以假设我们要尝试搭建 int8 推理的话，倡导我们可以先做一些 benchmark，以便更好地从性能和开发难易水平做取舍。

在 FasterTransformer 中我们驳回了 IMMA-specific 规划。所以接上去，我们以 IMMA-specific 规划为例，繁难引见了一下 cublasLt int8 矩阵乘法的基本搭建流程，以及一些开发技巧。

cublasLt int8 矩阵乘法的基本搭建流程，一共可以分为 5 步：

上述引见了 IMMA-specific 规划下的搭建流程，可以看到外面会有不少限度。为了防止这些限度对性能的影响，我们在 Faster Transformer 中驳回了以下技巧：

以下是我们在 Faster Transformer 中驳回的的 INT8 流程的示用意，可以看到，一切矩阵乘都变为了 int8 数据类型，每个 int8 矩阵乘法前后都会拔出对应的量化和反量化节点，而后关于加 bias，加残差或 layernorm 等操作，我们还是保管原始的 FP32 或 FP16 数据类型，当然它的 I/O 或许是 int8 的，从而会比 FP16 或 FP32 I/O 性能要好。

这里展现的是 Swin Transformer int8 量化的精度状况，经过 QAT我们可以保障精度损失在千分之 5 以内。

而在 PTQ 那一列，我们可以看到 Swin-Large 的掉点比拟重大，普通对应掉点重大的疑问，我们都可以思考驳回缩小一些量化节点的方式来优化量化精度，当然这样或许会带来减速成果的削弱。

在 FT 中，我们可以经过禁用 FC2 和 PatchMerge 中 int 8 矩阵乘法的 int8 输入前的反量化和量化结点（即驳回 int32 输入），来进一步优化量化精度，可以看到在此优化操作下，swin-large 的 PTQ 精度也清楚优化了。

接上去是我们推理侧取得的减速成果，我们区分在不同型号的 GPUT4、A10、A100 上启动了跟 pytorch FP16 成功的性能对比。

其中下图左侧是优化后跟 pytorch 的 latency 对比，右图为优化后 FP16 下跟 pytorch 以及 INT8 优化跟 FP16 优化的减速比。可以看到，经过优化，在 FP16 精度上，我们可以取得，相关于 pytorch 2.82x ~ 7.34x 的减速，结合 INT8 量化，我们可以在此基础上进一步取得 1.2x ~ 1.5x 的减速。

4. Swin Transformer 优化总结

最后，我们总结一下，本次分享中我们引见了如何经过 nsight system 性能剖析工具发现性能瓶颈，而后针对性能瓶颈，引见了一系列训练推理减速技巧，其中包括1. 混合精度训练 / 低精度推理，2. 算子融合，3. cuda kernel 优化技巧 ：如矩阵补零，向量化读写，巧用寄存器数组等，4. 推理优化上驳回一些预处置，来完善我们的计算流程；我们也引见了 multi-stream，cuda graph 的一些运行 。

结合上述优化，我们在训练上，以 Swin-Large 模型为例取得了单卡2.85x 的减速比，8 卡 2.32x 的减速比；在推理上，以 Swin-tiny 模型为例，在 FP16 精度下取得了 2.82x~7.34x 的减速比，结合 INT8 量化，进一步取得 1.2x~1.5x 的减速比。

上述视觉大模型训练与推理的减速方法都曾经在百度百舸 AI 异构计算平台的 AIAK 减速配置中成功，欢迎大家经常使用。