深度學習技術在圖像識别、搜尋推薦等領域得到了廣泛應用。近年來各大 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連結位址:
—— 完 ——