LLaMA 2端到端推理打通！來自中國團隊

明敏 2023-11-23 13:34:39 來源：量子位

允中 發自 凹非寺

量子位 | 公眾號 QbitAI

Buddy Compiler 端到端 LLaMA2-7B 推理示例已經合并到 buddy-mlir倉庫[1]主線。

我們在 Buddy Compiler 的前端部分實現了面向 TorchDynamo 的第三方編譯器，從而結合了 MLIR 和 PyTorch 的編譯生態。

目前，前端部分可以覆蓋 LLaMA 計算圖，轉換到 MLIR 后我們集成了部分向量化和并行優化，并在 AVX512 平臺上進行了測試。整個推理過程可以跑通但還需要大量優化。以下是相關鏈接和現狀：

• [E2E] Buddy Compiler 端到端 LLaMA2-7B 推理示例[2]
• [E2E] 上述端到端推理示例目的是展示編譯棧設計，并非完備的 LLaMA 問答工具
• [Frontend] Buddy Dynamo Compiler[3]
• [Midend] 集成面向矩陣乘法的向量化以及面向循環的并行優化
• [Backend] 端到端示例在 X86 AVX512 機器上進行測試（Ubuntu 22.04）
• [WIP] 開發并集成各種優化（現在速度太慢）
• [WIP] 在多種 CPU 后端進行測試（Arm Neon, RVV, etc.）
• [WIP] 對接多種硬件平臺（GPU, Gemmini Accelerator, etc.）
• [WIP] 增加前端覆蓋程度（Stable Diffusion, CLIP, Whisper, etc.）

概述

AI 大模型的爆發為軟硬件設計帶來了新的抓手和機會。隨著模型的規模、類型、模態的增加和發散，單一軟硬件技術棧覆蓋各種場景的能力越來越受限。這種趨勢也加深了我對軟硬件生態重要性的認識。在我看來，整個生態的最重要的三條設計原則如下（重要程度由高到低）：

• 技術路線標準化（生態的立足之本）
• 上手門檻低（足夠多的貢獻者和玩家是生態繁榮的關鍵）
• 優化上限高（決定了生態的發展潛力）

我們基于以上思考打造了 Buddy Compiler， 致力于實現軟硬件協同設計生態。我們的目標是實現從領域特定編程語言（DSL）到領域特定硬件架構（DSA）的編譯流程和協同設計。

本文將介紹如何使用 Buddy Compiler 完成 LLaMA 2 的端到端推理。同時，我們也會分享 Buddy Compiler 的整體設計理念和未來的規劃。具體的構建流程請參閱此處的文檔[4]（大模型的構建需要十足的耐心和一臺性能不錯的機器）我們基于 PyTorch 和 MLIR 打通了 LLaMA 端到端的推理通路，但是尚未進行完整的性能分析和優化。

目前，我們只應用了針對矩陣乘法的向量化優化，以及針對循環的并行計算優化。

優化和調優的策略仍在開發階段，因此目前的性能還處于較低水平。我們的初步目標并不是追求極致的性能，而是建立一個標準的端到端通路，在此基礎上再做各層級性能優化和調優。

技術路線

技術路線的標準化不僅是我們努力追求的核心原則，而且我們堅信這能夠吸引更多的貢獻者和合作伙伴，同時還能夠有效降低用戶的學習門檻和后續的維護成本。

如圖所示，我們選擇 PyTorch 生態對接各種 AI 模型，選擇 MLIR 生態作為 Buddy Compiler 的中間表示。

我們將 Buddy Compiler 作為 Torch Dynamo 的自定義編譯器組成整個編譯棧，從而希望實現將各種 AI 模型映射到多種硬件架構上的愿景。以下是一些設計點：

使用TorchDynamo作為Trace工具對接AI模型

TorchDynamo 使用 CPython 的 Frame Evaluation API 特性能夠做到更加準確的 PyTorch 圖的捕捉，詳情可以參閱 PyTorch 的文檔[5]。除此之外，PyTorch 2.x 提供了全面的編譯支持，尤其是提供了非常友好的自定義編譯器對接方案[6]。因此，PyTorch 2.x TorchDynamo 成為了我們對接 AI 模型的不二選擇。

選擇Aten IR作為對接層級

根據 PyTorch 的文檔[5]，Core Aten IR 是服務于對接后端的算子集合，它相比 Prime IR 抽象級別也更高。我們傾向于使用較高的抽象級別映射到 MLIR 從而有更多的空間和信息來進行多層優化，因此我們選擇 Aten IR 作為對接層級，可以很方便地映射到 Linalg Dialect 和 TOSA Dialect 的各種操作。

使用MLIR Python Bindings實現Dynamo Compiler生成TOSA/Linalg Ops

Buddy Compiler 前端中的 Dynamo Compiler（或者叫做 Dynamo Importer）的作用是將 PyTorch 的 Aten IR 轉換到 MLIR，同時也負責將模型參數進行處理和打包。

Dynamo Compiler 遍歷從 TorchDynamo 傳入的 FX Graph，并且針對每個 Aten IR 節點進行轉換，轉換過程使用 MLIR Python Bindings 生成目標 MLIR 代碼。

Dynamo Compiler 支持對接 Dialect 的優先級設定，即可以選擇 Linalg Dialect 或者 TOSA Dialect 作為首選項進行轉換。

最終，Dynamo Compiler 負責輸出轉換后的 MLIR Module 和權重參數。

使用Buddy Compiler工具鏈進行優化和下降

我們的整個編譯通路目前并沒有完全使用 Python 腳本完成，而是使用 CMake 將前、中、后端集成起來。

這一定程度上簡化了并解耦了前、中、后端的開發流程。編譯優化和下降的所有 Pass 注冊到 buddy-opt 工具中，它包括所有的上游 Pass 和 Buddy Compiler 中的優化 Pass，因此 buddy-opt 是上游 mlir-opt 的超集。面向通用硬件的編譯流程我們使用 MLIR Core Dialect 進行實現，從而達成最大化的復用，我們的工具也因此和所有 LLVM/MLIR 的工具兼容，例如 mlir-translate, llc 等等。

目前我們針對循環采用并行計算的優化，其中相關中間表示和優化 Pass 完全來自上游的 OMP Dialect，可以直接復用并帶來不錯的優化效果。從此也可以看出統一生態的優勢。此外，我們針對粗顆粒度的 Operations 設計了向量化算法進行編譯優化，使用 Vector Dialect 也可以實現跨 SIMD/Vector 平臺的效果。

如果希望面向特定加速器（例如 Gemmini Accelerator）生成代碼，也可以使用 buddy-translate 和 buddy-llc生成到特定于硬件的 LLVM IR，從而最終生成加速器的硬件指令。

Forward函數搭配Buddy Compiler Text Container完成端到端推理

在完成編譯優化和下降之后，模型的 Forward 函數將會被構建為共享庫。

由于我們很多場景需要在模擬器和開發平臺上測試，因此我們沒有選擇將執行的流程交給 Python 生態，而是進行 AOT 編譯，生成可執行文件。

為了配合從 MLIR 構建出來的 Forward 函數實現端到端的推理，我們提供了 C++ 版本的 Text Container 和 MemRef Container 來作為文本輸入輸出的 Tokenizer 和數據容器。

最終，在 C++ 的 main 函數中將輸入的文本和權重參數加載到數據容器，然后調用 Forward 函數進行推理，輸出的 Token 再交由 Text Container 進行后處理即可得到最終的文本。

未來工作

我們目前打通了 LLaMA 到 CPU SIMD/Vector 平臺的通路，使用 X86 AVX512 進行初步的測試。用于 Vector Dialect 的跨平臺性，Arm Neon 和 RISC-V Vector Extesion 也是可以天然支持的，我們正在進行廣泛測試。

同時我們也在支持嘗試將 Buddy Compiler 中的 Gemmini 加速器支持對接到大模型推理的通路上。

此外，GPU 的優化 Pass 也在開發中。在模型層面，我們希望測試并對接更多的多模態大模型，進一步提升 Buddy Dynamo Compiler 前端的覆蓋程度。在完成上述工作后，下一個里程碑是將多模態大模型編譯到多種硬件平臺上。

總的來說，接下來前、中、后端將會相對解耦地進行開發，最終進行對接。

• 前端：對接更多模型完善算子覆蓋程度。
• 中端：進行更詳細的性能分析，面向各種計算負載開發不同層級的優化。
• 后端：對 CPU 各 SIMD/Vector 平臺進行測試，完成面向 GPU / Gemmini 加速器的端到端通路。

三個部分均相對完備的時候就考慮用 Python 接口將所有工具鏈包裝起來，形成一個更為優雅的使用流程。

總結

如今大模型推理的軟件棧也層出不窮，技術路線也各不相同。我們使用的 Torch 2.x + Buddy Compiler 的編譯棧設計策略實際上是希望融合 PyTorch 和 MLIR 兩大生態進行 AI 計算。當然，在 MLIR 生態里面做大模型推理，我們認為 Torch-MLIR + IREE 目前是相對比較完備的解決方案。nod.ai 的 SHARK Turbine[7] 就使用了這種技術路線。

相比于 Torch-MLIR 搭配 IREE 的組合，Buddy Compiler 更強調 Simple But Powerful 設計，采用極致復用策略和完全代碼生成策略。相比于 Torch-MLIR 的 Torch Dialect 層級，Buddy Compiler 更偏向直接復用 TOSA/Linalg 對接 Aten IR；相比于 IREE 覆蓋一切后端的 Runtime 和 Execution Engine 的設計，Buddy Compiler 更偏向進行完全代碼生成。

Shark Turbine 封面[8]是一個飛機渦輪發動機，這和他們的技術路線非常契合，TorchDynamo + Torch-MLIR + IREE 是一個極其精密且重型的編譯棧，這樣的“發動機“理論上可以帶著他們飛躍任何高山溝壑。相比而言，Buddy Compiler 更像是電動汽車的三電平臺，可以以此為基礎打造各種性格的電動汽車。對我們來說，LLaMA 的支持不是起點也不是終點，是在探索路上偶遇的一座高山，我們希望翻過它，看看山那邊的世界，尤其是開著自己造的車！

致謝

感謝所有 Buddy Compiler 的貢獻者，特別感謝一起努力跑通 LLaMA 的伙伴：zhanghb97，weilinquan，xTayEx，EllisLambda，Lester-1，LAJIidea，SForeKeeper，LHY-24，xlinsist，qingqing12138. 同時感謝 OSPP 組委會提供的開源項目席位。

關于Intelligent Computing

Intelligent Computing是由之江實驗室和美國科學促進會（AAAS）共同創辦，是《科學》合作期刊框架中智能計算領域的第一本開放獲取（Open Access）國際期刊。期刊以“面向智能的計算、智能驅動的計算”以及“智能、數據與計算驅動的科學發現”為主題，主要刊載原創研究論文、綜述論文和觀點論文。

參考鏈接：

[1]https://github.com/buddy-compiler/buddy-mlir

[2]https://github.com/buddy-compiler/buddy-mlir/tree/main/examples/BuddyLlama#buddy-compiler-llama-example

[3]https://github.com/buddy-compiler/buddy-mlir/tree/main/frontend/Python

[4]https://github.com/buddy-compiler/buddy-mlir/tree/main/examples/BuddyLlama

[5]https://pytorch.org/docs/stable/torch.compiler.html#torch-compiler

[6]https://pytorch.org/docs/stable/torch.compiler_custom_backends.html#custom-backends

[7]https://github.com/nod-ai/SHARK-Turbine#shark-turbine

[8]https://github.com/nod-ai/SHARK-Turbine#shark-turbineg