百度智算峰會精彩回顧：視覺大模型訓練與推理優(yōu)化

在昨日舉行的 “2022 百度云智峰會·智算峰會” 上，百度智能云發(fā)布國內首個全棧自研的 AI 基礎設施 “百度 AI 大底座”，引發(fā)行業(yè)的廣泛關注。其中，百度百舸·AI 異構計算平臺作為 AI IaaS 層的重要構成，是 AI 大底座堅實的基礎設施。百度百舸·AI 異構計算平臺整合了百度自研的 AI 芯片 “昆侖芯”，在 AI 計算、存儲、加速、容器方面進行系統(tǒng)優(yōu)化，發(fā)揮算力的最佳性能和效率。

在本次峰會上，NVIDIA GPU 計算專家陳庾分享了以 “視覺大模型訓練與推理優(yōu)化” 為題的演講，以 Swin Transformer 模型為例，介紹了 NVIDIA 在視覺大模型訓練與推理優(yōu)化上的探索和實踐經(jīng)驗，包括在 Swin Transformer 的訓練和推理上采用到的加速技巧以及對應的加速效果。這些加速方法目前已經(jīng)應用到百度百舸·AI 異構計算平臺的 AIAK 加速功能中，助力各行業(yè)客戶的 AI 業(yè)務加速落地。

以下為演講內容概要。

視覺大模型簡介：Swim Transformer

在介紹 Swin Transformer 之前，我們可以先回顧一下 Vision Transformer (ViT) 的模型結構。如左圖所示，ViT 會將一個圖像分割成一系列的圖像塊（patch），然后通過基于 Transformer 的編碼器對這一系列圖像塊進行編碼，最后得到可用于分類等任務的特征。

而 Swin Transformer，它引入了 window attention 的概念，不同于 ViT 對整個圖像進行注意力計算（attention）。Swin Transformer 會先將圖像劃分成若干個窗口（window），然后僅對窗口內部的圖像塊進行對應的注意力計算，從而減少計算量。為了彌補窗口帶來的邊界問題，Swin Transformer 進一步引入 window shift 的操作。同時為了使得模型有更豐富的位置信息，Swin Transformer 還在注意力計算時引入了 relative position bias。

在 Swin Transformer 的網(wǎng)絡結構中，我們也可以看到， Swin Transformer 被劃分為多個階段（stage），其中每個階段的分辨率是不一樣的，從而形成了一個分辨率金字塔，這樣也使得每個階段的計算復雜程度也逐漸降低。每個階段中對應若干個 transformer block。

接下來，我們從具體操作的角度來對 Swin Transformer 進行算子的劃分?？梢钥吹剑粋€ transformer block 中涉及到三大部分，第一部分是如 window shift、partition、reverse 的一些窗口相關的操作，第二部分是注意力計算，第三部分是 FFN 計算。而 attention 和 FFN 部分又可以進一步細分為若個算子，最終如圖所示，我們可以將整個模型細分為幾十個算子的組合。這樣的算子劃分，對于我們進行性能分析，定位性能瓶頸以及開展加速優(yōu)化而言，都是非常重要的。

視覺大模型 Swim Transformer 的訓練加速

接下來，介紹 NVIDIA 在 Swim Transformer 模型訓練加速上進行的優(yōu)化工作。對于 Swin Transformer 而言，我們一般采用數(shù)據(jù)并行的方式進行多 GPU 訓練，我們發(fā)現(xiàn)卡間通訊的開銷占比會相對較少，所以在這里，我們優(yōu)先對單 GPU 上的計算瓶頸進行優(yōu)化。

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

那么對于矩陣乘法而言，我們的一個優(yōu)化手段，就是充分利用 tensor core 進行加速，可以考慮直接采用 tf32 tensor core 或者在混合精度下采用 fp16 tensor core。在這里，我們采用了混合精度的方案，可以看到右圖中，矩陣乘法的 kernel 由于采用了 fp16 tensor core，其耗時占比有了明顯下降。此外，采用混合精度的方式，我們也可以在 kernel io 上取得一定的收益。整體而言，在 Swin-Large 模型上，通過混合精度的方式，我們可以取得了 1.63 倍的吞吐提升。

下一個優(yōu)化方案是算子融合。算子融合一般而言可以為我們帶來兩個好處：一個是減少 kernel launch 的開銷。如左圖所示，兩個 cuda kernel 的執(zhí)行需要兩次 launch，那樣可能會導致 kernel 之間存在一定的空隙，使得 GPU 空閑。

如果我們將兩個 cuda kernel 融合成一個 cuda kernel，一方面節(jié)省了一次 launch，同時也可以避免 gap 的產(chǎn)生。另外一個好處是減少了 global memory 的訪問，因為 global memory 的訪問是非常耗時的。而兩個獨立的 cuda kernel 之間要進行結果傳遞，都需要通過 global memory，將兩個 cuda kernel 融合成一個 kernel，我們就可以在寄存器或者 shared memory 上進行結果傳遞，從而避免了一次 global memory 寫和讀，從而提升性能。

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

除了利用現(xiàn)有的 apex 庫，NVIDIA 也進行了手工的融合算子開發(fā)。這里展示的是，我們針對 Swin Transformer 特有的 window 相關操作，如 window partition、shift、merge 等進行的融合算子開發(fā)。可以看到，通過引入我們的定制化融合算子，可以將加速比進一步提升至 2.19 倍。

接下來展示的是，NVIDIA 對 mha 部分的融合工作。Mha 部分是 transformer 模型中一個占比很大的模塊，因此對它的優(yōu)化往往可以帶來較大的加速效果。從圖中可以看到，在沒有進行算子融合之前，mha 部分的操作占比為 37.69%，其中包括了不少 elementwise 的 kernel。我們通過將相關操作融合成一個獨立的 fmha kernel，從而將加速比進一步提升 2.85 倍。

上述是我們在單卡上取得的訓練加速效果。我們再來看一下單機 8 卡的訓練情況，可以看到，通過上述優(yōu)化，我們可以將訓練吞吐從 1,612 提升至 3,733，取得 2.32 倍的加速。

視覺大模型 Swim Transformer 的推理加速

接下來，介紹 NVIDIA 在 Swim Transformer 模型推理加速上進行的優(yōu)化工作。

跟訓練一樣，推理的加速離不開算子融合這一方案。不過相對于訓練而言，在推理上進行算子融合有兩個好處：一，推理上的算子融合不需要考慮反向，所以 kernel 開發(fā)過程中不需要考慮保存計算梯度所需要的中間結果；二，推理過程允許預處理，我們可以對一些只需要一次計算便可重復使用的操作，提前算好，保留結果，每次推理時直接調用從而避免重復計算。

在推理側，NVIDIA 進行了不少算子融合，在這里挑選了兩個加速效果比較明顯的算子進行介紹。

首先是 mha 部分的算子融合，我們將 position bias lookup 這一操作提前到預處理部分，從而避免每次推理時都進行 lookup，然后將 batch gemm、softmax、batch gemm 融合成一個獨立的 fmha kernel?？梢钥吹剑?strong>融合后，該部分取得了 10 倍的加速，而端到端也取得了 1.58 倍的加速。

另一個算子融合是 QKV gemm + bias 的融合，gemm 和 bias 的融合是一個十分常見的融合手段。在這里為了配合我們前面提到的 fmha kernel，我們需要對 weight 和 bias 提前進行格式上的變換，這種提前變換也體現(xiàn)了上文提到的，推理上進行算子融合的靈活性。最后，通過 QKV gemm+bias 的融合，我們可以進一步取得 1.1 倍的端到端加速。

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

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

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

以 half2 為例，可以帶來以下三個好處：

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

二， 結合 half2 專有的高吞吐的數(shù)學指令，可以減低 kernel 的時延。

三，在進行 reduction 時，采用 half2 數(shù)據(jù)類型意味著一個 cuda 線程同時處理兩個元素，可以有效減少空閑的線程數(shù)。最終通過使用 half2，我們可以觀察到圖示的 addBiasResidual 這一個 kernel，latency 從 20.96us 下降為 10.78us，加速 1.94 倍。

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

由于寄存器是每個 cuda 線程獨占的，所以在進行 kernel 設計時，我們需要提前設定好每個 cuda 線程所需要緩存的元素個數(shù)，從而開辟對應大小的寄存器數(shù)組，并且在分配每個 cuda 線程所負責元素時，需要確保我們可以做到合并訪問。

如右圖所示，當我們有 8 個線程時，0 號線程可以處理 0 號元素，當我們有 4 個線程是，0 號線程則處理 0 號和 4 號元素，如此類推。我們一般建議可以采用模板函數(shù)的方式，通過模板參數(shù)來控制每個 cuda 線程的寄存器數(shù)組大小。

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

最后一個我想介紹優(yōu)化手段是 INT8 量化。INT8 量化是推理加速非常重要的加速手段，對于 Transformer based 的模型而言，INT8 量化往往可以帶來不錯的加速效果。

而對于 Swin Transformer 來說，通過結合合適的 PTQ 或 QAT 量化方案，可以在取得良好加速效果的同時，保證量化精度。通過 QAT，我們可以保證精度損失在千分之 5 以內。

而在 PTQ 那一列，我們可以看到 swin large 的掉點比較嚴重。我們可以通過禁用 FC2 和 PatchMerge 中的部分量化結點，來進一步提升量化精度，這其實就是一個性能和精度的 tr 平衡了。

接下來介紹 NVIDIA 在推理側取得的加速效果。

左上圖優(yōu)化后 跟 pytorch 的時延對比，右下圖為優(yōu)化后 FP16 下跟 pytorch 以及 INT8 優(yōu)化跟 FP16 優(yōu)化的加速比?？梢钥吹?，通過優(yōu)化，在 FP16 精度上，我們可以取得，相對于 pytorch 2.82x~7.34x 的加速，結合 INT8 量化，我們可以在此基礎上進一步取得 1.2x~1.5x 的加速。

總結：視覺大模型 Swim Transformer 訓練與推理加速技巧

本次分享中我們介紹了一系列訓練與推理加速技巧，其中包括 ：

一，混合精度訓練、低精度推理，使用更高性能的 Tensor Core 進行計算矩陣乘法或卷積計算；

二、算子融合；

三、cuda kernel 優(yōu)化技巧 ：如矩陣補零，向量化讀寫，巧用寄存器數(shù)組等。

結合上述優(yōu)化技巧，加速效果：

在訓練上，以 Swin-Large 模型為例，取得了單卡 ~2.85x 的加速比，8 卡 2.32x 的加速比；

在推理上，以 Swin-tiny 模型為例，在 FP16 精度下取得了 2.82x - 7.34x 的加速比，結合 INT8 量化，進一步取得 1.2x - 1.5x 的加速比。

上述介紹的視覺大模型訓練與推理的加速方法，已經(jīng)在百度百舸·AI 異構計算平臺的 AIAK 加速功能中實現(xiàn)。