比快更快:混合精度计算加速的实现

发布时间:2021-11-29 07:00

引言/背景

 

深度学习应用的发展极大推动了计算加速硬件的发展。这些新型的计算硬件也被广泛用于各类科学计算中。世界排名前列的超算summit就集成了两万多张GPU加速卡,这些加速卡与近万颗CPU构成的异构系统可以达到200PFLOPS的计算性能。

 

传统科学计算一般都在FP64等高数值精度上实现。为了满足机器学习的需求,各类计算加速卡都强化了低数值精度下的计算能力。如表1所示,在Tensor Core(TC)等技术帮助下,新一代的GPU如V100, A100等半精度(FP16)计算速度比双精度(FP64)计算速度高了一个数量级。如何在数值计算中充分利用新型加速硬件的强大低精度算力成为科学计算领域一个很现实的需求。除了计算速度提升外,低精度带来的另一个好处是减少了数据表达位宽,从而降低了计算中对设备带宽的需求,进而提升了数据通讯和存储的效率。

 

表1  GPU各浮点精度的计算能力(in TFLOPS): 英伟达V100和A100 [1]

 

因此,充分利用不同精度计算能力成为学术界和业界不可忽视的课题。我们在之前的公众号文章中已经介绍了混合精度实现的研究和进展。本文以数值计算测试中常用的线性方程求解为例简要分析混合精度计算实现的一般方法。事实上这些实现已经集成在了各类数学计算库中。设备厂商如英伟达等开发了如cuBLAS等软件库,提供了对自身硬件多精度上数学运算的支持。开源库如MAGMA、LAPACK也提供了针对各类加速卡优化的混合精度实现。

 

混合精度计算的实现

 

阻碍低精度表达在科学计算中应用的重要障碍是低精度带来的计算误差。通过研究,学者们发现计算过程中并非所有环节都需要高精度数值格式。通过合理分配不同环节的数值精度,我们就可在保证最终计算结果准确度的条件下,实现整体计算的加速。

 

科学计算中最常见的数值计算问题是线性方程求解。给全球超级计算中心排名的TOP500 List采用了LINPACK Benchmark作为测试标准。LINPack的核心就是测试机器对线性方程求解的计算能力。我们以线性方程求解来简单理解混合精度计算的实现。

 

混合精度线性方程求解

 

众所周知,由于量化误差,线性方程Ax=b常通过迭代法求解,即先计算解的初值x0,再计算方程误差r=b-Ax0并求解Az=r,计算更新的方程解x=x+z。反复执行上述迭代,直到满足误差条件。

对于大型线性方程,我们常采用LU分解法求解。该方法先将矩阵A分解为下三角阵L和上三角阵U,再对原方程求解。整个步骤可分解为:

 

 

尽管看起来复杂了,但LU方法将计算做了分解。其中第一步分解的计算复杂度为O(n3),而后续回代求解的两步的计算复杂度均为O(n2)。更重要的是,计算量最大的矩阵分解过程可以在低精度上执行。

 

综上所述,一个混合精度线性方程求解算法可以描述如下:

 

 

上述描述中的矩阵P代表矩阵分解中的行交换操作。这个算法红色部分的计算都在单精度上实现,只有蓝色部分才在更高的双精度上执行。实践证明,尽管迭代次数由于低精度的引入可能会增加,但由于低精度带来的计算速度和带宽的大幅提升,整体求解效率仍然可以提高。值得注意是该方法中精度转换操作也会带来一些执行的代价和性能的影响,本文不再展开。

 

混合精度矩阵分解

 

进一步分析发现,矩阵分解如LU或QR本身也可以用混合精度实现,比如可以用FP16-32混合精度实现FP32矩阵的分解。

 

 

图1 矩阵LU分解

 

如图1所示,大型矩阵的LU分解可由一系列的两阶段(panel分解,后缘更新)三步计算迭代完成:panel分解(橙色),U阵更新(紫色),剩余A阵更新(绿色)。LINPACK benchmark就要求测试者按照LU求解的方法来测试系统算力。通过分析我们可以发现,该分解中A阵更新的计算代价最大,复杂度为O(n3),消耗了主要的计算量。经过实验,研究者发现该更新可以用更低精度完成。一种可行的混合精度LU分解算法如下:

 

矩阵A 按图1分割, Panel从左到右,迭代执行:

 

 

其中红色部分以矩阵A表达精度更低的数值精度执行。对于FP32精度的A矩阵,我们可以用FP16来执行L,U矩阵乘法,结果为FP32,用于更新FP32的A阵。

 

结合上一小节介绍的迭代求解方法,我们就可以实现FP16-FP32-FP64三级精度的混合精度线性方程求解方法。其中计算复杂度最高的部分可以在FP16上执行,从而充分利用计算加速设备上强大的半精度计算能力,提升数值计算的整体效率。

 

混合精度计算性能

 

运用上述方法,研究者们在MAGMA等开源或商业计算软件包中集成了各种混合精度的计算加速方法。如图所示,同基于固定FP64精度的方法对比,利用英伟达V100 tensor Core FP16加速器(FP16-TC)混合精度LU分解实现的计算速度增加了三至四倍。

 

图2 混合精度LU分解的测试(Fig 5 in [2])

 

线性方程求解测试中,他们发现混合精度实现仅增加数次迭代优化就可以使结果达到相同计算准确度。尽管混合精度方法求解需要多次迭代优化以达到FP64方法的精度,但计算量最大的LU分解仅需执行一次,而单次迭代计算量的overhead仅为约2%。因此同FP64方法相比,基于FP16-TC的计算速度提升也几乎能达到四倍。

 

此外,由实验结果可知,矩阵尺寸越大,混合精度方法线性方程求解的计算加速效果越显著。

 

图3 混合精度线性方程求解性能测试(Fig 6 in [2])

 

混合精度运算的另一个优势是能耗效率。文[3]统计了实数和复数线性求解任务中不同精度实现时计算机系统中CPU和GPU(单个V100)的总能耗。对于相同的大型矩阵求解(n=40000)任务,由于总运算时间减少,利用了Tensor Core FP16能力的FP16-TC方法的能耗仅为FP64方法的1/4 – 1/3。

 

 

图4 混合精度计算的能耗测试(Fig 16 in [3])

 

HPL-AI:混合精度计算Benchmark

 

如前文所述,集成线性方程求解的LINPack是计算机系统性能测试的重要标准。目前世界超算TOP500排名正是依据其报告的每秒浮点运算次数(FLOPS)数据而定。然而LINPack所代表的指标多关注传统科学计算的双精度计算能力,不能准确地代表包含强大低精度计算能力的新型硬件加速器。比如Summit超算就集成了大量英伟达V100。这种卡拥有远超FP64算力的半精度算力。为此美国橡树岭国家实验室(ORNL)的一些学者于2019年提出了HPL-AI指标[4]来评测超算的混合精度计算能力。

 

同LINPACK benchmark一样,HPL-AI也用线性方程求解来评测系统的混合精度计算性能。所不同的是,它要求使用混合精度的LU分解与迭代优化方法来实现64位浮点精度计算的准确度。为统一比较,该指标对LU分解的计算复杂度和最终求解的误差精度设定了具体的标准。

 

测试结果每半年更新一次,目前已经有超过10家全球的超算中心参加了测试,多集中在欧美和日本,其中就包括了目前TOP500前两位的富岳和Summit。

 

最新一轮(2021年11月)公布的结果中,富岳的HPL-AI达到了2.0E Flops,其混合精度和双精度计算能力均高于Summit的指标。但值得一提的是,Summit的混合精度和双精度计算指标之比是7.7倍,高于富岳的4.5倍。 Summit较高的混合精度计算加速一个可能原因是二者架构上的区别。富岳超算集成由同构的Arm架构CPU构成。该CPU的FP16计算能力是其双精度的4.5倍。Summit是CPU与GPU的混合结构,其GPU可以用Tensor core提供高一个数量级的FP16计算能力。

 

另一个有趣的地方是混合精度计算表中第三位的Selene的FP64计算能力弱于第四位的Perlmutter,但是混合精度计算能力反而超出。这一现象在表中还有多处,有兴趣的读者可查看全部排名。

 

Rank

Site

Computer

Cores

HPL-AI (Eflop/s)

TOP500 Rank

HPL Rmax (Eflop/s)

Speedup

1

RIKEN

Fugaku

7,630,848

2.000

1

0.4420

4.5

2

DOE/SC/ORNL

Summit

2,414,592

1.411

2

0.1486

9.5

3

NVIDIA

Selene

555,520

0.630

6

0.0630

9.9

4

DOE/SC/LBNL

Perlmutter

761,856

0.590

5

0.0709

8.3

5

FZJ

JUWELS BM

449,280

0.470

8

0.0440

10.0

表2 HPL-AI排名(2021.11)[4]

 

结论与思考

 

本文简述了在GPU等多精度计算加速设备上实现混合精度计算的方法,讨论了混合精度计算对超算系统在科学计算上性能的影响。本文还介绍了混合精度计算benchmark和发展情况。随着计算加速硬件日新月异的发展,我们相信混合精度计算将取得越来越多的进展并应用到更多的科学研究和工程实施领域。

 

参考文献

 

[1] NVIDIA Ampere Architecture In-Depth, at https://developer.nvidia.com/blog/nvidia-ampere-architecture-in-depth/

[2] A. Abdelfattah, et.al., A Survey of Numerical Linear Algebra Methods Utilizing Mixed Precision Arithmetic,  Intel. J. HPC App. 35(4), 2021

[3] A. Haidar, et.al., Mixed-precision iterative refinement using tensor cores on GPUs to accelerate solution of linear systems, Proc. Royal Soci. A, 476(2243), 2020.

[4] HPL-AI Mixed-Precision Benchmark, at https://hpl-ai.org/

上一个: 视觉神经场:NeRF研究的新视角

下一个: GAN在因果推理中的应用研究

近期文章

AI 智能体:应用前景与对算力需求的影响

人工智能技术的迅猛发展,尤其是 AI 智能体这一核心部分。本文讨论了AI智能体的定义、应用前景和其对算力需求的影响。

2023-11-13

查看更多