利用Arm Kleidi技術(shù)實(shí)現(xiàn)PyTorch優(yōu)化

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

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

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

PyTorch 即時(shí)執(zhí)行模式

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

PyTorch 即時(shí)執(zhí)行模式使用 oneDNN，針對(duì)具有 ACL 內(nèi)核的 Arm Neoverse 處理器進(jìn)行了優(yōu)化?？梢酝ㄟ^(guò)以下的 PyTorch 軟件棧圖進(jìn)行了解。

圖 1：PyTorch 軟件棧

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

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

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

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

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

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

PyTorch 圖模式（使用 torch.compile）

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

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

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

下一步通過(guò) KleidiAI 庫(kù)

提升生成式 AI 推理性能

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

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

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

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

Kleidi 與 PyTorch 的集成

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

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

圖 5：通過(guò)在 PyTorch 中集成 KleidiAI

來(lái)提升 4 位量化 LLM 模型性能

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

結(jié)論

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

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