从Tensor Core恢复单精度精度,同时超越FP32理论峰值性能
Recovering single precision accuracy from Tensor Cores while surpassing the FP32 theoretical peak performance
AI总结:
针对Tensor Core矩阵乘法精度不足问题,提出基于CUTLASS的校正方法,处理舍入与下溢,实现与FP32 SIMT Core相同精度且吞吐量超越其理论峰值。
摘要:
Tensor Core是NVIDIA GPU上的混合精度矩阵乘法单元,在Ampere架构上理论峰值性能超过300 TFlop/s。Tensor Core是为了响应机器学习对稠密矩阵乘法的高需求而开发的。然而,科学计算中的许多应用,如迭代求解器的预条件子和低精度傅里叶变换,也可以利用这些Tensor Core。为了在Tensor Core上计算矩阵乘法,我们需要将输入矩阵转换为半精度,这会导致精度损失。为了避免这种情况,我们可以使用额外的半精度变量来保留转换过程中的尾数损失,并用它们来校正矩阵乘法的精度。即使进行了这种校正,与FP32 SIMT Core相比,使用Tensor Core仍能获得更高的吞吐量。然而,这种方法单独校正能力有限,所得精度无法与FP32 SIMT Core上的矩阵乘法相匹配。我们解决了这个问题,开发了一种使用Tensor Core的高精度、高性能、低功耗的矩阵乘法实现,其精度与FP32 SIMT Core完全匹配,同时实现更高的吞吐量。该实现基于NVIDIA的CUTLASS。我们发现,实现这一精度的关键在于如何处理Tensor Core内部的舍入以及校正计算过程中的下溢概率。我们的实现在NVIDIA A100 GPU上,使用FP16 Tensor Core在有限指数范围内达到51TFlop/s,使用TF32 Tensor Core在FP32的完整指数范围内达到33TFlop/s,超过了FP32 SIMT Core理论峰值性能19.5TFlop/s。
英文摘要:
Tensor Core is a mixed-precision matrix-matrix multiplication unit on NVIDIA GPUs with a theoretical peak performance of more than 300 TFlop/s on Ampere architectures. Tensor Cores were developed in response to the high demand of dense matrix multiplication from machine learning. However, many applications in scientific computing such as preconditioners for iterative solvers and low-precision Fourier transforms can exploit these Tensor Cores. To compute a matrix multiplication on Tensor Cores, we need to convert input matrices to half-precision, which results in loss of accuracy. To avoid this, we can keep the mantissa loss in the conversion using additional half-precision variables and use them for correcting the accuracy of matrix-matrix multiplication. Even with this correction, the use of Tensor Cores yields higher throughput compared to FP32 SIMT Cores. Nevertheless, the correcting capability of this method alone is limited, and the resulting accuracy cannot match that of a matrix multiplication on FP32 SIMT Cores. We address this problem and develop a high accuracy, high performance, and low power consumption matrix-matrix multiplication implementation using Tensor Cores, which exactly matches the accuracy of FP32 SIMT Cores while achieving superior throughput. The implementation is based on NVIDIA's CUTLASS. We found that the key to achieving this accuracy is how to deal with the rounding inside Tensor Cores and underflow probability during the correction computation. Our implementation achieves 51TFlop/s for a limited exponent range using FP16 Tensor Cores and 33TFlop/s for full exponent range of FP32 using TF32 Tensor Cores on NVIDIA A100 GPUs, which outperforms the theoretical FP32 SIMT Core peak performance of 19.5TFlop/s.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。