ECS實例——Arm芯片的 Python-AI算力優(yōu)化

深度學習技術(shù)在圖像識別、搜索推薦等領(lǐng)域得到了廣泛應(yīng)用。近年來各大 CPU 廠商也逐漸把 AI 算力納入了重點發(fā)展方向，通過《Arm 芯片 Python-AI 算力優(yōu)化》我們將看到龍蜥社區(qū) Arm 架構(gòu) SIG（Special Interest Group） 利用最新的 Arm 指令集優(yōu)化 Python-AI 推理 workload 的性能。

倚天 ECS 實例的 AI 推理軟件優(yōu)化

阿里云推出的倚天 Arm ECS 實例，擁有針對 AI 場景的推理加速能力，我們將了解加速的原理以及以及相關(guān)的軟件生態(tài)適配。

卷積神經(jīng)網(wǎng)絡(luò)(CNN)在圖像和語音領(lǐng)域使用廣泛，神經(jīng)網(wǎng)絡(luò)算法相比傳統(tǒng)的算法消耗了更多算力。為了探索對計算的優(yōu)化，我們進一步看到 AlexNet 模型(一種 CNN)的推理過程的各個層的計算資源消耗占比。

可以看到名為 conv[1-5] 的 5 個卷積層消耗了 90% 的計算資源，因此優(yōu)化 CNN 推理的關(guān)鍵就是優(yōu)化卷積層的計算。

3、對由圖像塊組成的矩陣和由多個卷積核展開組成的矩陣應(yīng)用矩陣乘法。

上面一頁的計算應(yīng)用了矩陣乘法操作，為什么我們不采用更加直接的迭代計算方式，而是采用需要額外內(nèi)存的矩陣乘法呢？這里有兩個關(guān)鍵因素：

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

• 計算機科學家們已經(jīng)深入探索了矩陣乘法操作，矩陣乘法操作可以被優(yōu)化得非?？臁?/span>

在 fortran 世界中，GEMM(general matrix multiplication)已經(jīng)成為一個通用操作：該操作通過對數(shù)據(jù)重新排列，精心設(shè)計計算過程，利用多線程和向量指令，可以比自己實現(xiàn)的樸素版本快十倍以上。因此使用矩陣運算帶來的收益相比額外的開銷是值得的。

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

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

• Intel AMX(Advanced Matrix Extensions) 通過脈動陣列在硬件層面支持矩陣乘。

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

雖然在 AI 算力上 GPU 要遠高于 CPU，但是 CPU 因為其部署方便，且無需在主機-設(shè)備間拷貝內(nèi)存，在 AI 推理場景占有一席之地。目前市面上尚沒有可以大規(guī)模使用的支持 AMX 或者 SME 的硬件，在這個階段我們應(yīng)該如何優(yōu)化 CPU 上的 AI 推理算力？我們首先要了解 BF16 數(shù)據(jù)類型。

BF16(Brain Float 16)是由 Google Brain 開發(fā)設(shè)計的 16 位浮點數(shù)格式。相比傳統(tǒng)的 IEEE16 位浮點數(shù)，BF16 擁有和 IEEE 單精度浮點數(shù)(FP32)一樣的取值范圍，但是精度較差。研究人員發(fā)現(xiàn)，在 AI 訓(xùn)練和推理中，使用 BF16 可以節(jié)約一半的內(nèi)存，獲得和單精度浮點數(shù)接近的準確率。

根據(jù)右圖，BF16 指數(shù)的位數(shù)和 FP32 是一致的，因此 BF16 和 FP32 的相互轉(zhuǎn)換只要截斷尾數(shù)即可，左下角圖上便是 tensorflow 源碼中的轉(zhuǎn)換實現(xiàn)。

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

如今我們雖然無法大規(guī)模使用到支持 AMX/SME 的硬件，但是 Armv8.6-A 提供了 bf16 擴展，該擴展利用了有限的 128bit 向量寄存器，通過 BFMMLA 指令執(zhí)行矩陣乘法運算：

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

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

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

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

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

倚天 ECS 實例是市面上少數(shù)可以支持 bf16 指令擴展的 Arm服務(wù)器。目前已經(jīng)支持了 Tensorflow 和 Pytorch 兩種框架的 AI 推理：

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

• Pytorch 已經(jīng)支持了 OneDNN + ACL，但是目前還在試驗狀態(tài)，無法很好地發(fā)揮性能。但是 Pytorch 同時支持 OpenBLAS 作為其計算后端，因此可以通過 OpenBLAS 來享受 ARM bf16 擴展帶來的性能收益。

可以看到相比默認的 eigen 實現(xiàn)，開啟 OneDNN + ACL 后，perf 獲得的計算熱點已經(jīng)從 fmla(向量乘加)轉(zhuǎn)換到了 bfmmla，算力顯著提升。

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

• g7：Intel IceLake 實例。

• g8m：倚天 Arm 服務(wù)器。

左邊柱狀圖中藍色柱子表示算力對比，橙色柱子表示考慮性價比后使用倚天處理器獲得的收益?？梢钥吹皆?Resnet50和 BERT-Large 模型的推理場景下，軟件優(yōu)化后的倚天處理器皆可獲得一倍左右的性價比收益。

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

• 未適配 Arm 架構(gòu)，安裝失敗。
• 軟件未適配 bf16 擴展或者環(huán)境參數(shù)有誤，無法發(fā)揮硬件的全部算力，性能打折。

• 需要精心選擇計算后端，例如目前 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 架構(gòu)提高 CPU 利用率。Python 作為動態(tài)語言，其模塊是動態(tài)導(dǎo)入的，因此啟動速度不是 Python 的強項，這也制約了 Python workload 在 Serverless 場景的普及。

3、執(zhí)行代碼對象。

其中的第二步在首次加載模塊時，要對 .py 文件進行編譯，獲得 code_object, 為了降低將來加載的開銷，Python 解釋器會序列化并緩存 code_object 到 .pyc 文件。

即便模塊導(dǎo)入過程已經(jīng)通過緩存機制優(yōu)化過了，但是讀取 .pyc 文件并反序列化依舊比較耗時。

在這里我們借助了 OpenJDK 的 AppCDS 的思路：將 heap 上的 code_object 復(fù)制到內(nèi)存映射文件中(mmap)。在下次加載模塊時，直接使用 mmap 中的 code_object。

這種框架下有兩個難點：

1、Python 的 code_object 是散落在 heap 的各處且不連續(xù)的，因此 mmap 復(fù)制整個 heap 是行不通的。我們采用的方式是以 code_object 為根，遍歷對象圖，對感興趣的內(nèi)容復(fù)制并緊湊排布。

2、Python 的 code_object 會引用 .data 段的變量，在 Linux 的隨機地址安全機制下，.data 段的數(shù)據(jù)的地址在每次運行時都會隨機變化，這樣 mmap 中的指針就失效了。我們的解決方式是遍歷所有對象，針對 .data 段的指針進行偏移量修復(fù)。

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

我們測試了 bota3、numpy、flask 等常用的 Python 庫，平均可以節(jié)省 20% 的模塊導(dǎo)入耗時。

對于現(xiàn)有的 Python 應(yīng)用可以輕易地使用 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，修改環(huán)境變量運行和使用 cdsAPI 做 dump 即可對現(xiàn)有的應(yīng)用啟動進行加速了。

code-data-share-for-python 是一個新項目，需要大家的參與和反饋，歡迎通過以下鏈接了解和使用：

https://github.com/alibaba/code-data-share-for-python

https://pypi.org/project/code-data-share-for-python

ARM 架構(gòu) SIG鏈接地址：