介紹

我們正處在大型語言模型和生成式人工智能的時(shí)代，其應(yīng)用場(chǎng)景有可能改變每個(gè)人的生活。開放式大型語言模型為提供定制化和領(lǐng)域特定部署帶來了重大機(jī)會(huì)。

我們正處在一個(gè)激動(dòng)人心的開放模型開發(fā)年份。一方面，我們看到了在（云）服務(wù)器部署方面取得的令人振奮的進(jìn)展，這些解決方案可以為具有多個(gè)GPU的更大模型提供同時(shí)服務(wù)的能力。同時(shí)，我們也開始看到在設(shè)備本地部署方面取得了進(jìn)展，能夠?qū)⑿阅軓?qiáng)大的量化模型部署到筆記本電腦、瀏覽器和手機(jī)上。未來將會(huì)走向何方？我們相信未來是混合的，因此重要的是要使任何人都能在云端和本地環(huán)境中運(yùn)行LLM。

許多LLM推斷項(xiàng)目，包括我們過去版本的MLC LLM工作，為服務(wù)器和本地使用案例提供了不同的解決方案，具有不同的實(shí)現(xiàn)和優(yōu)化。例如，服務(wù)器解決方案通常支持連續(xù)批處理和更好的多GPU支持，而本地解決方案則帶來了跨平臺(tái)更好的可移植性。然而，我們認(rèn)為有必要將所有技術(shù)匯集到一起。在一個(gè)使用案例的一側(cè)出現(xiàn)的許多技術(shù)直接適用于另一側(cè)。雖然像連續(xù)批處理這樣的技術(shù)在當(dāng)前某些本地使用案例中可能并不實(shí)用，但隨著LLM成為操作系統(tǒng)的關(guān)鍵組件并支持多個(gè)請(qǐng)求以啟用代理任務(wù)，它們將變得有價(jià)值。我們想提出一個(gè)問題：是否可能構(gòu)建一個(gè)單一統(tǒng)一的LLM引擎，可在服務(wù)器和本地使用案例中運(yùn)行？

在這篇文章中，我們介紹MLC LLM引擎（簡稱MLCEngine），這是一個(gè)通用的LLM部署引擎。MLCEngine引入了一個(gè)單一引擎，用于在服務(wù)器上進(jìn)行高吞吐量、低延遲的服務(wù)，同時(shí)將小型而性能強(qiáng)大的模型無縫集成到多樣化的本地環(huán)境中。

圖1. MLCEngine：通用LLM部署引擎 實(shí)現(xiàn)通用部署的路徑具有獨(dú)特的挑戰(zhàn)。首先，它需要支持廣泛的GPU編程模型和運(yùn)行時(shí)，以在每個(gè)支持的平臺(tái)上實(shí)現(xiàn)加速。這樣的過程通常需要大量重復(fù)的工程工作。我們需要在可能的情況下利用有效的供應(yīng)商庫，但也需要支持缺乏標(biāo)準(zhǔn)庫支持的新興平臺(tái)，如Vulkan或WebGPU。此外，每個(gè)應(yīng)用平臺(tái)都有不同的編程語言環(huán)境，這增加了將引擎引入到Swift、Kotlin、Javascript等環(huán)境中的復(fù)雜性。

為了克服平臺(tái)支持的挑戰(zhàn)，我們利用了通過Apache TVM進(jìn)行機(jī)器學(xué)習(xí)編譯的技術(shù)，自動(dòng)生成便攜式GPU庫，以適用于各種硬件和平臺(tái)。此外，我們構(gòu)建了一個(gè)便攜式運(yùn)行時(shí)架構(gòu)，將業(yè)界領(lǐng)先的LLM服務(wù)優(yōu)化（如連續(xù)批處理、推測(cè)解碼、前綴緩存等）與對(duì)云端和本地平臺(tái)的最大可移植性相結(jié)合。

圖2. MLC LLM工作流程 編譯的模型庫與一個(gè)通用運(yùn)行時(shí)共同工作，該運(yùn)行時(shí)在不同的主機(jī)應(yīng)用語言中提供直接的語言綁定，這些語言共享相同的OpenAI風(fēng)格的聊天完成API。所有本地應(yīng)用環(huán)境共享與云服務(wù)器設(shè)置相同的核心引擎，但針對(duì)每個(gè)個(gè)體設(shè)置進(jìn)行了專門配置。

通過Chat開始

開始使用MLCEngine的最簡單方式是通過聊天CLI界面。我們?yōu)榇颂峁┝艘粋€(gè)命令行聊天界面。下面是一個(gè)示例命令，它在終端中啟動(dòng)聊天CLI，運(yùn)行一個(gè)4位量化的Llama3 8B模型。

mlc_llm chat HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC

圖3. Chat CLI 在幕后，我們執(zhí)行即時(shí)模型編譯，為每個(gè)平臺(tái)生成GPU代碼。因此，它可以跨GPU類型和操作系統(tǒng)工作。

我們支持廣泛的模型系列，包括Llama3、Mistral/Mixtral、Phi3、Qwen2、Gemma等等。

通過OpenAI風(fēng)格API使用MLCEngine

隨著我們引入通用LLM引擎，設(shè)計(jì)一組開發(fā)人員熟悉且易于使用的API是非常重要的。我們選擇在所有環(huán)境中采用OpenAI風(fēng)格API。本節(jié)將逐步介紹這些API。

• 云端的REST API服務(wù)器

LLM服務(wù)是最典型的LLM應(yīng)用之一。LLM服務(wù)器通過運(yùn)行內(nèi)部LLM引擎來處理從網(wǎng)絡(luò)接收到的所有請(qǐng)求。我們提供了一個(gè)具有完整OpenAI API的REST服務(wù)器，用于處理LLM生成請(qǐng)求。下面是一個(gè)示例命令，它在本地主機(jī)上啟動(dòng)REST服務(wù)器，以提供4位量化的Llama3 8B模型。

圖4. 啟動(dòng)LLM服務(wù)器并發(fā)送聊天完成請(qǐng)求 為了適應(yīng)不同的服務(wù)場(chǎng)景，我們?yōu)榉?wù)器提供了三種模式：“服務(wù)器”，“本地”和“交互式”。 “服務(wù)器”模式最大化了GPU內(nèi)存的使用，并支持更高的并發(fā)性，而“本地”和“交互式”模式則限制了并發(fā)請(qǐng)求的數(shù)量，并且使用更少的GPU內(nèi)存。我們建議在專用服務(wù)器GPU（如A100/H100）上使用“服務(wù)器”模式，利用完整的FP16模型和FP8變體。對(duì)于本地環(huán)境，建議使用“本地”模式和4位量化模型設(shè)置。

• Python API

MLC LLM在Python中提供MLCEngine和AsyncMLCEngine，分別用于同步和異步LLM生成。這兩種引擎支持與OpenAI Python軟件包相同的API。

圖5. AsyncMLCEngine的Python API

• iOS SDK

為了使LLM在iPhone和iPad上運(yùn)行，我們?cè)赟wift SDK中公開了MLCEngine。雖然沒有來自O(shè)penAI的官方Swift API，但我們緊密模擬了Python API，因此我們可以找到與Python代碼鏡像的代碼，具有結(jié)構(gòu)化的輸入和輸出。Swift API還有效地利用了AsyncStream，以啟用生成內(nèi)容的異步流式傳輸。

圖6. MLCEngine的iOS/Swift API

• Android SDK

除了iOS，我們還希望為Android提供相同的體驗(yàn)。我們選擇使用基于OpenAI API的Kotlin API。下面是一個(gè)在三星S23上運(yùn)行4位量化Phi-3模型的聊天完成Kotlin API和MLC Chat應(yīng)用程序截圖的示例。

圖7. MLCEngine的Android/Kotlin API

• WebLLM SDK

近年來，我們看到了WebGPU生態(tài)系統(tǒng)的不斷成熟。WebGPU通過將WGSL（WebGPU著色語言）著色器轉(zhuǎn)換為本地GPU著色器來工作。這使得在瀏覽器中執(zhí)行GPU計(jì)算成為可能，這是快速的（本地GPU加速），方便的（無需任何環(huán)境設(shè)置）和私密的（100％客戶端計(jì)算）。MLCEngine也可以通過WebLLM項(xiàng)目在JavaScript中訪問。

圖8. MLCEngine的JavaScript/TypeScript API

• 討論

我們可能意識(shí)到的一件事是，這些API遵循完全相同的風(fēng)格engine.chat.completions.create，因此我們可以在這些平臺(tái)上擁有相同的開發(fā)體驗(yàn)。我們盡可能多地利用語言本地功能來創(chuàng)建結(jié)構(gòu)化的輸入和輸出，并啟用異步流式傳輸，以便引擎活動(dòng)不會(huì)阻塞UI。

高效的結(jié)構(gòu)化生成

除了構(gòu)建聊天應(yīng)用程序外，我們還將使應(yīng)用程序超越聊天。結(jié)構(gòu)化生成大大提高了LLM的能力，超越了基本的聊天或純文本生成。通過可控的結(jié)構(gòu)化生成，LLM能夠成為標(biāo)準(zhǔn)工具，并且可以更好地集成到生產(chǎn)中的其他應(yīng)用程序中。在各種場(chǎng)景中，JSON是最常用的結(jié)構(gòu)化格式之一，并且對(duì)于LLM引擎支持JSON字符串的結(jié)構(gòu)化生成具有重要意義。

MLCEngine提供了最先進(jìn)的JSON模式結(jié)構(gòu)化生成。對(duì)于每個(gè)請(qǐng)求，MLCEngine運(yùn)行和維護(hù)一個(gè)高性能的GrammarBNF狀態(tài)機(jī)，該狀態(tài)機(jī)在自回歸生成期間約束響應(yīng)格式。

MLCEngine支持兩種級(jí)別的JSON模式：通用JSON響應(yīng)和JSON模式定制。通用JSON模式將響應(yīng)約束為符合JSON語法。要使用通用JSON模式，請(qǐng)將參數(shù)response_format={"type": "json_object"}傳遞給聊天完成。以下是使用JSON模式的請(qǐng)求示例：

圖9. MLCEngine中的JSON模式 

此外，MLCEngine允許為每個(gè)單獨(dú)的請(qǐng)求定制響應(yīng)的JSON模式。當(dāng)提供了JSON模式時(shí)，MLCEngine將生成嚴(yán)格符合該模式的響應(yīng)。以下是帶有自定義JSON模式的請(qǐng)求示例：

圖10. MLCEngine中指定JSON模式的JSON模式 重要的是，結(jié)構(gòu)化生成支持內(nèi)置于引擎中，這意味著它可以在MLCEngine支持的所有API平臺(tái)上使用。

使LLM在各種平臺(tái)上可訪問

我們的使命之一是使LLM在各種平臺(tái)上可訪問。我們已經(jīng)討論了服務(wù)器GPU、Android和iOS上的用例。我們驗(yàn)證了MLCEngine可以在各種平臺(tái)上運(yùn)行。包括但不限于：

• NVIDIA RTX 4090，
• NVIDIA Jetson Orin，
• NVIDIA T4（在Google Colab中），
• AMD 7900 XTX，
• Steam Deck，
• Orange Pi。

其中，Steam Deck是一個(gè)有趣的平臺(tái)，因?yàn)樗惶峁┯邢薜腉PU驅(qū)動(dòng)程序支持。我們生成了Vulkan代碼，在其上運(yùn)行了一個(gè)Llama3-8B模型。

圖11. MLCEngine在Steam Deck上運(yùn)行Llama 3 8B

我們發(fā)現(xiàn)機(jī)器學(xué)習(xí)編譯幫助我們減少了在構(gòu)建MLCEngine上的工程工作量，并幫助我們?cè)诟鼜V泛的硬件范圍內(nèi)提供高性能解決方案，使這些模型對(duì)更多人可用。

優(yōu)化引擎性能

雖然這是MLCEngine的第一個(gè)發(fā)布里程碑，但我們已經(jīng)付出了大量努力來優(yōu)化其性能。引擎結(jié)合了幾個(gè)關(guān)鍵的系統(tǒng)改進(jìn)，如連續(xù)批處理、推測(cè)解碼、分頁KV管理、常見前綴緩存和級(jí)聯(lián)推斷。我們利用FlashInfer庫在CUDA平臺(tái)上計(jì)算快速的注意力，并通過編譯器代碼生成擴(kuò)展其支持到更廣泛的平臺(tái)。

MLCEngine支持多GPU開箱即用。以下命令在兩個(gè)GPU上啟動(dòng)REST服務(wù)器：

mlc_llm serve HF://mlc-ai/Qwen2-72B-Instruct-q0f16-MLC --overrides "tensor_parallel_shards=2"

圖12. 2個(gè)NVIDIA RTX 4090上的REST服務(wù)器

圖13. Python API 2x AMD 7900 XTX

MLCEngine在服務(wù)器使用案例中可以實(shí)現(xiàn)高競(jìng)爭(zhēng)性的性能，特別是在高吞吐低延遲設(shè)置下，對(duì)多個(gè)GPU的強(qiáng)大擴(kuò)展。我們還在維護(hù)著在蘋果GPU等平臺(tái)上的最新性能，并支持新興平臺(tái)，如WebGPU。

重要的是，我們發(fā)現(xiàn)在各平臺(tái)上構(gòu)建統(tǒng)一的解決方案存在很多協(xié)同作用。例如，擁有相同的引擎使我們能夠?qū)⒏呒?jí)優(yōu)化（如注意力管理、前綴緩存、推測(cè)執(zhí)行和結(jié)構(gòu)化生成）帶到這些平臺(tái)上，在過去的解決方案中由于有限的工程帶寬而缺乏支持。我們還可以通過利用代碼生成在平臺(tái)之間轉(zhuǎn)移優(yōu)化知識(shí)。我們?cè)赪ebGPU上的性能受益于本地GPU運(yùn)行時(shí)的優(yōu)化，然后直接編譯到目標(biāo)WebGPU環(huán)境。

總結(jié)

我們介紹了MLCEngine，一個(gè)統(tǒng)一的高效LLM引擎，可以在不同的硬件和平臺(tái)上運(yùn)行LLM，從云服務(wù)器到邊緣設(shè)備。這標(biāo)志著MLC LLM項(xiàng)目的新篇章。我們樂于與社區(qū)合作，帶來更多改進(jìn)，包括提供更好的模型覆蓋、更多的系統(tǒng)優(yōu)化以及先進(jìn)的機(jī)器學(xué)習(xí)編譯，以實(shí)現(xiàn)更加高效的通用部署。

本項(xiàng)目得以實(shí)現(xiàn)，多虧了開源社區(qū)的許多貢獻(xiàn)，包括CMU Catalyst、OctoAI、UW SAMPL、SJTU以及更廣泛的MLC社區(qū)的貢獻(xiàn)。我們計(jì)劃繼續(xù)與開源社區(qū)合作，為所有人帶來開放的基礎(chǔ)模型。

譯自（有刪改）：https://blog.mlc.ai/2024/06/07/universal-LLM-deployment-engine-with-ML-compilation

本文轉(zhuǎn)載自公眾號(hào)AIGC最前線   

原文鏈接：??https://mp.weixin.qq.com/s/GAVn6vm_rWUBOACnm8_EHA??