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

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

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

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

演示應用

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

圖：演示架構

演示應用在 LLM 推理結束后測定并顯示以下性能指標：

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

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

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

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

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

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

圖：包含示例響應和指標的演示

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

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

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

圖：針對 PyTorch 實(shí)現的 KleidiAI 優(yōu)化

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

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

性能

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

圖：性能比較

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

結論

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