从Tensor Cores恢复单精度精度并超越FP32理论峰值性能
分布式、并行与集群计算
2023-10-19 v3
摘要
Tensor Core是NVIDIA GPU上的混合精度矩阵-矩阵乘法单元，在Ampere架构上理论峰值性能超过300 TFlop/s。Tensor Cores是为响应机器学习对稠密矩阵乘法的巨大需求而开发的。然而，科学计算中的许多应用，如迭代求解器的预条件子与低精度傅里叶变换，也可利用这些Tensor Cores。要在Tensor Cores上计算矩阵乘法，我们需要将输入矩阵转换为半精度，这会导致精度损失。为避免此问题，我们可在转换中使用额外的半精度变量保留尾数损失，并用它们来修正矩阵-矩阵乘法的精度。即使有此修正，使用Tensor Cores仍比FP32 SIMT Cores具有更高吞吐。然而，该方法单独的修正能力有限，所得精度无法匹敌FP32 SIMT Cores上的矩阵乘法。我们解决了这一问题，并基于Tensor Cores开发了高精度、高性能且低功耗的矩阵-矩阵乘法实现，其精确匹配FP32 SIMT Cores的精度同时实现更优吞吐。该实现基于NVIDIA的CUTLASS。我们发现达到此精度的关键在于如何处理Tensor Cores内部的舍入以及修正计算中的下溢概率。我们的实现在使用FP16 Tensor Cores的有限指数范围下达到51TFlop/s，在使用TF32 Tensor Cores的FP32全指数范围下达到33TFlop/s（于NVIDIA A100 GPU上），优于19.5TFlop/s的FP32 SIMT Core理论峰值性能。
引用
@article{arxiv.2203.03341,
title = {Recovering single precision accuracy from Tensor Cores while surpassing the FP32 theoretical peak performance},
author = {Hiroyuki Ootomo and Rio Yokota},
journal= {arXiv preprint arXiv:2203.03341},
year = {2023}
}