比快更快:混合精度計算加速的實現

摘要
深度學習推動了計算裝置如GPU的計算能力,尤其是低精度數值計算能力的快速提升。然而傳統科學計算任務如數值天氣預報都是在FP64等高精度模式上進行。為充分利用現代計算裝置的低精度計算能力,研究者們開發了很多混合精度計算方法。本文簡要介紹混合精度計算的實現方法和在現代GPU上的表現,最後討論混合精度計算的評測指標。
引言
深度學習應用的發展極大推動了計算加速硬體的發展。這些新型的計算硬體也被廣泛用於各類科學計算中。世界排名前列的超算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的核心就是測試機器對線性方程求解的計算能力。我們以線性方程求解來簡單理解混合精度計算的實現。

混合精度線性方程求解

眾所周知,由於量化誤差,線性方程  常透過迭代法求解,即先計算解的初值 ,再計算方程誤差  並求解  ,計算更新的方程解  。反覆執行上述迭代,直到滿足誤差條件。
對於大型線性方程,我們常採用LU分解法求解。該方法先將矩陣A分解為下三角陣L和上三角陣U,再對原方程求解。整個步驟可分解為:
儘管看起來複雜了,但LU方法將計算做了分解。其中第一步分解的計算複雜度為  ,而後續回代求解的兩步的計算複雜度均為  。更重要的是,計算量最大的矩陣分解過程可以在低精度上執行。
綜上所述,一個混合精度線性方程求解演算法可以描述如下:
上述描述中的矩陣P代表矩陣分解中的行交換操作。這個演算法紅色部分的計算都在單精度上實現,只有藍色部分才在更高的雙精度上執行。實踐證明,儘管迭代次數由於低精度的引入可能會增加,但由於低精度帶來的計算速度和頻寬的大幅提升,整體求解效率仍然可以提高。值得注意是該方法中精度轉換操作也會帶來一些執行的代價和效能的影響,本文不再展開。

混合精度矩陣分解

進一步分析發現,矩陣分解如LU或QR本身也可以用混合精度實現,比如可以用FP16-32混合精度實現FP32矩陣的分解。

圖1 矩陣LU分解
如圖1所示,大型矩陣的LU分解可由一系列的兩階段(panel分解,後緣更新)三步計算迭代完成:panel分解(橙色),U陣更新(紫色),剩餘A陣更新(綠色)。LINPACK benchmark就要求測試者按照LU求解的方法來測試系統算力。透過分析我們可以發現,該分解中A陣更新的計算代價最大,複雜度為,消耗了主要的計算量。經過實驗,研究者發現該更新可以用更低精度完成。一種可行的混合精度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.0Eflop/s,其混合精度和雙精度計算能力均高於Summit的指標。但值得一提的是,Summit的混合精度和雙精度計算指標之比是7.7倍,高於富嶽的4.5倍。Summit較高的混合精度計算加速一個可能原因是二者架構上的區別。富嶽超算整合由同構的Arm架構CPU構成。該CPU的FP16計算能力是其雙精度的4.5倍。Summit是CPU與GPU的混合結構,其GPU可以用Tensor core提供高一個數量級的FP16計算能力。
另一個有趣的地方是混合精度計算表中第三位的Selene的FP64計算能力弱於第四位的Perlmutter,但是混合精度計算能力反而超出。這一現象在表中還有多處,有興趣的讀者可檢視全部排名。
表2 HPL-AI排名(2021.11)[4]
結論與思考
本文簡述了在GPU等多精度計算加速裝置上實現混合精度計算的方法,討論了混合精度計算對超算系統在科學計算上效能的影響。本文還介紹了混合精度計算benchmark和發展情況。隨著計算加速硬體日新月異的發展,我們相信混合精度計算將取得越來越多的進展並應用到更多的科學研究和工程實施領域。
參考文獻
[1] Austin Derrow-Pinion, Jennifer She, David Wong, et al. ETA Predictionwith Graph Neural Networks in Google Maps. 2021
[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/
關於壁仞科技研究院
壁仞科技研究院作為壁仞科技的前沿研究部門,旨在研究新型智慧計算系統的關鍵技術,重點關注新型架構,先進編譯技術和設計方法學,並將逐漸拓展研究方向,探索未來智慧系統的各種可能。壁仞科技研究院秉持開放的原則,將積極投入各類產學研合作並參與開源社群的建設,為相關領域的技術進步做出自己的貢獻。
掃碼關注我們

相關文章