Arm KleidiAI 助力提升 PyTorch 上 LLM 推理性能

作者：Arm 基礎(chǔ)設(shè)施事業(yè)部軟件工程師 Nobel Chowdary Mandepudi

生成式人工智能 (AI) 正在科技領(lǐng)域發(fā)揮關(guān)鍵作用，許多企業(yè)已經(jīng)開始將大語(yǔ)言模型 (LLM) 集成到云端和邊緣側(cè)的應(yīng)用中。生成式 AI 的引入也使得許多框架和庫(kù)得以發(fā)展。其中，PyTorch 作為熱門的深度學(xué)習(xí)框架尤為突出，許多企業(yè)均會(huì)選擇其作為開發(fā) AI 應(yīng)用的庫(kù)。通過(guò)部署 Arm Kleidi 技術(shù) ，Arm 正在努力優(yōu)化 PyTorch，以加速在基于 Arm 架構(gòu)的處理器上運(yùn)行 LLM 的性能。Arm 通過(guò)將 Kleidi 技術(shù)直接集成到 PyTorch 中，簡(jiǎn)化了開發(fā)者訪問(wèn)該技術(shù)的方式。

在本文中，我們將通過(guò)一個(gè)演示應(yīng)用來(lái)展示 Arm KleidiAI 在 PyTorch 上運(yùn)行 LLM 實(shí)現(xiàn)的性能提升。該演示應(yīng)用在基于 Arm Neoverse V2 的亞馬遜云科技 (AWS) Graviton4 R8g.4xlarge EC2 實(shí)例上運(yùn)行 Llama 3.1。如果你感興趣，可以使用以下 Learning Path，自行重現(xiàn)這個(gè)演示。 (Link: https://learn.arm.com/learning-paths/servers-and-cloud-computing/pytorch-llama)

演示應(yīng)用

我們的演示應(yīng)用是一個(gè)基于 LLM 的聊天機(jī)器人，可以回答用戶提出的各種問(wèn)題。該演示使用 Arm 平臺(tái)上的 PyTorch 框架運(yùn)行 Meta Llama 3.1 模型，并被設(shè)計(jì)成一個(gè)使用 Streamlit 前端的瀏覽器應(yīng)用。Streamlit 將信息提供給 Torchat 框架，后者運(yùn)行 PyTorch 并作為 LLM 后端。Torchat 輸出的信息進(jìn)入注意力層并生成詞元 (token)。這些詞元使用 OpenAI 框架流式傳輸功能發(fā)送到前端，并在瀏覽器應(yīng)用上顯示給用戶。該演示的架構(gòu)下圖所示。

圖：演示架構(gòu)

演示應(yīng)用在 LLM 推理結(jié)束后測(cè)定并顯示以下性能指標(biāo)：

生成首個(gè)詞元的用時(shí)（秒） ：對(duì)于 LLM 推理，需要快速生成首個(gè)詞元，以盡量減少延遲并向用戶提供即時(shí)輸出。

解碼速度/文本生成（詞元/秒） ：每秒詞元數(shù)是指生成式 AI 模型生成詞元的速率。生成下一個(gè)詞元的時(shí)間最長(zhǎng)不超過(guò) 100 毫秒，這是交互式聊天機(jī)器人的行業(yè)標(biāo)準(zhǔn)。這意味著解碼速度至少為 10 個(gè)詞元/秒。這對(duì)于提升實(shí)時(shí)應(yīng)用的用戶體驗(yàn)至關(guān)重要。

生成百萬(wàn)詞元的成本（美元） ：根據(jù) AWS 云端 EC2 實(shí)例的解碼速度和每小時(shí)成本，我們可以計(jì)算出生成 100 萬(wàn)個(gè)詞元的成本，這也是一個(gè)常用的比較指標(biāo)。由于每小時(shí)成本是固定的，解碼速度越快，生成百萬(wàn)詞元的成本就越低。

生成提示詞的總用時(shí)（秒） ：這是使用所有詞元生成提示詞所花費(fèi)的總時(shí)間。

生成提示詞的總成本（美元） ：這是根據(jù)使用所有詞元生成完整提示詞的總時(shí)間、解碼速度和云端機(jī)器成本計(jì)算得出的。

下圖顯示了示例響應(yīng)，可作為使用所示指標(biāo)驗(yàn)證聊天機(jī)器人的示例。生成首個(gè)詞元的時(shí)間短于 1 秒，解碼速率為 33 個(gè)詞元/秒，這兩項(xiàng)數(shù)據(jù)都非常令人滿意，并且滿足交互式聊天機(jī)器人的行業(yè)標(biāo)準(zhǔn)。

圖：包含示例響應(yīng)和指標(biāo)的演示

針對(duì) PyTorch 的 KleidiAI 優(yōu)化

KleidiAI 庫(kù)為 Arm 平臺(tái)提供了多項(xiàng)優(yōu)化。Kleidi 在 Torch ATen 層中提供了一個(gè)新算子以加載模型。該層將模型權(quán)重以特定格式打包在內(nèi)存中，使得 KleidiAI GEMM 內(nèi)核可用來(lái)提高性能。同樣地，針對(duì)模型執(zhí)行的優(yōu)化使用了 ATen 層中的另一個(gè)算子。該算子對(duì)先前打包的模型權(quán)重進(jìn)行 matmul 運(yùn)算的量化。

在我們的演示中，該模型是從 Meta Hugging Face 庫(kù)下載的。該模型使用 INT4 內(nèi)核布局打包在內(nèi)存中，然后使用針對(duì) PyTorch 優(yōu)化的 INT4 KleidiAI 內(nèi)核進(jìn)行量化。該演示的架構(gòu)如下圖所示。

圖：針對(duì) PyTorch 實(shí)現(xiàn)的 KleidiAI 優(yōu)化

使用我們 Learning Path 中包含的補(bǔ)丁 [注] ，可將這些 KleidiAI 優(yōu)化應(yīng)用到 PyTorch、Torchchat 和 Torchao 中。你可以使用這些補(bǔ)丁來(lái)查看 Arm 平臺(tái)上的 PyTorch 為工作負(fù)載帶來(lái)的 LLM 推理性能提升。

注： Arm KleidiAI 的 PyTorch 補(bǔ)丁正在與上游 PyTorch 合并，并將在未來(lái)的 PyTorch 官方版本中提供。

性能

為了印證 KleidiAI 的性能優(yōu)勢(shì)，我們使用 PyTorch 運(yùn)行相同的聊天機(jī)器人應(yīng)用，并測(cè)定了 KleidiAI 優(yōu)化前后的每秒生成詞元數(shù)和生成首個(gè)詞元的用時(shí)，結(jié)果如下圖所示。

圖：性能比較

可以看到，將 KleidiAI 庫(kù)應(yīng)用到現(xiàn)有的生成式 AI 技術(shù)棧中可以大大提高詞元生成速率，并縮短為不同生成式 AI 模型生成首個(gè)詞元的時(shí)間。

結(jié)論

對(duì)于聊天機(jī)器人等實(shí)時(shí)工作負(fù)載來(lái)說(shuō)，在 CPU 上運(yùn)行 LLM 推理可行且有效。我們?cè)谥?nbsp;《在基于 Arm Neoverse 的 AWS Graviton3 CPU 上實(shí)現(xiàn)出色性能》 文章中使用 Llama.cpp 演示了這一點(diǎn)。在本文中，我們展示了如何使用 KleidiAI 庫(kù)為 Arm 平臺(tái)上的 PyTorch 實(shí)現(xiàn)良好的 LLM 推理性能。通過(guò)使用搭載 Neoverse V2 核心且基于 AWS Graviton4 的 R8g 實(shí)例進(jìn)行演示，印證了 KleidiAI 為在 Arm 平臺(tái)上使用 PyTorch 運(yùn)行 LLM 推理實(shí)現(xiàn)了顯著的性能提升。開發(fā)者現(xiàn)在可以利用 Arm 針對(duì) PyTorch 的 KleidiAI 優(yōu)化來(lái)運(yùn)行新的或現(xiàn)有的 AI 應(yīng)用。