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

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

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

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

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

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

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

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

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

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

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

因?yàn)?AI 推理大量使用了矩陣乘法，如今也有許多硬件對(duì)矩陣運(yùn)算進(jìn)行了加速：

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

從 workload 角度評(píng)測(cè)，上圖對(duì)比了兩種機(jī)型：

• g7：Intel IceLake 實(shí)例。

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

左邊柱狀圖中藍(lán)色柱子表示算力對(duì)比，橙色柱子表示考慮性價(jià)比后使用倚天處理器獲得的收益。可以看到在 Resnet50和 BERT-Large 模型的推理場(chǎng)景下，軟件優(yōu)化后的倚天處理器皆可獲得一倍左右的性價(jià)比收益。

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

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

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

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

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

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

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

這種框架下有兩個(gè)難點(diǎn)：

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

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

因?yàn)樵擁?xiàng)目共享了 Python 的 code_object，因此名字是 code-data-share-for-python，簡(jiǎn)稱 pycds。

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

對(duì)于現(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


# 對(duì)比基線
time python -c 'import numpy'
real 0m0.105s
user 0m0.216s
sys 0m0.476s

我們僅僅通過安裝 PyPI，修改環(huán)境變量運(yùn)行和使用 cdsAPI 做 dump 即可對(duì)現(xiàn)有的應(yīng)用啟動(dòng)進(jìn)行加速了。

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

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

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

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