深入理解混合精度训练：从 Tensor Core 到 CUDA 编程

背景

近年来，自动混合精度（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 内会包含各种计算核心以及存