共计 7561 个字符，预计需要花费 19 分钟才能阅读完成。
作者：陈振寰 | 旷视科技 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
参考
- [1]NVIDIA TURING GPU ARCHITECTURE White Paper
- [2]warp matrix functions – Programming Guide :: CUDA Toolkit Documentation (nvidia.com)（wmma CUDA API）
- [3] 总体流程设计（1）-CUDA 程序的等级构造
- [4]J. Burgess,“RTX on – The NVIDIA turing GPU,”IEEE Micro, vol. 40, no. 2, pp. 36–44, 2020.
- [5]cudnnConvolutionForward – API Reference :: NVIDIA Deep Learning cuDNN Documentation（convforward CuDNN API）
