最近，科技領(lǐng)域有很多人都在為算力發(fā)愁。

OpenAI CEO 奧特曼：我整天在都想著 flops。

自預(yù)訓(xùn)練大模型興起以來，人們面臨的算力挑戰(zhàn)就變得越來越大。為此，人們?yōu)榇笳Z言模型（LLM）提出了許多訓(xùn)練和推理的解決方案。顯然，大多數(shù)高性能推理解決方案都基于 CUDA 并針對英偉達(dá) GPU 進(jìn)行了優(yōu)化。

但在動(dòng)輒千億參數(shù)的模型體量，多家科技公司激烈競爭，以及單一供應(yīng)商的合力作用下，想搶到 GPU 又變成了一件難事。

最近，微軟、OpenAI 等公司都表示正在采取必要措施來緩解用于 AI 任務(wù)的 H100、A100 專用 GPU 的短缺問題。微軟正在限制員工訪問 GPU 的時(shí)間，Quora 首席執(zhí)行官表示，硬件短缺掩蓋了人工智能應(yīng)用程序的真正潛力。伊隆?馬斯克還開玩笑說，企業(yè)級 GPU 比買「藥」還難。

馬斯克預(yù)測 GPT-5 大概需要三到五萬塊 H100 來訓(xùn)練。

旺盛的需求除了推動(dòng)英偉達(dá)的股價(jià)，使其改變生產(chǎn)計(jì)劃之外，也讓人們不得不去尋求其他替代方式。好消息是，圖形芯片市場上并不只有 N 卡一家。

昨天，卡耐基梅隆大學(xué)博士生侯博涵（Bohan Hou）放出了使用 AMD 顯卡進(jìn)行大模型推理的新方案，立刻獲得了機(jī)器學(xué)習(xí)社區(qū)的關(guān)注。

在 CMU，侯博涵的導(dǎo)師是 TVM、MXNET、XGBoost 的作者陳天奇。對于這項(xiàng)新實(shí)踐，陳天奇表示，解決 AI 硬件短缺問題的方法還是要看軟件，讓我們帶來高性能、通用部署的開源大模型吧。

在知乎上，作者對于實(shí)現(xiàn)高性能 LLM 推理進(jìn)行了詳細(xì)介紹：

通過這種優(yōu)化方法，在最新的 Llama2 的 7B 和 13B 模型中，如果用一塊 AMD Radeon RX 7900 XTX 速度可以達(dá)到英偉達(dá) RTX 4090 的 80%，或是 3090Ti 的 94%。

除了 ROCm 之外，這種 Vulkan 支持還允許我們把大模型的部署推廣到其他 AMD 芯片類型上，例如具有 AMD APU 的 SteamDeck。

如果粗略的比較一下規(guī)格，我們可以看到 AMD 的 RX 7900 XTX 與英偉達(dá)的 RTX 4090 和 RTX 3090 Ti 處于相近級別。

它們的顯存都在 24GB，這意味著它們可以容納相同尺寸的模型，它們都具有相似的內(nèi)存帶寬。

但是在算力上，RTX 4090 的 FP16 性能比 7900 XTX 高兩倍，而 3090 Ti 的 FP16 性能比 7900 XTX 高 1.3 倍。如果只考慮延遲敏感的大模型推理，其性能主要受內(nèi)存限制，因此 FP16 性能不是這里的瓶頸。

而看價(jià)格的話，RX 7900 XTX 比 RTX 4090 便宜 40% 還多（京東上看甚至有 50%），在消費(fèi)級領(lǐng)域里前者幾乎是和 RTX 4080 對標(biāo)的。

3090Ti 的價(jià)格則很難比較，畢竟那是上一代產(chǎn)品。但從純硬件規(guī)格的角度來看，AMD 7900 XTX 似乎與 RTX 3090 Ti 相當(dāng)。

我們知道，硬件層的算力并不一定是 AMD 長期以來在機(jī)器學(xué)習(xí)上落后的原因 —— 主要差距在于缺乏相關(guān)模型的軟件支持和優(yōu)化。從生態(tài)角度來看，有兩個(gè)因素已開始改變現(xiàn)狀：

• AMD 正在努力在 ROCm 平臺上增加投入。
• 機(jī)器學(xué)習(xí)編譯等新興技術(shù)現(xiàn)在有助于降低跨后端的，更通用軟件支持的總體成本。

研究人員深入討論了 AMD GPU 體系與目前流行的英偉達(dá) GPU 上高性能 CUDA 解決方案相比的表現(xiàn)如何。

用 ROCm 進(jìn)行機(jī)器學(xué)習(xí)編譯

機(jī)器學(xué)習(xí)編譯

機(jī)器學(xué)習(xí)編譯是一種用于編譯和自動(dòng)優(yōu)化機(jī)器學(xué)習(xí)模型的新興技術(shù)。MLC 解決方案不是為每個(gè)后端（如 ROCm 或 CUDA）編寫特定的算子 ，而是自動(dòng)生成適用于不同后端的代碼。在這里，作者利用 MLC-LLM，一種基于機(jī)器學(xué)習(xí)編譯的解決方案，提供了 LLM 的高性能通用部署。MLC-LLM 建立在 Apache TVM Unity 之上，后者是一個(gè)機(jī)器學(xué)習(xí)編譯軟件棧，提供了基于 Python 的高效開發(fā)和通用部署。MLC-LLM 為各種后端（包括 CUDA、Metal、ROCm、Vulkan 和 OpenCL）提供了最先進(jìn)的性能，涵蓋了從服務(wù)器級別 GPU 到移動(dòng)設(shè)備（iPhone 和 Android）。

整體而言，MLC-LLM 允許用戶使用基于 Python 的工作流程獲取開源的大語言模型，并在包括轉(zhuǎn)換計(jì)算圖、優(yōu)化 GPU 算子的張量 layout 和 schedule 以及在感興趣的平臺上本地部署時(shí)進(jìn)行編譯。

面向 ROCm 的機(jī)器學(xué)習(xí)編譯技術(shù)棧。

針對 AMD GPU 和 APU 的 MLC

人們對于 A 卡用于機(jī)器學(xué)習(xí)的探索其實(shí)并不鮮見，支持 AMD GPU 有幾種可能的技術(shù)路線：ROCm、OpenCL、Vulkan 和 WebGPU。ROCm 技術(shù)棧是 AMD 最近推出的，與 CUDA 技術(shù)棧有許多相應(yīng)的相似之處。Vulkan 是最新的圖形渲染標(biāo)準(zhǔn)，為各種 GPU 設(shè)備提供了廣泛的支持。WebGPU 是最新的 Web 標(biāo)準(zhǔn)，允許在 Web 瀏覽器上運(yùn)行計(jì)算。

雖然有這么多可能的路線，但很少有解決方案支持除了 CUDA 之外的方法，這在很大程度上是因?yàn)閺?fù)制新硬件或 GPU 編程模型的技術(shù)棧的工程成本過高。MLC-LLM 支持自動(dòng)代碼生成，無需為每個(gè) GPU 算子重新定制，從而為以上所有方法提供支持。但是，最終性能仍然取決于 GPU 運(yùn)行時(shí)的質(zhì)量以及在每個(gè)平臺上的可用性。

在這個(gè)案例中，作者選擇 Radeon 7900 XTX 的 ROCm 和 Steamdeck 的 APU 的 Vulkan，可以發(fā)現(xiàn) ROCm 技術(shù)棧是開箱即用的。由于 TVM unity 中具有高效的基于 Python 的開發(fā)流程，花費(fèi)了若干小時(shí)來進(jìn)一步提供 ROCm 的性能優(yōu)化。具體來說，研究人員采取了以下措施來提供 ROCm 支持：

• 重用現(xiàn)有后端（如 CUDA 和 Metal）的整個(gè) MLC 流水線，包括內(nèi)存規(guī)劃、算子融合等。
• 重用 TVM TensorIR 中的通用 GPU 算子優(yōu)化空間，并將其后端選為 AMD GPU
• 重用 TVM 的 ROCm 代碼生成流程，通過 LLVM 生成 ROCm 代碼。
• 最后，將生成的代碼導(dǎo)出為可以由 CLI、Python 和 REST API 調(diào)用的共享或靜態(tài)庫。

使用 MLC Python 包進(jìn)行性能測試

作者使用 4 bit 量化對 Llama 2 7B 和 13B 進(jìn)行了性能測試。通過設(shè)置 prompt 長度為 1 個(gè) token 并生成 512 個(gè) token 來測量 decoding 的性能。所有結(jié)果都是在 batch size=1 的情況下測試。

AMD RX 7900 XTX 與 NVIDIA RTX 4090 和 3090 Ti 的性能對比。

基于 ROCm5.6，AMD 7900 XTX 可以達(dá)到 NVIDIA 4090 速度的 80%。

關(guān)于 CUDA 性能說明：在這里 CUDA baseline 的性能如何？據(jù)我們所知，MLC-LLM 是 CUDA 上大語言模型推理的最優(yōu)解決方案。但作者相信它仍然有改進(jìn)的空間，例如通過更好的 attention 算子優(yōu)化。一旦這些優(yōu)化在 MLC 中實(shí)現(xiàn)，預(yù)計(jì) AMD 和 NVIDIA 的數(shù)據(jù)都會(huì)有所改善。

如果這些優(yōu)化僅在 N 卡那里實(shí)施，將使差距從 20% 增加到 30%。因此，在查看這些數(shù)字時(shí)，作者建議放置 10% 的誤差。

自行嘗試

該項(xiàng)目提供了預(yù)構(gòu)建的安裝包和使用說明，以便用戶在自己的設(shè)備上復(fù)現(xiàn)新的結(jié)果。要運(yùn)行這些性能測試，請確保你的 Linux 上有安裝了 ROCm 5.6 或更高版本的 AMD GPU。按照這里的說明（https://mlc.ai/mlc-llm/docs/get_started/try_out.html）安裝啟用了 ROCm 的預(yù)構(gòu)建 MLC pacakge。

運(yùn)行以下 Python 腳本，需要使用 MLC package 來復(fù)現(xiàn)性能數(shù)據(jù)：

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 還提供了一個(gè)命令行界面 CLI，允許用戶與模型進(jìn)行交互式聊天。對于 ROCm，需要從源代碼構(gòu)建 CLI。請按照這里的說明（https://mlc.ai/mlc-llm/docs/deploy/cli.html#option-2-build-mlc-runtime-from-source）從源代碼構(gòu)建 CLI。

使用統(tǒng)一內(nèi)存在 SteamDeck 上運(yùn)行 Vulkan

作者表示，還將考察更廣泛的 AMD 設(shè)備，更具體地說，是搭載了 AMD APU 的 Steam Deck。雖然在 BIOS 中，ROCm 中可用的 GPU VRAM 被限制為 4GB，但 Mesa Vulkan 驅(qū)動(dòng)程序具有強(qiáng)大的支持，允許緩沖區(qū)超越上限，使用統(tǒng)一內(nèi)存最多可達(dá) 16GB，足以運(yùn)行 4 位量化的 Llama-7B。

在 Steam Deck 上與大語言模型進(jìn)行交互。

這些結(jié)果為支持更多不同類型的消費(fèi)者提供了一些啟示。

討論和未來的方向

身處生成式 AI 的時(shí)代，硬件可用性已經(jīng)成為一個(gè)迫切需要解決的問題。ML 編譯可以通過在硬件后端之間提供高性能的通用部署，從而提高硬件的可用性。

鑒于本文所展現(xiàn)的數(shù)據(jù)，作者認(rèn)為在適當(dāng)?shù)膬r(jià)格和可用性條件下，AMD GPU 可以開始用于 LLM 推理。

在陳天奇團(tuán)隊(duì)，研究目前的重點(diǎn)關(guān)注消費(fèi)級 GPU。作者表示，根據(jù)過往經(jīng)驗(yàn)，針對消費(fèi)級 GPU 型號的 MLC 優(yōu)化通?？梢酝茝V到云 GPU（例如從 RTX 4090 到 A100 和 A10g），有信心讓該解決方案在云和消費(fèi)級 AMD 和 NVIDIA GPU 之間具有普適性，并將在獲得更多 GPU 訪問權(quán)限后更新研究。與此同時(shí)，作者期待研究社區(qū)在 MLC 通用部署流程的基礎(chǔ)上構(gòu)建解決方案。

本文是通過 MLC 來支持高效通用的機(jī)器學(xué)習(xí)部署的研究的一個(gè)階段性努力，研究人員也正積極地在以下幾個(gè)方向上努力推廣成果：

• 啟用 batching 和多 GPU 支持（對服務(wù)器端推理尤為重要）；
• 與 PyTorch 生態(tài)系統(tǒng)集成；
• 支持更多量化和模型架構(gòu)；
• 在更多硬件后端上進(jìn)行更多自動(dòng)優(yōu)化。

機(jī)器學(xué)習(xí)系統(tǒng)工程是一個(gè)持續(xù)的問題，在持續(xù)創(chuàng)新方面，英偉達(dá)仍然處于領(lǐng)先地位，作者預(yù)計(jì)隨著新硬件（如 H100）以及更重要的軟件演變，整個(gè)領(lǐng)域?qū)l(fā)生變化。因此，關(guān)鍵問題不僅是現(xiàn)在構(gòu)建正確的解決方案，還包括如何不斷趕上并將機(jī)器學(xué)習(xí)工程引入新平臺。在這個(gè)過程中，機(jī)器學(xué)習(xí)工程的生產(chǎn)力是關(guān)鍵。

由于基于 Python 的 ML 編譯開發(fā)流程，我們可以在幾小時(shí)內(nèi)獲得 ROCm 優(yōu)化的支持。預(yù)計(jì)此次提出的新方法，在探索更多關(guān)于通用部署的想法并解決硬件可用性問題時(shí)會(huì)變得更加有用。

相關(guān)資源

該項(xiàng)目已經(jīng)在 GitHub 上發(fā)布。有關(guān)如何嘗試 MLC LLM 部署的詳細(xì)指南，請參閱項(xiàng)目頁面。MLC LLM 的源代碼可在官方 GitHub 上找到。

• 項(xiàng)目頁面：https://mlc.ai/mlc-llm/docs/
• GitHub：https://github.com/mlc-ai/mlc-llm/