作者：Arm 基礎(chǔ)設(shè)施事業(yè)部高級產(chǎn)品經(jīng)理 Ashok Bhat

PyTorch 是一個廣泛應(yīng)用的開源機(jī)器學(xué)習(xí) (ML) 庫。近年來，Arm 與合作伙伴通力協(xié)作，持續(xù)改進(jìn) PyTorch 的推理性能。本文將詳細(xì)介紹如何利用 Arm Kleidi 技術(shù)提升 Arm Neoverse 平臺上的 PyTorch 推理表現(xiàn)。Kleidi 技術(shù)可以通過 Arm Compute Library (ACL) 和 KleidiAI 庫獲取。

PyTorch 提供兩種主要的執(zhí)行模式：即時執(zhí)行模式 (Eager Mode) 和圖模式 (Graph Mode)。即時執(zhí)行模式是一種動態(tài)執(zhí)行模式，操作會以 Python 代碼編寫的方式立即執(zhí)行，該模式非常適用于實驗與調(diào)試。而圖模式則是在執(zhí)行前將一系列 PyTorch 操作編譯成靜態(tài)計算圖，從而實現(xiàn)性能優(yōu)化和高效的硬件加速。通過使用 torch.compile 函數(shù)，可以方便地將 PyTorch 代碼轉(zhuǎn)換為圖模式，通常能夠顯著提升執(zhí)行速度。

PyTorch 即時執(zhí)行模式

CPU 推理性能提升高達(dá)三倍

PyTorch 即時執(zhí)行模式使用 oneDNN，針對具有 ACL 內(nèi)核的 Arm Neoverse 處理器進(jìn)行了優(yōu)化。可以通過以下的 PyTorch 軟件棧圖進(jìn)行了解。

圖 1：PyTorch 軟件棧

PyTorch 中的 FX Graph 是用于可視化和優(yōu)化 PyTorch 模型的一種中間表示形式。

ATen 是支撐 PyTorch 框架的基礎(chǔ)張量庫。它提供了核心張量類別和大量數(shù)學(xué)運算，構(gòu)成了 PyTorch 模型的基本組件。

oneDNN 是一個性能庫，為包括 Arm 和 x86 在內(nèi)的各種硬件架構(gòu)提供常見深度學(xué)習(xí)原語的優(yōu)化實現(xiàn)方案。在這些架構(gòu)上，ATen 使用 oneDNN 作為性能增強(qiáng)后端。這意味著當(dāng) PyTorch 遇到支持的操作時，會將計算委托給 oneDNN，后者可以使用針對特定硬件的優(yōu)化來提高執(zhí)行效率。

Arm Compute Library 于 2016 年首次發(fā)布，提供針對 Arm 進(jìn)行優(yōu)化的關(guān)鍵 ML 原語，包括卷積、池化、激活函數(shù)、全連接層、歸一化。這些原語利用 Arm Neoverse 核心上針對特定 ML 和特定硬件的功能和指令來實現(xiàn)高性能。我們已將 Arm Compute Library 集成到 oneDNN 中，以便在 Arm 平臺上加速 ATen 操作。

Arm Neoverse CPU 包含有助于加速 ML 的硬件擴(kuò)展，其中包括 Neon、SVE/SVE2、BF16 和 I8MM，通過有效地進(jìn)行向量處理、BF16 運算和矩陣乘法來加速 ML 任務(wù)。

圖 2：各種模型在即時執(zhí)行模式下實現(xiàn)的性能提升

PyTorch 圖模式（使用 torch.compile）

比 PyTorch 即時執(zhí)行模式進(jìn)一步提高兩倍

PyTorch 2.0 引入了 torch.compile，與默認(rèn)的即時執(zhí)行模式相比，可提高 PyTorch 代碼的速度。與即時執(zhí)行模式不同，torch.compile 將整個模型預(yù)編譯成針對特定硬件平臺優(yōu)化的單圖。從 PyTorch 2.3.1 開始，官方 AArch64 安裝包均包含 torch.compile 優(yōu)化。在基于亞馬遜云科技 (AWS) Graviton3 的 Amazon EC2 實例上，對于各種自然語言處理 (NLP)、計算機(jī)視覺 (CV) 和推薦模型，這些優(yōu)化可以為 TorchBench 模型推理帶來比即時執(zhí)行模式高出兩倍的性能。

圖 3：各種模型在編譯模式下實現(xiàn)的性能提升

下一步通過 KleidiAI 庫

提升生成式 AI 推理性能

目前，我們已經(jīng)研究了 Arm Compute Library 如何在即時執(zhí)行模式和編譯模式下提升 PyTorch 推理性能。接下來，我們來看 PyTorch 即將推出什么新功能。Arm 目前正在努力提升 PyTorch 中的大語言模型 (LLM) 推理性能，并以 Llama 和 Gemma 為主要 LLM 示例。

經(jīng)優(yōu)化的 INT4 內(nèi)核

今年早些時候，Arm 軟件團(tuán)隊和合作伙伴共同優(yōu)化了 llama.cpp 中的 INT4 和 INT8 內(nèi)核，以利用更新的 DOT 和 MLA 指令。在 AWS Graviton3 處理器上，這些內(nèi)核在即時評估方面比現(xiàn)有 GEMM MMLA 內(nèi)核提升了 2.5 倍，并且在文本生成方面比默認(rèn)的 vec_dot 內(nèi)核提升了兩倍。這些經(jīng)優(yōu)化的新內(nèi)核也是 Arm KleidAI 庫的一部分。

今年早些發(fā)布的 KleidiAI 庫是一個開源庫，具有針對 Arm CPU 上的 AI 任務(wù)進(jìn)行優(yōu)化的微內(nèi)核。對于微內(nèi)核，可將它視為能夠提升特定 ML 操作性能的軟件。開發(fā)者可以通過包含相關(guān)的 .c 和 .h 文件及公共頭文件來使用這些微內(nèi)核。無需包含庫的其余部分。

Kleidi 與 PyTorch 的集成

圖 4：Kleidi 技術(shù)與 PyTorch 的集成

我們引入了兩種新的 ATen 操作：torch.ops.aten._kai_weights_pack_int4() 和 torch.ops.aten._kai_input_quant_mm_int4()，兩者均使用 KleidiAI 庫中高度優(yōu)化的打包技術(shù)和 GEMM 內(nèi)核。gpt-fast 利用這些 PyTorch 算子來 (1) 使用對稱的每通道量化將權(quán)重量化為 INT4，并添加包含量化尺度的額外數(shù)組；(2) 動態(tài)量化激活矩陣并使用 AArch64 I8MM 擴(kuò)展來執(zhí)行激活矩陣和權(quán)重的 INT8 矩陣乘法。

圖 5：通過在 PyTorch 中集成 KleidiAI

來提升 4 位量化 LLM 模型性能

通過這種方法，與目前默認(rèn)的 PyTorch 實現(xiàn)方案相比，我們可以將 Llama 的推理性能提高 18 倍，將 Gemma 的推理性能提高 14 倍。

結(jié)論

Arm 及其合作伙伴利用 Arm Compute Library 中的 Kleidi 技術(shù)提高了 Arm Neoverse 平臺上的 PyTorch 推理性能。在即時執(zhí)行模式下可實現(xiàn)高達(dá)兩倍的性能提升，在圖模式下（使用 torch.compile）可再提升兩倍。此外，我們還在努力將生成式 AI 模型（Llama 和 Gemma）的推理性能提升高達(dá) 18 倍。

Arm 通過部署 Kleidi 技術(shù)來實現(xiàn)PyTorch 上的優(yōu)化，以加速在基于 Arm 架構(gòu)的處理器上運行 LLM 的性能。Arm 技術(shù)專家在基于 Neoverse V2 的 AWS Graviton4 R8g.4xlarge EC2 實例上運行 Llama 3.1 展示了所實現(xiàn)的性能提升。如果你對這一演示感興趣，可閱讀《Arm KleidiAI 助力提升 PyTorch 上 LLM 推理性能》了解。