大家好，我是來自 NVIDIA GPU 計(jì)算專家團(tuán)隊(duì)的陶礪，很高興今天有機(jī)會在這里跟大家分享一下我和我的同事陳庾，在 Swin Transformer 這個視覺大模的型訓(xùn)練和推理優(yōu)化上的一些工作。其中一些的方法與策略，在其他的模型訓(xùn)練、推理的優(yōu)化上都可以使用，來提高模型的吞吐、提升 GPU 的使用效率、加快模型的迭代。

我會介紹 Swin Transformer 模型的訓(xùn)練部分的優(yōu)化，在推理優(yōu)化部分的工作，將由我的同事來做詳細(xì)的介紹

這里是我們今天分享的目錄，主要分為四個部分，既然是針對特定模型進(jìn)行的優(yōu)化，那么我們首先會簡單介紹一下 Swin Transformer 模型。然后，我會結(jié)合 profiling 的工具，也就是 nsight system 對訓(xùn)練的流程進(jìn)行分析和優(yōu)化。在推理部分，我的同事會給出推理優(yōu)化的策略和方法，包含較為細(xì)節(jié)的 cuda 層面的優(yōu)化。最后，是今天優(yōu)化內(nèi)容的一個總結(jié)。

首先是第一部分，也就是 Swin Transformer 的介紹。

1. Swin Transformer 簡介

從模型的名稱我們可以看出，這是一個基于 transformer 的模型，我們先對 transformer 進(jìn)行一下簡單的回顧。

Transformer 模型從 attention is all you need 這篇文章中被提出后，在自然語言處理領(lǐng)域的很多任務(wù)上大放異彩。

Transformer 模型的核心就是所謂的注意力機(jī)制，也就是 attention mechanism。對于注意力模塊，通常的輸入是 query，key 和 value 三個張量。通過 query 和 key 的作用，加上 softmax 的計(jì)算，可以得到通常被稱為 attention map 的注意力結(jié)果，根據(jù) attention map 中的數(shù)值的高低，模型就可以學(xué)習(xí)到需要更加注意 value 中的哪些區(qū)域，或者說模型可以學(xué)習(xí)到，value 中的哪些數(shù)值對我們的任務(wù)有很大的幫助。這就是最基礎(chǔ)的單頭注意力模型。

我們通過增加這樣單頭注意力的模塊的數(shù)量，也就可以構(gòu)成常見的多頭注意力模塊。常見的 encoder、decoder 都是基于這樣的多頭注意力模塊搭建的。

很多模型通常包含了 self-attention，cross-attention 這兩種注意力模塊，或者是一個或多個模塊的堆疊。如著名的 BERT 就是由多個 encoder 模塊組成，現(xiàn)在大熱的 diffusion 模型通常同時包含了 self-attention 和 cross-attention。

在 Swin Transformer 之前， Vision Transformer (ViT) 首先將 transformer 應(yīng)用到了計(jì)算機(jī)視覺領(lǐng)域。ViT 的模型結(jié)構(gòu)，如下圖左側(cè)所示，ViT 會將一個圖像分割成一系列的 patch，每一個 patch 類比于自然語言處理中的 token，然后通過一個 Transformer-based 的 encoder 對這一系列 patch 進(jìn)行 encode，最后得到可用于分類等任務(wù)的 feature。

而來到 Swin Transformer，它引入了 window attention 的概念，不同于 ViT 對整個圖像進(jìn)行 attention，Swin Transformer 會先將圖像劃分成若干個 window，然后僅對 window 內(nèi)部的 patch 進(jìn)行 attention，從而減少計(jì)算量。

為了彌補(bǔ) window 帶來的邊界問題，Swin Transformer 進(jìn)一步引入 window shift 的操作。同時為了使得模型有更豐富的位置信息，還在 attention 時引入了 relative position bias。其實(shí)這里的 window attention 和 window shift，就是 Swin Transformer 中的 Swin 名稱的由來。

這里給出的是 Swin Transformer 的網(wǎng)絡(luò)結(jié)構(gòu)，大致的一個網(wǎng)絡(luò)結(jié)構(gòu)和傳統(tǒng)的 CNN 如 ResNet 十分相近。

可以看到整個網(wǎng)絡(luò)結(jié)構(gòu)被劃分為多個 stage，在不同 stage 中間，會有對應(yīng)的降采樣的過程。每個 stage 的分辨率是不一樣的，從而形成了一個分辨率金字塔，這樣也使得每個 stage 的計(jì)算復(fù)雜程度也逐漸降低。

然后每個 stage 中會有若干個 transformer block。每一個 transformer block 中，就會用到上面提到的 window attention 模塊。

接下來，我們從具體操作的角度來對 Swin Transformer 進(jìn)行解構(gòu)。

可以看到，一個 transformer block 中涉及到三大部分，第一部分是 window shift/partition/reverse 的 window 相關(guān)的操作，第二部分是 attention 計(jì)算，第三部分是 FFN 計(jì)算；而 attention 和 FFN 部分又可以進(jìn)一步細(xì)分為若個 op，最終我們可以將整個模型細(xì)分為幾十個 op 的組合。

這樣的算子劃分對于我們進(jìn)行性能分析，定位性能瓶頸以及開展加速優(yōu)化而言，都是非常重要的。

以上就是第一部分的介紹。接下來，我們來介紹一下在訓(xùn)練上我們進(jìn)行的一些優(yōu)化工作，特別的，我們結(jié)合 profiling 工具，也就是 nsight system，對整體的訓(xùn)練流程做一個分析和優(yōu)化。

2. Swin Transformer 訓(xùn)練優(yōu)化

對于大模型的訓(xùn)練而言，通常會用到多卡、多節(jié)點(diǎn)的計(jì)算資源。針對 Swin Transformer，我們發(fā)現(xiàn)卡間通訊的開銷占比會相對較少，隨著卡數(shù)的增長，整體速度的提升幾乎呈現(xiàn)線性的增長，所以在這里，我們優(yōu)先對單 GPU 上的計(jì)算瓶頸進(jìn)行分析和優(yōu)化。

nsight system 是一個系統(tǒng)層面的性能分析工具，通過這個工具，我們可以很方便的看到模型的各個模塊的 GPU 的使用情況，是否存在數(shù)據(jù)等待等可能存在的性能瓶頸和優(yōu)化空間，可以便于我們合理的規(guī)劃 CPU、GPU 之間的負(fù)載。

nsight system 可以捕捉到 CUDA，以及一些 gpu 計(jì)算庫如 cublas，cudnn，tensorRT 等調(diào)用的核（kernel）函數(shù)的調(diào)用和運(yùn)行情況，以及可以方便用戶添加一些標(biāo)記，來統(tǒng)計(jì)標(biāo)記范圍內(nèi)對應(yīng) gpu 的運(yùn)行情況。

一個標(biāo)準(zhǔn)的模型優(yōu)化流程如下圖所示，我們對模型進(jìn)行 profiling，拿到性能分析報告，發(fā)現(xiàn)性能優(yōu)化點(diǎn)，然后有針對性的去做性能調(diào)優(yōu)。

這里是一個 nsight system 的界面，我們可以很清晰地看到核函數(shù)的發(fā)射，也就是 kernel launch；核函數(shù)的運(yùn)行，也就是這里的 runtime 部分。對于具體的核函數(shù)，我們可以看到在整個流程里的時間占比，以及 gpu 是否存在空閑等信息。在添加完 nvtx 標(biāo)記之后，我們可以看到模型前向，反向所需要的時間。

在前向部分，如果放大，我們也可以清晰地看到具體每個 SwinTransformer Block 的計(jì)算需要的時間。

我們首先通過 nsight system 性能分析工具來看一下整個 baseline 的性能表現(xiàn)，下圖中展示的就是 FP32 的 baseline，可以看到它的 GPU 利用率是很高的，而其中占比最高的是矩陣乘的 kernel。

那么對于矩陣乘法而言，我們的一個優(yōu)化手段，就是充分利用 tensor core 進(jìn)行加速。

我們知道 NVIDIA 的 GPU 內(nèi)有 cuda core 和 tensor core 這樣的硬件資源，tensor core 是專門為了矩陣乘法的加速的模塊。我們可以考慮直接采用 tf32 tensor core 或者混合精度下，采用 fp16 tensor core。要知道，使用 fp16 的 tensor core 在矩陣乘法上的吞吐，會比 tf32 要高，對比純 fp32 的矩陣乘也會有很高的加速效果。

在此，我們采用了混合精度的方案。通過采用 torch.cuda.amp 的混合精度的模式，我們可以取得了 1. 63 倍的吞吐提升。

在 profiling 的結(jié)果里也能夠很清晰地看到，原本占最高的矩陣乘，經(jīng)過優(yōu)化后，在整個 timeline 中的占比降到了 11.9%。至此，占比較高的 kernel 都是 elementwise kernel。

對于 elementwise kernel，我們首先要了解哪里會用到 elementwise 的 kernel。

Elementwise kernel 里，比較常見的 unrolled elementwise kernel 和 vectorized elementwise kernel。其中 unrolled elementwise kernel 廣泛存在于一些有偏置的卷積，或者線性層中，以及一些保證數(shù)據(jù)在內(nèi)存連續(xù)性的 op 中。

vectorized elementwise kernel 則經(jīng)常出現(xiàn)在一些激活函數(shù)，如 ReLU 的計(jì)算中。如果想要減少這里大量的 elementwise kernel，一個常見的做法是做算子融合，比如矩陣乘法中，我們可以通過將 elementwise的操作與矩陣乘法的算子融合在一起，來降低這部分的時間開銷。

對于算子融合，一般而言可以為我們帶來兩個好處：

一個是減少 kernel launch 的開銷，如下圖所示，兩個 cuda kernel 的執(zhí)行需要兩次 launch，那樣可能會導(dǎo)致 kernel 之間存在 gap，使得 GPU 空閑，那么如果我們將兩個 cuda kernel 融合成一個 cuda kernel，一方面節(jié)省了一次 launch，同時也可以避免 gap 的產(chǎn)生。

另外一個好處是減少了 global memory 的訪問，因?yàn)?global memory 的訪問是非常耗時的，而兩個獨(dú)立的 cuda kernel 之間要進(jìn)行結(jié)果傳遞，都需要通過 global memory，將兩個 cuda kernel 融合成一個 kernel，我們可以在寄存器或者 share memory 上進(jìn)行結(jié)果傳遞，從而避免了一次 global memory 寫和讀，提升性能。

對于算子融合，我們第一步是采用現(xiàn)成的 apex 庫來進(jìn)行 Layernorm 和 Adam 中操作的融合，可以看通過簡單的指令替換，我們可以使能 apex 的 fused layernorm 和  fused Adam，從而使得加速從 1.63 倍提升至 2.11 倍。

從 profling 的日志我們也可以看到，經(jīng)過算子融合之后，elementwise kernel 在這個 timeline 的占比大幅降低，矩陣乘法重新成為時間占比最大的 kernel。

除了利用現(xiàn)有的 apex 庫，我們也進(jìn)行了手工的融合算子開發(fā)。

通過觀察 timeline，以及對模型的理解，我們發(fā)現(xiàn) Swin Transformer 中有特有的 window 相關(guān)操作，如 window partition/shift/merge 等，這里的一次 window shift，需要調(diào)用兩個 kernel，并在 shift 完成之后調(diào)用 elementwise 的 kernel。并且，attention 模塊前如果需要做一次這樣的操作，那么之后會有對應(yīng)的 reverse 操作。這里單單 window shift 調(diào)用的 roll_cuda_kernel 就在整個 timeline 中占比 4.6%。

剛才提到的這些操作，其實(shí)只是對數(shù)據(jù)進(jìn)行了劃分，即對應(yīng)的數(shù)據(jù)會被劃分到一個 window 中去，對應(yīng)的原始代碼如下圖所示。

我們發(fā)現(xiàn)，這部分的操作其實(shí)本質(zhì)上只是 index mapping，因此，我們對這一部分進(jìn)行的融合算子開發(fā)。開發(fā)的過程，我們需要掌握 CUDA 編程的相關(guān)知識，并且編寫算子的前向計(jì)算和反向計(jì)算的相關(guān)代碼。

如何向 pytorch 中引入自定義算子，官方給出了教程，我們可以按照教程編寫 CUDA 代碼，編譯好后就可以作為一個模塊引入原始的模型。可以看到，通過引入我們的定制化融合算子，我們可以將加速比進(jìn)一步提升至 2.19 倍。

接下來展示的是，我們對 mha 部分的融合工作。

Mha 部分是 transformer 模型中一個占比很大的模塊，因此對它的優(yōu)化往往可以帶來較大的加速效果。從圖中可以看到，在沒有進(jìn)行算子融合之前，mha 部分的操作占比為 37.69%，其中包括了不少 elementwise 的 kernel。如果我們能夠?qū)⑾嚓P(guān)操作融合成一個獨(dú)立的 kernel，并具有更快的速度，加速比可以得到進(jìn)一步提升。

對于 Swin Transformer，這部分的模塊除了 query，key 和 value 外，mask 和 bias 都是以 tensor 的形式傳入的，我們開發(fā)了 fMHA 這樣的一個模塊，可以將原本的若干 kernel 融合起來。從 fMHA 這個模塊涉及到的計(jì)算來看，針對 Swin Transformer 中遇到的一些 shape，該模塊都有比較顯著的提升。

模型用上 fMHA 模塊后，我們可以將加速比進(jìn)一步提升 2. 85 倍。上述是我們在單卡上取得的訓(xùn)練加速效果，那么我們來看一下單機(jī) 8 卡的訓(xùn)練情況，可以看到，通過上述優(yōu)化，我們可以將訓(xùn)練吞吐從 1612 提升至 3733，取得 2.32 倍的加速。

對于訓(xùn)練優(yōu)化而言，加速比我們希望越高越好，對應(yīng)的，我們也希望加速后的性能能夠與加速前保持一致。

疊加上上述若干加速方案后，可以看到，模型的收斂性與原始的 baseline 保持一致，優(yōu)化前后的模型的收斂、精度的一致性，在 Swin-Tiny，Swin-Base 以及 Swin-Large 上都得到了驗(yàn)證。

關(guān)于訓(xùn)練部分，一些其他的加速策略包括 CUDA graph、multi-stream 等，都能對 Swin Transformer 的性能有進(jìn)一步提升；其他方面，目前我們介紹的是使用混合精度的方案，也就是 Swin Transformer 官方 repo 采用的策略；使用純 fp16 的方案（即 apex O2 模式）可以達(dá)到更快的加速效果。

雖然 Swin 對通信的要求不高，但是對于多節(jié)點(diǎn)大模型的訓(xùn)練，相比于原始的分布式訓(xùn)練，使用合理的策略去隱藏通信的開銷，能夠在多卡訓(xùn)練上獲得進(jìn)一步的收益。

接下來，有請我的同事來介紹一下我們在推理上的加速方案和效果。

3. Swin Transformer 推理優(yōu)化

大家好，我是來自英偉達(dá) GPU 計(jì)算專家團(tuán)隊(duì)的陳庾，非常感謝陶礪在訓(xùn)練加速上的介紹，接下來由我來介紹一下推理上的加速。

跟訓(xùn)練一樣，推理的加速離不開算子融合這一方案。不過相對于訓(xùn)練而言，在推理上進(jìn)行算子融合有更好的靈活性，主要體現(xiàn)有兩點(diǎn)：

• 推理上的算子融合不需要考慮反向，所以 kernel 開發(fā)過程中不需要考慮保存計(jì)算梯度所需要的中間結(jié)果；
• 推理過程允許預(yù)處理，我們可以對一些只需要一次計(jì)算便可重復(fù)使用的操作，提前算好，保留結(jié)果，每次推理時直接調(diào)用從而避免重復(fù)計(jì)算。

在推理側(cè)，我們可以進(jìn)行不少的算子融合，這里給出的是我們在 Transformer 模型中常見的一些算子融合的 pattern 以及實(shí)現(xiàn)相關(guān) pattern 所需要用到的工具。

首先，我們單獨(dú)列出矩陣乘法和卷積，是因?yàn)橛幸淮箢愃阕尤诤鲜菄@他們進(jìn)行的，對于矩陣乘法相關(guān)的融合，我們可以考慮采用 cublas，cutlass，cudnn 這三個庫；對于卷積，我們可以采用 cudnn 或者 cutlass。那么對于矩陣乘法的算子融合而言，在 Transformer 模型中，我們歸納為 gemm + elementwise 的操作，比如 gemm + bias, gemm + bias + 激活函數(shù)等，這一類的算子融合，我們可以考慮直接調(diào)用 cublas 或 cutlass 來實(shí)現(xiàn)。

此外，如果我們 gemm 之后的 op 操作比較復(fù)雜，比如 layernorm，transpose 等，我們可以考慮將 gemm 和 bias 分開，然后把 bias  融合到下一個 op 中，這樣可以更為容易地調(diào)用 cublas 來實(shí)現(xiàn)簡單的矩陣乘法，當(dāng)然這種 bias 和下一個 op 進(jìn)行融合的 pattern 一般是需要我們手寫 cuda kernel 來實(shí)現(xiàn)。

最后，有一些特定 op，同樣需要我們以手寫 cuda kernel 的方式進(jìn)行融合，比如 layernorm + shift + window partition。

由于算子融合需要我們比較巧妙地設(shè)計(jì) cuda kernel，所以我們一般建議先通過 nsight system 性能分析工具對整體 pipeline 進(jìn)行分析，優(yōu)先針對熱點(diǎn)模塊進(jìn)行算子融合優(yōu)化，以達(dá)到性能和工作量的平衡。

那么在眾多的算子融合優(yōu)化中，我們挑選了兩個加速效果比較明顯的算子進(jìn)行介紹。

首先是 mha 部分的算子融合，我們將 position bias lookup 這一操作提前到預(yù)處理部分，從而避免每次推理時都進(jìn)行 lookup。

然后將 batch gemm，softmax，batch gemm 融合成一個獨(dú)立的 fMHA kernel，同時我們把 transpose 相關(guān)的操作融合到了 fMHA kernel I/O 操作中，通過一定的數(shù)據(jù)讀寫的 pattern 來避免顯式的 transpose 操作。

可以看到，融合后該部分取得了 10 倍的加速，而端到端也取得了 1.58 倍的加速。

另一個我想介紹一下的算子融合是 QKV gemm + bias 的融合。

gemm 和 bias 的融合是一個十分常見的融合手段，在這里為了配合我們前面提到的 fMHA kernel，我們需要對 weight 和 bias 提前進(jìn)行格式上的變換。

我之所以在這里選擇介紹這個算子融合，也正是因?yàn)檫@種提前變換體現(xiàn)了我們前面提到的，推理上進(jìn)行算子融合的靈活性，我們可以對模型的推理流程做一些不影響其精度的變化，從而實(shí)現(xiàn)更好算子融合 pattern，取得更好的加速效果。

最后，通過 QKV gemm+bias 的融合，我們可以進(jìn)一步取得 1.1 倍的端到端加速。

下一個優(yōu)化手段是矩陣乘法 padding。

在 Swin Transformer 的計(jì)算中，有時候我們會遇到主維為奇數(shù)的矩陣乘法，這時候并不利于我們的矩陣乘法 kernel 進(jìn)行向量化讀寫，從而使得 kernel 的運(yùn)行效率變低，此時我們可以考慮對參與運(yùn)算的矩陣主維進(jìn)行 padding 操作，使其變?yōu)?8 的倍數(shù)，這樣一來，矩陣乘 kernel 就可以以 alignment=8，一次讀寫 8 個元素的方式來進(jìn)行向量化讀寫，提升性能。

如下表所示，我們將 n 從 49 padding 到 56 后，矩陣乘法的 latency 從 60.54us 下降為 40.38us，取得了 1.5 倍的加速比。

下一個優(yōu)化手段是巧用 half2 或者 char4 這樣的數(shù)據(jù)類型。

以下的代碼是一個 half2 優(yōu)化的示例，它實(shí)現(xiàn)的是一個簡單的加 bias 再加殘差這樣的算子融合操作，可以看到通過使用 half2 數(shù)據(jù)類型，相對于 half 數(shù)據(jù)類，我們可以將 latency 從 20.96us  下降到 10.78us，加速 1.94 倍。

那么采用 half2 數(shù)據(jù)類型一般有什么好處呢？主要有三點(diǎn)：

第一個好處是向量化讀寫可以提升 memory 的帶寬利用效率并降低訪存指令數(shù)；如下圖右側(cè)所示，通過 half2 的使用，訪存指令減少了一半，同時 memory 的 SOL 也有顯著提升；

第二個好處是結(jié)合 half2 專有的高吞吐的數(shù)學(xué)指令，可以減低 kernel 的 latency。這兩點(diǎn)都已經(jīng)體現(xiàn)在了這個示例程序中；

第三個好處是在進(jìn)行 reduction 相關(guān) kernel 開發(fā)時，采用 half2 數(shù)據(jù)類型意味著一個 cuda 線程同時處理兩個元素，可以有效減少空閑的線程數(shù)，也可以減少線程同步的 latency。

下一個優(yōu)化手段是巧用寄存器數(shù)組。

在我們進(jìn)行 layernorm 或者 softmax 等 Transformer 模型常見的算子優(yōu)化時，我們經(jīng)常需要在一個 kernel 中多次使用同一個輸入數(shù)據(jù)，那么相對于每次都從 global memory 讀取，我們可以采用寄存器數(shù)組來緩存數(shù)據(jù)，從而避免重復(fù)讀取 global memory。

由于寄存器是每個 cuda 線程獨(dú)占的，所以在進(jìn)行 kernel 設(shè)計(jì)時，我們需要提前設(shè)定好每個 cuda 線程所需要緩存的元素個數(shù)，從而開辟對應(yīng)大小的寄存器數(shù)組，并且在分配每個 cuda 線程所負(fù)責(zé)元素時，需要確保我們可以做到合并訪問，如下圖右上側(cè)所示，當(dāng)我們有 8 個線程時，0 號線程可以處理 0 號元素，當(dāng)我們有 4 個線程是，0 號線程則處理 0 號和 4 號元素，如此類推。

我們一般建議可以采用模板函數(shù)的方式，通過模板參數(shù)來控制每 個cuda 線程的寄存器數(shù)組大小。

此外，在使用寄存器數(shù)組時，需要保證我們的下標(biāo)是常量，如果是循環(huán)變量作為下標(biāo)，我們應(yīng)該盡量保證可以進(jìn)行循環(huán)展開，這樣可以避免編譯器將數(shù)據(jù)放到了 latency 很高的 local memory 中，如下圖所示，我們在循環(huán)條件中添加限制，通過 ncu report 可以看到，避免了 local memory 的使用。

最后一個我想介紹優(yōu)化手段是 INT8 量化。

INT8 量化是推理加速非常重要的加速手段，對于 Transformer based 的模型而言，INT8 量化可以在減少顯存消耗的同時帶來更好的性能。

而對于 Swin 來說，通過結(jié)合合適的 PTQ 或 QAT 量化方案，可以在取得良好加速的同時，保證量化精度。一般我們進(jìn)行 int8 量化，主要是對矩陣乘法或者卷積進(jìn)行量化，比如 int8 矩陣乘法中，我們會先將原始的 FP32 或 FP16 的 input 和 weight 量化為 INT8 然后再進(jìn)行 INT8 矩陣乘法，累加到 INT32 數(shù)據(jù)類型上，這是我們會進(jìn)行反量化操作，得到 FP32 或 FP16 的結(jié)果。

比較常見調(diào)用 INT8 矩陣乘法的工具是 cublasLt，為了可以取得更好的性能，我們有必要深入地了解一下 cublasLt api 的一些特性。

cublasLt 對于 int8 矩陣乘法，提供了兩種輸出類型，分別是下圖左側(cè)所示，以 INT32 輸出，或者下圖右側(cè)所示，以 INT8 輸出，圖中藍(lán)框所示的 cublasLt 的計(jì)算操作。

可以看到相對于 INT32 輸出而言， INT8 輸出會多了一對反量化和量化操作，這樣一來一般會帶來更多的精度損失，但是由于 INT8 輸出，在寫出到 global memory 時相對 INT32 輸出少了 3/4 的數(shù)據(jù)量，性能會更好，所以這里面存在著精度和性能 tradeoff。

那么對于 Swin Transformer 而言，我們發(fā)現(xiàn)配合 QAT，以 INT8 輸出會在取好的加速比的前提下，保證精度，因?yàn)槲覀儾捎昧?INT8 輸出的方案。

另外，關(guān)于 cublasLt  中 INT8  矩陣乘法，還需要考慮數(shù)據(jù)的布局問題，cublasLt 支持兩種布局，一種 IMMA-specific 的布局，其中涉及到一些比較復(fù)雜的格式，而且在這種布局只支持 NT-gemm，另外一種是常規(guī)的列優(yōu)先的布局，在此布局下支持 TN-gemm。

一般來說，采用列優(yōu)先的布局，會更有利于整個 pipeline 代碼的開發(fā)，因?yàn)槿绻覀冇?IMMA-specific 布局的話，我們?yōu)榱思嫒葸@種布局可能需要很多額外的操作，以及上下游 kernel 也需要為這種特殊布局做兼容。但是在一些尺寸的矩陣乘法上，IMMA-specific 布局可能會有更好的性能，所以如果我們要嘗試搭建 int8 推理的話，建議咱們可以先做一些 benchmark，以便更好地從性能和開發(fā)難易程度做取舍。

在 FasterTransformer 中我們采用了 IMMA-specific 布局。所以接下來，我們以 IMMA-specific 布局為例，簡單介紹了一下 cublasLt int8 矩陣乘法的基本搭建流程，以及一些開發(fā)技巧。

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

• 首先我們需要創(chuàng)建句柄和乘法描述符；
• 接下來我們?yōu)槊總€矩陣創(chuàng)建一個矩陣描述符；
• 因?yàn)橐话阄覀兊妮斎攵际浅Ｒ?guī) layout 的，所以我們需要對常規(guī)布局的矩陣進(jìn)行布局轉(zhuǎn)換，使其變?yōu)?IMMA-specific 的布局；
• 然后再進(jìn)行 int8 矩陣乘法，得到結(jié)果之后，我們可以考慮繼續(xù)用這個結(jié)果進(jìn)行下游的乘法計(jì)算，這樣可以避免轉(zhuǎn)變會常規(guī)布局的開銷；
• 只有最后一個矩陣乘法的結(jié)果，我們需要轉(zhuǎn)換常規(guī)布局以便輸出。

上述介紹了 IMMA-specific 布局下的搭建流程，可以看到里面會有不少限制。為了避免這些限制對性能的影響，我們在 Faster Transformer 中采用了以下技巧：

• 首先 IMMA-specific 布局對矩陣是有特定的尺寸要求，為了避免推理過程中需要額外分配空間的操作，我們會提前分配好符合 IMMA-specific 布局尺寸的 buffer；
• 然后，由于 weight 可以一次處理重復(fù)使用，所以我們會提前對 weight（相當(dāng)于乘法中的 B 矩陣）進(jìn)行布局變換，避免在推理過程中反復(fù)變換 weight；
• 第三個技巧是，對于不得不進(jìn)行特殊布局變換的 A 和 C，我們會把變換和上游或下游 op 進(jìn)行算子融合，以便隱藏這部分的開銷；
• 最后一點(diǎn)，是與布局無關(guān)，而是 int8 矩陣乘法必有的量化和反量化的操作，我們同樣會采用算子融合的方式，把它的 latency 隱藏起來。

以下是我們在 Faster Transformer 中采用的的 INT8 流程的示意圖，可以看到，所有矩陣乘都變?yōu)榱?int8 數(shù)據(jù)類型，每個 int8 矩陣乘法前后都會插入對應(yīng)的量化和反量化節(jié)點(diǎn)，然后對于加 bias，加殘差或 layernorm 等操作，我們還是保留原始的 FP32 或 FP16 數(shù)據(jù)類型，當(dāng)然它的 I/O 可能是 int8 的，從而會比 FP16 或 FP32 I/O 性能要好。

這里展示的是 Swin Transformer int8 量化的精度情況，通過 QAT 我們可以保證精度損失在千分之 5 以內(nèi)。

而在 PTQ 那一列，我們可以看到 Swin-Large 的掉點(diǎn)比較嚴(yán)重，一般對應(yīng)掉點(diǎn)嚴(yán)重的問題，我們都可以考慮采用減少一些量化節(jié)點(diǎn)的方式來提升量化精度，當(dāng)然這樣可能會帶來加速效果的減弱。

在 FT 中，我們可以通過禁用 FC2 和 PatchMerge 中 int 8 矩陣乘法的 int8 輸出前的反量化和量化結(jié)點(diǎn)（即采用 int32 輸出），來進(jìn)一步提升量化精度，可以看到在此優(yōu)化操作下，swin-large 的 PTQ 精度也明顯提升了。

接下來是我們推理側(cè)取得的加速效果，我們分別在不同型號的 GPU T4、A10、A100 上進(jìn)行了跟 pytorch FP16 實(shí)現(xiàn)的性能對比。

其中下圖左側(cè)是優(yōu)化后跟 pytorch 的 latency 對比，右圖為優(yōu)化后 FP16 下跟 pytorch 以及 INT8 優(yōu)化跟 FP16 優(yōu)化的加速比。可以看到，通過優(yōu)化，在 FP16 精度上，我們可以取得，相對于 pytorch 2.82x ~ 7.34x 的加速，結(jié)合 INT8 量化，我們可以在此基礎(chǔ)上進(jìn)一步取得 1.2x ~ 1.5x 的加速。

4. Swin Transformer 優(yōu)化總結(jié)

最后，我們總結(jié)一下，本次分享中我們介紹了如何通過 nsight system 性能分析工具發(fā)現(xiàn)性能瓶頸，然后針對性能瓶頸，介紹了一系列訓(xùn)練推理加速技巧，其中包括 1. 混合精度訓(xùn)練 / 低精度推理，2. 算子融合，3. cuda kernel 優(yōu)化技巧 ：如矩陣補(bǔ)零，向量化讀寫，巧用寄存器數(shù)組等，4. 推理優(yōu)化上采用一些預(yù)處理，來完善我們的計(jì)算流程；我們也介紹了 multi-stream，cuda graph 的一些應(yīng)用 。

結(jié)合上述優(yōu)化，我們在訓(xùn)練上，以 Swin-Large 模型為例取得了單卡 2.85x 的加速比，8 卡 2.32x 的加速比；在推理上，以 Swin-tiny 模型為例，在 FP16 精度下取得了 2.82x ~ 7.34x 的加速比，結(jié)合 INT8 量化，進(jìn)一步取得 1.2x ~ 1.5x 的加速比。

上述視覺大模型訓(xùn)練與推理的加速方法都已經(jīng)在百度百舸 AI 異構(gòu)計(jì)算平臺的 AIAK 加速功能中實(shí)現(xiàn)，歡迎大家使用。