技術解讀倚天 ECS 例項——Arm 晶片的 Python-AI 算力最佳化 | 龍蜥技術

深度學習技術在影像識別、搜尋推薦等領域得到了廣泛應用。近年來各大 CPU 廠商也逐漸把 AI 算力納入了重點發展方向，透過《Arm 晶片 Python-AI 算力最佳化》我們將看到龍蜥社群 Arm 架構 SIG（Special Interest Group） 利用最新的 Arm 指令集最佳化 Python-AI 推理 workload 的效能。

倚天 ECS 例項的 AI 推理軟體最佳化

阿里雲推出的倚天 Arm ECS 例項，擁有針對 AI 場景的推理加速能力，我們將瞭解加速的原理以及以及相關的軟體生態適配。

卷積神經網路(CNN)在影像和語音領域使用廣泛，神經網路演算法相比傳統的演算法消耗了更多算力。為了探索對計算的最佳化，我們進一步看到 AlexNet 模型(一種 CNN)的推理過程的各個層的計算資源消耗佔比。

可以看到名為 conv[1-5] 的 5 個卷積層消耗了 90% 的計算資源，因此最佳化 CNN 推理的關鍵就是最佳化卷積層的計算。

我們進一步來看如何對影像應用卷積核：

1、使用 im2col 根據卷積核尺寸，將影像轉化為若干塊(patch)。

2、將多個卷積核展開成若干向量。

3、對由影像塊組成的矩陣和由多個卷積核展開組成的矩陣應用矩陣乘法。

上面一頁的計算應用了矩陣乘法操作，為什麼我們不採用更加直接的迭代計算方式，而是採用需要額外記憶體的矩陣乘法呢？這裡有兩個關鍵因素：

• 深度學習的卷積計算量很大，典型計算需要涉及 5000 萬次乘法和加法操作，因此對計算的最佳化十分重要。

• 電腦科學家們已經深入探索了矩陣乘法操作，矩陣乘法操作可以被最佳化得非常快。

在 fortran 世界中，GEMM(general matrix multiplication)已經成為一個通用操作：

該操作透過對資料重新排列，精心設計計算過程，利用多執行緒和向量指令，可以比自己實現的樸素版本快十倍以上。因此使用矩陣運算帶來的收益相比額外的開銷是值得的。

因為 AI 推理大量使用了矩陣乘法，如今也有許多硬體對矩陣運算進行了加速：

• NVIDIA Volta 架構引入了 tensor core，可以高效地以混合精度處理矩陣乘。

• Intel AMX(Advanced Matrix Extensions) 透過脈動陣列在硬體層面支援矩陣乘。

• Arm SME(Scalable Matrix Extension) 支援向量外積運算，加速矩陣乘。

雖然在 AI 算力上 GPU 要遠高於 CPU，但是 CPU 因為其部署方便，且無需在主機-裝置間複製記憶體，在 AI 推理場景佔有一席之地。目前市面上尚沒有可以大規模使用的支援 AMX 或者 SME 的硬體，在這個階段我們應該如何最佳化 CPU 上的 AI 推理算力？我們首先要了解 BF16 資料型別。

BF16(Brain Float 16)是由 Google Brain 開發設計的 16 位浮點數格式。相比傳統的 IEEE16 位浮點數，BF16 擁有和 IEEE 單精度浮點數(FP32)一樣的取值範圍，但是精度較差。研究人員發現，在 AI 訓練和推理中，使用 BF16 可以節約一半的記憶體，獲得和單精度浮點數接近的準確率。

根據右圖，BF16 指數的位數和 FP32 是一致的，因此 BF16 和 FP32 的相互轉換隻要截斷尾數即可，左下角圖上便是 tensorflow 原始碼中的轉換實現。

引入 BF16 的一大價值是如今的很多硬體計算的瓶頸在暫存器寬度或者訪問記憶體的速度上，更緊湊的記憶體表示往往可以獲得更高的計算吞吐，在理想情況下，BF16 相比 FP32 可以提高一倍的吞吐(FLOPS)。

如今我們雖然無法大規模使用到支援 AMX/SME 的硬體，但是 Armv8.6-A 提供了 bf16 擴充套件，該擴充套件利用了有限的 128bit 向量暫存器，透過 BFMMLA 指令執行矩陣乘法運算：

• 輸入 A：大小為 2*4 的 BF16 矩陣，按行儲存。

• 輸入 B：大小為 4*2 的 BF16 矩陣，按列儲存。

• 輸出 C：大小為 2*2 的 FP32 矩陣。

該指令單次執行進行了 16 次浮點數乘法和 16 次浮點數加法運算，計算吞吐非常高。

阿里巴巴向 OpenBLAS 專案貢獻了 sbgemm(s 表示返回單精度，b 表示輸入 bf16)的硬體加速實現，從 GEMM 吞吐上看，BF16 相比 FP32 GEMM 吞吐提升超過100%。

倚天 ECS 例項是市面上少數可以支援 bf16 指令擴充套件的 Arm 伺服器。目前已經支援了 Tensorflow 和 Pytorch 兩種框架的 AI 推理：

• Tensorflow下可以透過OneDNN + ACL(Arm Compute Library)來使用 BFMMLA 加速。

• Pytorch 已經支援了 OneDNN + ACL，但是目前還在試驗狀態，無法很好地發揮效能。但是 Pytorch 同時支援 OpenBLAS 作為其計算後端，因此可以透過 OpenBLAS 來享受 ARM bf16 擴充套件帶來的效能收益。

可以看到相比預設的 eigen 實現，開啟 OneDNN + ACL 後，perf 獲得的計算熱點已經從 fmla(向量乘加)轉換到了 bfmmla，算力顯著提升。

從 workload 角度評測，上圖對比了兩種機型：

• g7：Intel IceLake 例項。

• g8m：倚天 Arm 伺服器。

左邊柱狀圖中藍色柱子表示算力對比，橙色柱子表示考慮價效比後使用倚天處理器獲得的收益。可以看到在 Resnet50 和 BERT-Large 模型的推理場景下，軟體最佳化後的倚天處理器皆可獲得一倍左右的價效比收益。

在上文中，我們看到使用倚天處理器若想獲得較高收益，軟體版本的選擇十分重要。隨意選擇 tensorflow 或者 Pytorch 包可能遭遇：

• 未適配 Arm 架構，安裝失敗。

• 軟體未適配 bf16 擴充套件或者環境引數有誤，無法發揮硬體的全部算力，效能打折。

• 需要精心選擇計算後端，例如目前 Pytorch 下 OpenBLAS 較快。

因此我們提供了 Docker 映象，幫助雲上的使用者充分使用倚天 710 處理器的 AI 推理效能：

• accc-registry.cn-hangzhou.cr.aliyuncs.com/tensorflow/tensorflow

• accc-registry.cn-hangzhou.cr.aliyuncs.com/pytorch/pytorch

透過 Serverless 能力充分釋放算力

除了使能更多的硬體指令，另一種充分釋放硬體算力的方式就是透過 Serverless 架構提高 CPU 利用率。Python 作為動態語言，其模組是動態匯入的，因此啟動速度不是 Python 的強項，這也制約了 Python workload 在 Serverless 場景的普及。

Python 應用啟動的主要耗時在模組匯入，Python 模組匯入步驟為：

1、尋找到模組所在的檔案。

2、獲得程式碼物件 code_object。

3、執行程式碼物件。

其中的第二步在首次載入模組時，要對 .py 檔案進行編譯，獲得 code_object, 為了降低將來載入的開銷，Python 直譯器會序列化並快取 code_object 到 .pyc 檔案。

即便模組匯入過程已經透過快取機制最佳化過了，但是讀取 .pyc 檔案並反序列化依舊比較耗時。

在這裡我們藉助了 OpenJDK 的 AppCDS 的思路：將 heap 上的 code_object 複製到記憶體對映檔案中(mmap)。在下次載入模組時，直接使用 mmap 中的 code_object。

這種框架下有兩個難點：

1、Python 的 code_object 是散落在 heap 的各處且不連續的，因此 mmap 複製整個 heap 是行不通的。我們採用的方式是以 code_object 為根，遍歷物件圖，對感興趣的內容複製並緊湊排布。

2、Python 的 code_object 會引用 .data 段的變數，在 Linux 的隨機地址安全機制下，.data 段的資料的地址在每次執行時都會隨機變化，這樣 mmap 中的指標就失效了。我們的解決方式是遍歷所有物件，針對 .data 段的指標進行偏移量修復。

因為該專案共享了 Python 的 code_object，因此名字是 code-data-share-for-python，簡稱 pycds。

我們測試了 bota3、numpy、flask 等常用的 Python 庫，平均可以節省 20% 的模組匯入耗時。

對於現有的 Python 應用可以輕易地使用 pycds，且無需修改任何程式碼：

# 安裝pycds
pip install code-data-share # 安裝pycds
# 生成模組列表
PYCDSMODE=TRACE PYCDSLIST=mod.lst python -c 'import numpy’
# 生成 archive
python -c 'import cds.dump; cds.dump.run_dump("mod.lst", "mod.img")’
# 使用archive
time PYCDSMODE=SHARE PYCDSARCHIVE=mod.img python -c 'import numpy'
real 0m0.090s
user 0m0.180s
sys 0m0.339s
# 對比基線
time python -c 'import numpy'
real 0m0.105s
user 0m0.216s
sys 0m0.476s

我們僅僅透過安裝 PyPI，修改環境變數執行和使用 cds API 做 dump 即可對現有的應用啟動進行加速了。

code-data-share-for-python 是一個新專案，需要大家的參與和反饋，歡迎透過以下連結瞭解和使用：

ARM 架構 SIG連結地址：

—— 完 ——