大家好,我是來自NVIDIA GPU計算專家團隊的陶礪,很高興今天有機會在這里跟大家分享一下我和我的同事陳庾,在 Swin Transformer 這個視覺大模的型訓練和推理優化上的一些工作。其中一些的方法與策略,在其他的模型訓練、推理的優化上都可以使用,來提高模型的吞吐、提升 GPU 的使用效率、加快模型的迭代。
我會介紹 Swin Transformer 模型的訓練部分的優化,在推理優化部分的工作,將由我的同事來做詳細的介紹
這里是我們今天分享的目錄,主要分為四個部分,既然是針對特定模型進行的優化,那么我們首先會簡單介紹一下 Swin Transformer 模型。然后,我會結合 profiling 的工具,也就是 nsight system 對訓練的流程進行分析和優化。在推理部分,我的同事會給出推理優化的策略和方法,包含較為細節的 cuda 層面的優化。最后,是今天優化內容的一個總結。
首先是第一部分,也就是 Swin Transformer 的介紹。
1. 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 計算專家團隊的陳庾,非常感謝陶礪在訓練加速上的介紹,接下來由我來介紹一下推理上的加速。
跟訓練一樣,推理的加速離不開算子融合這一方案。不過相對于訓練而言,在推理上進行算子融合有更好的靈活性,主要體現有兩點:
推理上的算子融合不需要考慮反向,所以 kernel 開發過程中不需要考慮保存計算梯度所需要的中間結果;
推理過程允許預處理,我們可以對一些只需要一次計算便可重復使用的操作,提前算好,保留結果,每次推理時直接調用從而避免重復計算。
在推理側,我們可以進行不少的算子融合,這里給出的是我們在 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 步:
首先我們需要創建句柄和乘法描述符;
接下來我們為每個矩陣創建一個矩陣描述符;
因為一般我們的輸入都是常規 layout 的,所以我們需要對常規布局的矩陣進行布局轉換,使其變為 IMMA-specific 的布局;
然后再進行 int8 矩陣乘法,得到結果之后,我們可以考慮繼續用這個結果進行下游的乘法計算,這樣可以避免轉變會常規布局的開銷;
只有最后一個矩陣乘法的結果,我們需要轉換常規布局以便輸出。
上述介紹了 IMMA-specific 布局下的搭建流程,可以看到里面會有不少限制。為了避免這些限制對性能的影響,我們在 Faster Transformer 中采用了以下技巧:
首先 IMMA-specific 布局對矩陣是有特定的尺寸要求,為了避免推理過程中需要額外分配空間的操作,我們會提前分配好符合 IMMA-specific 布局尺寸的 buffer;
然后,由于 weight 可以一次處理重復使用,所以我們會提前對 weight(相當于乘法中的 B 矩陣)進行布局變換,避免在推理過程中反復變換 weight;
第三個技巧是,對于不得不進行特殊布局變換的 A 和 C,我們會把變換和上游或下游 op 進行算子融合,以便隱藏這部分的開銷;
最后一點,是與布局無關,而是 int8 矩陣乘法必有的量化和反量化的操作,我們同樣會采用算子融合的方式,把它的 latency 隱藏起來。
以下是我們在 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 加速功能中實現,歡迎大家使用。