A卡跑大模型，性能達到4090的80%，價格只有一半：陳天奇TVM團隊出品
最近，科技領域有很多人都在為算力發愁。
OpenAI CEO 奧特曼：我整天在都想著 flops。
自預訓練大模型興起以來，人們面臨的算力挑戰就變得越來越大。為此，人們為大語言模型（LLM）提出了許多訓練和推理的解決方案。顯然，大多數高性能推理解決方案都基于 CUDA 并針對英偉達 GPU 進行了優化。
但在動輒千億參數的模型體量，多家科技公司激烈競爭，以及單一供應商的合力作用下，想搶到 GPU 又變成了一件難事。
最近，微軟、OpenAI 等公司都表示正在采取必要措施來緩解用于 AI 任務的 H100、A100 專用 GPU 的短缺問題。微軟正在限制員工訪問 GPU 的時間，Quora 首席執行官表示，硬件短缺掩蓋了人工智能應用程序的真正潛力。伊隆?馬斯克還開玩笑說，企業級 GPU 比買「藥」還難。
馬斯克預測 GPT-5 大概需要三到五萬塊 H100 來訓練。
旺盛的需求除了推動英偉達的股價，使其改變生產計劃之外，也讓人們不得不去尋求其他替代方式。好消息是，圖形芯片市場上并不只有 N 卡一家。
昨天，卡耐基梅隆大學博士生侯博涵（Bohan Hou）放出了使用 AMD 顯卡進行大模型推理的新方案，立刻獲得了機器學習社區的關注。
在 CMU，侯博涵的導師是 TVM、MXNET、XGBoost 的作者陳天奇。對于這項新實踐，陳天奇表示，解決 AI 硬件短缺問題的方法還是要看軟件，讓我們帶來高性能、通用部署的開源大模型吧。
在知乎上，作者對于實現高性能 LLM 推理進行了詳細介紹：
通過這種優化方法，在最新的 Llama2 的 7B 和 13B 模型中，如果用一塊 AMD Radeon RX 7900 XTX 速度可以達到英偉達 RTX 4090 的 80%，或是 3090Ti 的 94%。
除了 ROCm 之外，這種 Vulkan 支持還允許我們把大模型的部署推廣到其他 AMD 芯片類型上，例如具有 AMD APU 的 SteamDeck。
如果粗略的比較一下規格，我們可以看到 AMD 的 RX 7900 XTX 與英偉達的 RTX 4090 和 RTX 3090 Ti 處于相近級別。
它們的顯存都在 24GB，這意味著它們可以容納相同尺寸的模型，它們都具有相似的內存帶寬。
但是在算力上，RTX 4090 的 FP16 性能比 7900 XTX 高兩倍，而 3090 Ti 的 FP16 性能比 7900 XTX 高 1.3 倍。如果只考慮延遲敏感的大模型推理，其性能主要受內存限制，因此 FP16 性能不是這里的瓶頸。
而看價格的話，RX 7900 XTX 比 RTX 4090 便宜 40% 還多（京東上看甚至有 50%），在消費級領域里前者幾乎是和 RTX 4080 對標的。
3090Ti 的價格則很難比較，畢竟那是上一代產品。但從純硬件規格的角度來看，AMD 7900 XTX 似乎與 RTX 3090 Ti 相當。
我們知道，硬件層的算力并不一定是 AMD 長期以來在機器學習上落后的原因 —— 主要差距在于缺乏相關模型的軟件支持和優化。從生態角度來看，有兩個因素已開始改變現狀：
- AMD 正在努力在 ROCm 平臺上增加投入。
- 機器學習編譯等新興技術現在有助于降低跨后端的，更通用軟件支持的總體成本。
研究人員深入討論了 AMD GPU 體系與目前流行的英偉達 GPU 上高性能 CUDA 解決方案相比的表現如何。
用 ROCm 進行機器學習編譯
機器學習編譯
機器學習編譯是一種用于編譯和自動優化機器學習模型的新興技術。MLC 解決方案不是為每個后端（如 ROCm 或 CUDA）編寫特定的算子 ，而是自動生成適用于不同后端的代碼。在這里，作者利用 MLC-LLM，一種基于機器學習編譯的解決方案，提供了 LLM 的高性能通用部署。MLC-LLM 建立在 Apache TVM Unity 之上，后者是一個機器學習編譯軟件棧，提供了基于 Python 的高效開發和通用部署。MLC-LLM 為各種后端（包括 CUDA、Metal、ROCm、Vulkan 和 OpenCL）提供了最先進的性能，涵蓋了從服務器級別 GPU 到移動設備（iPhone 和 Android）。
整體而言，MLC-LLM 允許用戶使用基于 Python 的工作流程獲取開源的大語言模型，并在包括轉換計算圖、優化 GPU 算子的張量 layout 和 schedule 以及在感興趣的平臺上本地部署時進行編譯。
面向 ROCm 的機器學習編譯技術棧。
針對 AMD GPU 和 APU 的 MLC
人們對于 A 卡用于機器學習的探索其實并不鮮見，支持 AMD GPU 有幾種可能的技術路線：ROCm、OpenCL、Vulkan 和 WebGPU。ROCm 技術棧是 AMD 最近推出的，與 CUDA 技術棧有許多相應的相似之處。Vulkan 是最新的圖形渲染標準，為各種 GPU 設備提供了廣泛的支持。WebGPU 是最新的 Web 標準，允許在 Web 瀏覽器上運行計算。
雖然有這么多可能的路線，但很少有解決方案支持除了 CUDA 之外的方法，這在很大程度上是因為復制新硬件或 GPU 編程模型的技術棧的工程成本過高。MLC-LLM 支持自動代碼生成，無需為每個 GPU 算子重新定制，從而為以上所有方法提供支持。但是，最終性能仍然取決于 GPU 運行時的質量以及在每個平臺上的可用性。
在這個案例中，作者選擇 Radeon 7900 XTX 的 ROCm 和 Steamdeck 的 APU 的 Vulkan，可以發現 ROCm 技術棧是開箱即用的。由于 TVM unity 中具有高效的基于 Python 的開發流程，花費了若干小時來進一步提供 ROCm 的性能優化。具體來說，研究人員采取了以下措施來提供 ROCm 支持：
- 重用現有后端（如 CUDA 和 Metal）的整個 MLC 流水線，包括內存規劃、算子融合等。
- 重用 TVM TensorIR 中的通用 GPU 算子優化空間，并將其后端選為 AMD GPU
- 重用 TVM 的 ROCm 代碼生成流程，通過 LLVM 生成 ROCm 代碼。
- 最后，將生成的代碼導出為可以由 CLI、Python 和 REST API 調用的共享或靜態庫。
使用 MLC Python 包進行性能測試
作者使用 4 bit 量化對 Llama 2 7B 和 13B 進行了性能測試。通過設置 prompt 長度為 1 個 token 并生成 512 個 token 來測量 decoding 的性能。所有結果都是在 batch size=1 的情況下測試。
AMD RX 7900 XTX 與 NVIDIA RTX 4090 和 3090 Ti 的性能對比。
基于 ROCm5.6，AMD 7900 XTX 可以達到 NVIDIA 4090 速度的 80%。
關于 CUDA 性能說明：在這里 CUDA baseline 的性能如何？據我們所知，MLC-LLM 是 CUDA 上大語言模型推理的最優解決方案。但作者相信它仍然有改進的空間，例如通過更好的 attention 算子優化。一旦這些優化在 MLC 中實現，預計 AMD 和 NVIDIA 的數據都會有所改善。
如果這些優化僅在 N 卡那里實施，將使差距從 20% 增加到 30%。因此，在查看這些數字時，作者建議放置 10% 的誤差。
自行嘗試
該項目提供了預構建的安裝包和使用說明，以便用戶在自己的設備上復現新的結果。要運行這些性能測試，請確保你的 Linux 上有安裝了 ROCm 5.6 或更高版本的 AMD GPU。按照這里的說明（https://mlc.ai/mlc-llm/docs/get_started/try_out.html）安裝啟用了 ROCm 的預構建 MLC pacakge。
運行以下 Python 腳本，需要使用 MLC package 來復現性能數據：
from mlc_chat import ChatModule
# Create a ChatModule instance that loads from `./dist/prebuilt/Llama-2-7b-chat-hf-q4f16_1`cm = ChatModule(model="Llama-2-7b-chat-hf-q4f16_1")
# Run the benchmarksoutput = cm.benchmark_generate("Hi", generate_length=512)
print(f"Generated text:\n{output}\n")
print(f"Statistics: {cm.stats()}")
# Reset the chat module by
# cm.reset_chat()MLC-LLM 還提供了一個命令行界面 CLI，允許用戶與模型進行交互式聊天。對于 ROCm，需要從源代碼構建 CLI。請按照這里的說明（https://mlc.ai/mlc-llm/docs/deploy/cli.html#option-2-build-mlc-runtime-from-source）從源代碼構建 CLI。
使用統一內存在 SteamDeck 上運行 Vulkan
作者表示，還將考察更廣泛的 AMD 設備，更具體地說，是搭載了 AMD APU 的 Steam Deck。雖然在 BIOS 中，ROCm 中可用的 GPU VRAM 被限制為 4GB，但 Mesa Vulkan 驅動程序具有強大的支持，允許緩沖區超越上限，使用統一內存最多可達 16GB，足以運行 4 位量化的 Llama-7B。
在 Steam Deck 上與大語言模型進行交互。
這些結果為支持更多不同類型的消費者提供了一些啟示。
討論和未來的方向
身處生成式 AI 的時代，硬件可用性已經成為一個迫切需要解決的問題。ML 編譯可以通過在硬件后端之間提供高性能的通用部署，從而提高硬件的可用性。
鑒于本文所展現的數據，作者認為在適當的價格和可用性條件下，AMD GPU 可以開始用于 LLM 推理。
在陳天奇團隊，研究目前的重點關注消費級 GPU。作者表示，根據過往經驗，針對消費級 GPU 型號的 MLC 優化通常可以推廣到云 GPU（例如從 RTX 4090 到 A100 和 A10g），有信心讓該解決方案在云和消費級 AMD 和 NVIDIA GPU 之間具有普適性，并將在獲得更多 GPU 訪問權限后更新研究。與此同時，作者期待研究社區在 MLC 通用部署流程的基礎上構建解決方案。
本文是通過 MLC 來支持高效通用的機器學習部署的研究的一個階段性努力，研究人員也正積極地在以下幾個方向上努力推廣成果：
- 啟用 batching 和多 GPU 支持（對服務器端推理尤為重要）；
- 與 PyTorch 生態系統集成；
- 支持更多量化和模型架構；
- 在更多硬件后端上進行更多自動優化。
機器學習系統工程是一個持續的問題，在持續創新方面，英偉達仍然處于領先地位，作者預計隨著新硬件（如 H100）以及更重要的軟件演變，整個領域將發生變化。因此，關鍵問題不僅是現在構建正確的解決方案，還包括如何不斷趕上并將機器學習工程引入新平臺。在這個過程中，機器學習工程的生產力是關鍵。
由于基于 Python 的 ML 編譯開發流程，我們可以在幾小時內獲得 ROCm 優化的支持。預計此次提出的新方法，在探索更多關于通用部署的想法并解決硬件可用性問題時會變得更加有用。
相關資源
該項目已經在 GitHub 上發布。有關如何嘗試 MLC LLM 部署的詳細指南，請參閱項目頁面。MLC LLM 的源代碼可在官方 GitHub 上找到。
- 項目頁面：https://mlc.ai/mlc-llm/docs/
- GitHub：https://github.com/mlc-ai/mlc-llm/
