編者按：在剛剛結束的 PyCon China 2022 大會上，龍蜥社區開發者朱宏林分享了主題爲《ARM 芯片的 Python+AI 算力優化》的技術演講。本次演講，作者將向大家介紹他們在倚天 710 ARM 芯片上開展的 Python+AI 優化工作，以及在 ARM 雲平臺上部署 Python+AI 任務的最佳實踐。
以下爲本次演講內容：
（圖/朱宏林現場演講）
我們的場景是 ARM 平臺的和 AI 相關的任務，主要的目標是進行性能優化，具體來說我們首先關注的是深度學習推理任務（inference task），主要原因也是來自於業務需求。
這裏說的 ARM 平臺不是我們理解的終端設備，比如手機之類的，而是指服務端平臺。在大家印象中，AI 任務，特別是深度學習的程序一般是跑在 GPU 或者 x86 的 CPU 上，出於功耗、成本、性能等因素的考慮，雲廠商逐步開始建設 ARM 架構的服務平臺，這是一個趨勢。當然 ARM 平臺還不是很成熟，許多軟件還無法成功跑起來，更不要說提升性能了。
我們想要吸引一部分用戶將AI應用從原先的 x86 平臺上遷移到 ARM 平臺上。這就要求 ARM 平臺能提供更好的性能，或者更好的性價比。所以說如何整合 Python+AI 的相關軟件使其發揮最好的性能成爲了我們關注的重點。
下文的分享整體分爲兩部分，一部分是介紹我們進行的優化工作，主要是跟矩陣乘法相關的優化，第二部分是關於 Python AI 應用在 ARM 雲平臺-倚天 710 上的最佳實踐。
一、優化工作介紹
前面說我們的優化是和矩陣乘法相關的，那首先需要說明爲什麼我們會關注到這個。
這裏有一個繞不開的場景就是深度學習，不管是前幾年知名的 AlphaGo，還是當前火熱的 ChatGPT，都用到了大量深度學習的技術，深度學習本身只是AI的一個分支，但卻影響廣泛，不容忽視。所以我們從深度學習開始切入，從當前最廣泛使用的深度學習框架，TensorFlow 和 PyTorch 開始。此外，我們還需要結合硬件場景，即前面說到的 ARM 服務端平臺，對於阿里雲來說就是結合倚天 710 芯片。
深度學習的實現中包含大量的矩陣乘法，甚至有文章直接寫出矩陣乘法是深度學習的核心。舉個例子，我們熟知的卷積操作，實際上經過一系列的轉換後，輸入特徵和卷積核會被轉換爲兩個矩陣，然後進行矩陣乘法，輸出的結果再解碼成特徵圖，就完成了卷積。除此以外，全連接層也由矩陣乘法實現，當前流行的 Transformers 結構，被包括 ChatGPT 在內的各類 NLP 模型所使用，也包含大量矩陣乘法操作。
我們可以看一些例子：
可以看到，像 AlexNet、ResNet-50 之類的模型，在進行推理時，大約 90% 計算耗時在執行矩陣乘法。即使對矩陣乘法做一些微小的優化，影響也是很廣泛的。
我們前面說的矩陣乘法，更準確的叫法是 GEMM，通用矩陣乘法，其實還包含係數和累加操作。但是時間複雜度仍然是 MNK 級別，主要還在於 AB 兩個矩陣相乘。直觀來看，深度學習涉及的矩陣乘法計算量很大，比如常見的卷積操作可能就涉及 5000 萬次計算，所以優化就顯得很有必要，右下圖是最樸素的三層循環迭代法，這種做法通常非常慢，計算機科學家做了許多努力，從優化內存佈局和利用向量指令出發，能夠將性能提升 10 倍以上。
內存佈局主要分兩步，第一步是對矩陣進行分塊，即對於一個超大的矩陣，我們並不是一個一個按順序計算，而是將矩陣切分爲一個一個小塊，分小塊計算。第二步是對分出的小塊，內部的元素序列進行重排，例如原來是按行排列的矩陣，那可能第一行四個計算好了，就需要取第二行的前四個，但是要取第二行就需要指針移動很長的距離，很容易造成 cache 不命中，於是需要重排，使得他們在內存上連續。優化內存佈局主要目的是爲了增加 cache 命中率，減少訪存次數。
其次是利用向量化指令，類似 AVX 對於 x86 設備，NEON 對於 ARM 設備。向量化指令本質上是爲了同時對多個數據進行計算，例如我們要對四組數據分別進行乘法，那麼常規情況下需要執行四次，如果將它們對應放入向量寄存器中，只需要一條向量化指令，就可以同時得出四個結果，計算效率得到提升。當然這個是需要硬件支持。
因爲 AI 推理大量使用了矩陣乘法，如今也有許多硬件對矩陣運算進行了加速：
- NVIDIA Volta 架構引入了tensor core，可以高效地以混合精度處理矩陣乘
- Intel AMX(Advanced Matrix Extensions) 通過脈動陣列在硬件層面支持矩陣乘
- ARM SME(Scalable Matrix Extension) 支持向量外積運算，加速矩陣乘
目前市面上尚沒有可以大規模使用的支持 AMX 或者 SME 的硬件，在這個階段我們應該如何優化 CPU 上的 AI 推理算力呢？我們首先要了解 BF16 數據類型。
BF16(全稱 Brain Floating Point)，是由 Google Brain 開發設計的 16 位浮點數格式。
相比傳統的 FP16 位浮點數，BF16 擁有和 FP32 一樣的取值範圍，但是精度較差。但對於深度學習來說，較低的精度並不顯著影響結果，而較低的表示範圍則會顯著影響模型訓練的好壞。
此外，BF16 還具有轉換方便的特點，BF16 和 FP32 的互轉只需要截斷或填充尾數即可。
使用 BF16 還可以節約一半的內存，緊湊的內存表示通常意味着更高的計算吞吐。
最後，我們也有了硬件指令支持，可以直接對 BF16 數據進行操作。
需要說明的是 BF16 的擴展包含在 ARMv8.6 設備上，當然倚天 710 是 ARMv9 的指令集，同樣支持。
我們主要通過 BFMMLA 來進行矩陣乘法計算，例如對於包含 128bit 的向量寄存器的設備來說：
- 輸入 A: 大小爲 2*4 的 BF16 矩陣，按行存儲
- 輸入 B: 大小爲 4*2 的 BF16 矩陣，按列存儲
- 輸出 C: 大小爲 2*2 的 FP32 矩陣
BFMMLA 單指令完成 16 次乘法和 16 次加法，計算吞吐非常高。
當然這時候如果我們需要 C 是 BF16 類型的話，就需要應用轉換指令，例如向量化指令 BFCVT，加速轉換過程。
我們的目標還是給 tensorflow 和 pytorch 用戶提供加速，這是整體的流程圖，對於一個 AI 推理任務，實際上不論是 TensorFlow 還是 PyTorch 都不會自己直接去計算，而是叫個專門的計算後端，在 ARM 主要是兩個，一個是 ARM Compute Library，另一個是 OpenBLAS，他們之間的關係如右圖。
TensorFlow 在最近的版本中開始採用 oneDNN + ACL 作爲計算後端，oneDNN 也是一層皮，實際的計算仍然是 ACL。用戶實際上只需要設置一個環境變量，就可以在不該動代碼的情況下獲得 BF16 加速。這個改進是由 ARM 公司的研發人員首先完成了。具體操作例子如下：
# 假設 resnet.py 包含用戶寫的模型推理的代碼
DNNL_DEFAULT_FPMATH_MODE=BF16 python3 resnet.py
PyTorch的情況比較複雜，PyTorch 支持 OneDNN + ACL，但無法很好的發揮性能，同時 PyTorch 支持 OpenBLAS 後端，因此可以通過 OpenBLAS 來享受 ARM bf16 擴展帶來的性能收益。
OpenBLAS 的 BF16 的 GEMM 優化是由龍蜥社區理事單位阿里巴巴貢獻的，於此同時，我們爲了方便用戶使用，也在 PyTorch 中加入了一個API，用戶在模型執行前添加一行torch.set_float32_fast_math_mode("BF16")，就可以獲得 BF16 加速，不必修改其他代碼（需要說明，這個api還沒有合入PyTorch，所以目前要使用我們提供的pytorch鏡像纔可以獲得）。操作例子如下：
# ...
# 在模型執行前設置fast math mode
torch.set_float32_fast_math_mode("BF16")
# ...
# 執行模型
pred = model(x)
# ...
之後是一些性能測試的展示，我們測試了 OpenBLAS 純矩陣計算的性能對比。分別記錄了 GFLOPS 和執行時間兩個指標。
然後測試 TensorFlow 和 PyTorch 的性能對比，在對比中，我們可以看到，得益於 BF16 擴展，最新的 ECS ARM 平臺上的性能優於 x86 平臺（g7）。
二、Python AI應用在ARM雲平臺-倚天710上的最佳實踐
現在介紹一下在 ARM 平臺，特別是倚天 710 的用戶，使用 TensorFlow 或 PyTorch 的最佳實踐。
要知道軟件版本的選擇十分重要，隨意選擇 tensorflow 或者 pytorch 包可能遭遇：
- 未適配 ARM 架構，安裝失敗
- 軟件未適配 BF16 擴展或者環境參數有誤，無法發揮硬件的全部算力，性能打折
- 需要精心選擇計算後端，例如目前 pytorch下OpenBLAS 較快
在 TensorFlow 上，我們可以選擇最新的兩個官方版本， 2.10.1 或者 2.11.0（最新版本），才能夠獲得 ACL 的 BF16 加速。用戶也可以選擇阿里雲的鏡像，這個和 pip 安裝的其實是一樣的，沒有區別。
對於 PyTorch 用戶，官方版本只有在最新的 1.13.0 才能夠獲得 ACL 加速，但是正如前面所說的，實際性能並不突出。阿里雲則提供了帶最新 OpenBLAS 的 PyTorch，在 docker 拉取時標註 torch_openblas 就可以獲得。此外，我們也提供了modelzoo 鏡像，包含模型的測試代碼和驗證代碼。
目前我們仍然在進行相關的工作，期待後續能爲大家提供更加完善的鏡像。歡迎大家入羣一起探索相關技術。
AI SIG 主頁地址：https://openanolis.cn/sig/AI_SIG
本文爲阿里雲原創內容，未經允許不得轉載。