作者：陳振寰 | 曠視科技 MegEngine 架構師
背景
近年來，自動混合精度（Auto Mixed-Precision，AMP）技術在各大深度學習訓練框架中作為一種使用簡單、代價低廉、效果顯著的訓練加速手段，被越來越廣泛地應用到算法研究中。然而大部分關於混合精度訓練的文章一般停留在框架接口介紹、如何避免 FP16 類型帶來的精度損失以及如何避免出現 NaN 等基礎原理和使用技巧方面，對於將深度學習框架視為黑盒工具的研究員來說確實足夠了，但是如果想要再往下多走一步，了解一點更底層的加速細節，那么 GPU 顯卡架構、CUDA 編程里的一個個專業名詞就很容易讓缺乏背景知識的人摸不着頭腦。
本文會以混合精度訓練背后涉及的 Tensor Core 為起點，結合代碼實例，幫助讀者對框架層面使用 Tensor Core 進行訓練加速的細節乃至 CUDA 編程有一些基本的認識。
Tensor Core 原理
首先還是簡單介紹一下 混合精度 和 Tensor Core 是什么。混合精度是指在底層硬件算子層面，使用半精度（FP16）作為輸入和輸出，使用全精度（FP32）進行中間結果計算從而不損失過多精度的技術，而不是網絡層面既有 FP16 又有 FP32。這個底層硬件層面其實指的就是 Tensor Core，所以 GPU 上有 Tensor Core 是使用混合精度訓練加速的必要條件。
圖 1
Tensor Core 直譯為張量核心，其物理含義是 NVIDIA GPU 上一塊特殊的區域（如圖 2 中大塊深綠色部分所示），與其地位類似的有普通的 CUDA Core（淺綠色和小塊深綠色部分）以及最新的 RT Core（Ray Tracing，光追核心，淺黃色部分）。CUDA Core 一般包含多個數據類型，每個數據類型包含多個小核心，比如圖中的 INT32 Core 和 FP32 Core 就各有 4×16 個，在計算專用卡上還可能會包含 FP64 Core（比如 V100 和 A100 顯卡），而 Tensor Core 在架構圖和接口上則沒有具體的區分，可以視作 GPU 上一塊較為獨立的計算單元（雖然實際內部有一定的區分）。
圖 2：Turing 架構 2080Ti 顯卡的 SM 圖 1
而在邏輯（數學）含義上，相比於 FP32 Core 一次只能對兩個數字進行計算（如圖 3 中兩張圖的左側部分），Tensor Core 能一次對兩個 4×4 的 FP16 Tensor 進行矩陣乘計算並累加到另一個 4×4 的 Tensor 上，即 D = A * B + C（如圖 3 中兩張圖的右側部分），這也是其取名為 Tensor Core 的原因。通過硬件上的特殊設計，Tensor Core 理論上可以實現 8 倍於 FP32 Core 的計算吞吐量（Volta 和 Turing 架構），並且沒有明顯的占用面積和功耗增加。混合精度也是利用 Tensor Core 的這一特性，才能夠實現訓練加速。
圖 3
這里的 8 倍是基於 SM（Streaming Multiprocessor）進行比較的，SM 是 GPU 進行完整計算任務的基本單元，一個 GPU 內一般包含若干 SM（比如 V100 包含 80 個 SM，A100 包含 128 個 SM），而一個 SM 內會包含各種計算核心以及存儲資源（圖 2 就是一個完整的 SM）。
在 V100 上，一個 SM 包含 8 個 Tensor Core 和 64 個 FP32 Core。一個 Tensor Core 一個時鍾周期內能進行 4×4×4=64 次 FMA（Fused-Multiply-Add，乘加計算），總計 64×8/clock，而 FP32 Core 則是 1×64/clock，故而為 8 倍。而在 A100 上，Tensor Core 的單個吞吐能力是上一代的 4 倍，一個時鍾周期能進行 256 次 FMA，在總個數減少為 4 個（占用面積更大）的情況下，總吞吐量相比 V100 提升為 2 倍，是 FP32 Core 的 16 倍。
從 CUDA 接口層面理解
前面介紹了 TensorCore 的物理含義和邏輯含義，但是還是有點抽象，所謂 “Talk is cheap. Show me the code.” 接下來就讓我們從代碼接口層面了解一下 Tensor Core 的含義與作用，從而理解混合精度訓練的底層加速原理。
要利用 Tensor Core 進行計算，需要使用 NVIDIA 提供的 CUDA Runtime API。既然在 Volta 架構中引入了 Tensor Core，那必然會有新的 CUDA 接口暴露出來。在 CUDA 9.0 中，引入了新的 WMMA（warp-level matrix multiply and accumulate）API，作用就是使用 Tensor Core 進行矩陣運算，與本文相關的主要是以下三個接口：
void load_matrix_sync(fragment<...> &a, const T* mptr, unsigned ldm, layout_t layout);
void store_matrix_sync(T* mptr, const fragment<...> &a, unsigned ldm, layout_t layout);
void mma_sync(fragment<...> &d, const fragment<...> &a, const fragment<...> &b, const fragment<...> &c, bool satf=false);
這里的 fragment 可以簡單理解為一個矩陣或 Tensor，三個接口的作用是通過 load_matrix_sync 將數據指針 mptr 里的數據加載到 fragment 中，再用 mma_sync 對四個 fragment 進行計算（d = a * b + c），最后通過 store_matrix_sync 將輸出 fragment 的數據返回到輸出指針 mptr 里。一個最簡單的對兩個 16×16 矩陣進行乘法並累加的例子如下所示2：
#include <mma.h>
using namespace nvcuda;
__global__ void wmma_ker(half *a, half *b, float *c) {
// Declare the fragments
wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::col_major> a_frag;
wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::row_major> b_frag;
wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag;
// Initialize the output to zero
wmma::fill_fragment(c_frag, 0.0f);
// Load the inputs
wmma::load_matrix_sync(a_frag, a, 16);
wmma::load_matrix_sync(b_frag, b, 16);
// Perform the matrix multiplication
wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
// Store the output
wmma::store_matrix_sync(c, c_frag, 16, wmma::mem_row_major);
}
但是到這里其實累積了一些問題，包括 warp 是什么意思？說好的 Tensor Core 接收 4×4 矩陣進行乘加，到這里為什么變成了 16×16？其實這都涉及到 GPU 進行並行計算的方式。
我們都知道 GPU 有非常多核心，比如一個 SM 里就有 64 個 FP32 Core。在管理這些核心時，為了提升效率，會將其進行分組，若干個核心在行為上進行綁定，執行一樣的命令，共同進退，而這樣的一個分組就稱為一個 warp（與 thread 相對應，都是紡織中概念的延伸 3）。在 CUDA 層面要得到一個多線程同步的結果必須以 warp 為單位，這也是上面三個函數都以"_sync"結尾的原因。
在硬件上其實也可以找到這種分組的跡象，比如我們再看上面 Turing SM 的結構（圖 2），可以發現其分為了四個一樣的部分（如下圖 4），稱作 Sub-Core，其中橙色的部分叫作 “Warp Scheduler”，其作用就是給 warp 分配任務。
圖 4：Turing 架構 SM 的 一個 Sub-Core
而分配任務一個時鍾周期只能進行一次，為了盡量讓各個部分都能一直運轉，這個任務一般需要多個時鍾周期執行（類似流水線並行）。在目前的 GPU 設計中，一個 warp scheduler 對應 32 個線程，可以理解為一個任務包含 32 個子任務，而每個 Sub-Core 只有 16 個 FP32 Core，所以需要兩個時鍾周期才能分配一次。
圖 5 Turing 架構 Sub-Core 里的指令流程4
對應到 Tensor Core 上算 FP16 的矩陣乘加，如果是 Volta 架構，一次會算 8×4 和 4×8 兩個矩陣的乘法和對應矩陣的累加（縮寫為 m8n8k4），需要 4 個時鍾周期才能分配一次，均攤下來一個時鍾周期恰好是兩個 4×4 矩陣的乘加，與宣稱的 TensorCore 性能一致。而實際上在 CUDA Runtime API 里，為了使指令 overlap 更高，提升並行效率，把這個 m8n8k4 提升為了最少 m16n16k16，這也就是為何 wmma::mma_sync 以 16×16 為最小單元了。
事實上 CUDA 里進行矩陣計算，往往都是把大的矩陣切分成一個個固定大小的分塊（tiling）進行計算，這其實也是接口的輸入叫 fragment （每個線程的 fragment 負責 tiling 的一部分）而非 Tensor 的一個原因。
總結下來就是 CUDA 通過 wmma 接口以 warp 為單位每 4 個時鍾周期向 Tensor Core 提交 m8n8k4 矩陣乘法的運算請求，待其執行完成后把 8×8 的結果進行返回，整個運算的過程都是基於 warp 層面的，即 warp-level。
到這里我們通過 wmma CUDA API 了解了 Tensor Core 的代碼含義，以及管中窺豹挖掘了一下硬件底層執行的流程，其中的用語和描述為了方便理解也許不是特別准確，不過相信能幫助大家對 GPU 如何執行並行計算有一個簡單的認識。
從框架使用層面進行理解
實際在框架層面一般不會直接基於 CUDA 接口來調用 Tensor Core 進行計算，而是基於 CuDNN 這一現成的 DNN 算子庫，一方面是因為 CuDNN 本身隱藏了很多硬件細節，可以保證在不同顯卡之間的兼容性（比如無論是否支持 Tensor Core 都可以運行），另一方面 CuDNN 的實現在大部分常見情況下是性能足夠的，也就無需重復造輪子。
下面我們以混合精度訓練中最常用的卷積操作來介紹一下計算過程，我們先看一下 CuDNN 里的卷積操作 API5：
cudnnStatus_t cudnnConvolutionForward(
cudnnHandle_t handle,
const void *alpha,
const cudnnTensorDescriptor_t xDesc,
const void *x,
const cudnnFilterDescriptor_t wDesc,
const void *w,
const cudnnConvolutionDescriptor_t convDesc,
cudnnConvolutionFwdAlgo_t algo,
void *workSpace,
size_t workSpaceSizeInBytes,
const void *beta,
const cudnnTensorDescriptor_t yDesc,
void *y)
這里面有一些名詞需要解釋一下：
- cudnnStatus_t，CuDNN 的接口一般采用在參數里包含輸出指針（比如這里的 y）進行結果寫入的設計，而返回值只包含成功失敗的狀態信息，即 status。
- cudnnHandle_t，handle 是與設備進行溝通的接口，類似的概念還有 file handle，直譯為句柄，任何接口都需要提供一個 cuda device 的 handle。
- cudnnTensorDescriptor_t 和 cudnnFilterDescriptor_t，都屬於數據描述符，包含 layout、dtype 等所有數據屬性信息，因為數據內容只由一個 void* 指針（比如這里的 x 和 w）提供。
- cudnnConvolutionDescriptor_t，操作描述符，與數據描述符類似，用於描述 Op 本身的一些參數和屬性，比如 conv 就包括 pad、stride、dilation 等。
- cudnnConvolutionFwdAlgo_t，直譯是前向卷積的算法，因為卷積操作的具體計算方式多種多樣，各自有其適合的數據場景，所以需要在這里指定采用什么算法。
- workSpace，相比於上層代碼可以隨時隨地創建數據對象，在設備層，一個計算需要的空間必須事前聲明，而 workspace 就是除了輸入輸出之外，進行這個計算所需的額外“工作空間”，也可以簡單理解為空間復雜度。
在看完 API 的參數介紹之后，其實如何使用這個接口進行計算也就自然明了了，我們不准備一步步教你如何用現成的接口填上這些內容，而是想讓你思考一下，你覺得這些參數之間的邏輯關系是什么，具體來說，你覺得什么參數能夠決定這個卷積操作是運行在 Tensor Core 上的呢？
首先我們結合前面 CUDA Runtime API 的接口進行分析，wmma 接口限制了矩陣的形狀都是 16×16，以及輸入數據都是 half 半精度類型（累加器 c 可以是 float），那么與數據相關的 x/w/y 的描述符必然是有影響的（數據指針本身沒有信息所以不影響），所以我們需要在數據描述符里指明數據類型為半精度，且需要數據的各個維度都是 8 的倍數（之所以不是 16 的倍數是因為內部實現還會做一些處理）。
然后我們分析卷積算子本身，就算數據類型和維度符合要求，也完全可以使用普通的 CUDA Core 進行運算，那么可以推斷出必然有控制算子行為的參數，對照上面的列表，不難猜出是操作描述符和算法兩個參數。對於算法，我們一般認為是運算的邏輯，而與實際運算的設備無關（比如一個算法在 GPU、CPU 上應該是同樣的流程），但是設備會限制能夠運行的算法。事實上，對於 NCHW 的二維卷積操作，FFT、GEMM、WINOGRAD 等算法都支持基於 Tensor Core 或 FP32 CUDA Core 的計算，但是有些算法則只能在 CUDA Core 上進行。
所以真正控制是否使用 Tensor Core 的參數就呼之欲出了，就是 Conv 的操作描述符。事實上，除了一般意義上的 param 參數比如 pad、stride、dilation，有一個重要參數 mathType 也包含在操作描述符內，這個參數的默認值是 CUDNN_DEFAULT_MATH，而如果要使用 Tensor Core 進行運算，必須要修改成 CUDNN_TENSOR_OP_MATH，從名字上看也是一個與 Tensor Core 強相關的值。
除此之外，還有一個參數值得一提，我們都知道混合精度訓練的重要特性是 FP16 的運算中間結果使用 FP32 存儲，直到最后才轉成 FP16，從而使得精度不會明顯下降，但是這其實不是 Tensor Core 的限制，Tensor Core 完全可以全程 FP16 運算，所以要實現混合精度，也需要我們在操作描述符內進行控制，這個參數就是操作描述符 convDesc 的 dataType 屬性，我們需要將其設置成單精度（CUDNN_DATA_FLOAT）而非半精度（CUDNN_DATA_HALF）才能實現保持精度的目的。
最后簡單看一下 convDesc 相關的設置代碼：
// 創建描述符
checkCudnnErr( cudnnCreateConvolutionDescriptor( &cudnnConvDesc ));
// 設定常見參數，包括 dataType（最后一項）
checkCudnnErr( cudnnSetConvolutionNdDescriptor(
cudnnConvDesc,
convDim,
padA,
convstrideA,
dilationA,
CUDNN_CONVOLUTION,
CUDNN_DATA_FLOAT) );
// 設置 mathType
checkCudnnErr( cudnnSetConvolutionMathType(cudnnConvDesc, CUDNN_TENSOR_OP_MATH) );
至於剩下的 workspace，其實是與前面所有參數都相關的，因為必須知道數據的屬性、計算的算法、算子的屬性和計算行為等所有實際計算所需的信息，才能得出所需的“工作空間”大小，這里就不過多介紹了。
綜上可以看出 NVIDIA 在接口的設計上還是非常老道的，簡明合理的參數設計使得我們可以在較高的抽象層次上控制底層硬件的計算邏輯。而通過分析接口設計上的邏輯，我們也對一個算子如何才能利用 Tensor Core 進行混合精度計算有了較為完整的理解。
總結
Tensor Core 作為混合精度訓練賴以加速的底層硬件支持，一直在大部分框架用戶或者說算法研究員眼中好似“雲霧山中人”，了解一些數學上的含義但又不清楚細節。本文則先從物理含義上將其與實際可見的 GPU 芯片進行了關聯，再從較底層的 CUDA 接口代碼層面如何控制 Tensor Core 做矩陣運算進行了講解，最后回到框架層面實際開發角度詳細介紹了使用卷積算子進行混合精度計算的過程。
通過這些介紹，相信大家都能理解之前熟知的一些 AMP 使用限制是為何存在了，比如為何我的顯卡沒有加速效果（必須要 Volta 架構及以上），為何要求維度都是 8 的倍數（Tensor Core 里需要矩陣分塊），而更進一步的關於硬件如何決定跑 FP16 還是 FP32 的問題，相信經過上面代碼層面的講解也能有所了解。
希望本文能讓從未接觸過 CUDA 編程的讀者能更加深入理解混合精度訓練的底層運算原理，也能對 GPU 計算和 CUDA 編程有一些簡單的認識。
附：
- GitHub：MegEngine 天元
- 官網：MegEngine-深度學習，簡單開發
- 歡迎加入 MegEngine 技術交流 QQ 群：1029741705
參考
- [2]warp matrix functions - Programming Guide :: CUDA Toolkit Documentation (nvidia.com)（wmma CUDA API）
- [4]J. Burgess, “RTX on - The NVIDIA turing GPU,” IEEE Micro, vol. 40, no. 2, pp. 36–44, 2020.