MLC-LLM的編譯部署流程

0x0. 前言

我的 ChatRWKV 學習筆記和使用指南 這篇文章是學習RWKV的第一步，然后學習了一下之后決定自己應該做一些什么。所以就在RWKV社區看到了這個將RWKV World系列模型通過MLC-LLM部署在各種硬件平臺的需求，然后我就開始了解MLC-LLM的編譯部署流程和RWKV World模型相比于MLC-LLM已經支持的Raven系列模型的特殊之處。

MLC-LLM的編譯部署流程在MLC-LLM的官方文檔已經比較詳細了，但這部分有一些隱藏的坑點需要你去發現，比如現在要支持的RWKV-World模型它的Tokenizer是自定義的，并不是Huggingface的格式，這就導致我們不能使用MLC-LLM去直接編譯這個模型，也不能使用預編譯好的MLC-LLM二進制庫去運行這個模型了。另外，在編譯MLC-LLM倉庫之前我們需要先編譯Relax倉庫而不是原始的TVM倉庫，Relax可以認為是TVM的一個fork，在此基礎上支持了Relax這個新一代的IR，這部分背景建議讀者看一下我這個倉庫的相關鏈接：

https://github.com/BBuf/tvm_mlir_learn

這個倉庫已經攬下1.4k star，謝謝讀者們支持。

從RWKV社區了解到，RWKV-World系列模型相比于Raven系列，推理代碼和模型都是完全一樣，不一樣的地方主要是tokenizer是自定義的，并且system prompt不同。

在編譯Relax的時候需要按需選擇自己的編譯平臺進行編譯，編譯完之后 MLC-LLM 會通過 TVM_HOME 這個環境變量來感知 Relax 的位置，并且Relax編譯時開啟的選項要和MLC-LLM編譯的選項匹配上，這樣才可以在指定平臺上進行正確的編譯和推理。

在適配 RWKV-World 1.5B時，由于模型比較小對逗號比較敏感，導致第一層就炸了精度，最終掛在sampler里面，這個地方我定位2個晚上，后來mlc-ai官方的馮思遠告訴我在 MLC-LLM 里如何逐層打印精度之后，我最終定位到了問題。并且在 RWKV 社區里面了解到了這個現象之前就出現過，那就是1.5B的模型第一層需要用FP32來計算，不然會炸精度，我后續實驗了RWKV-4-World 3B/7B，這個現象就沒有了。

另外，模型的組織格式也是值得注意的一點，并不是在任意位置編譯好模型都可以在運行時被 MLC-LLM 正確發現。我大概花了快一周工作外時間在 MLC-LLM 上來支持 RWKV-World 系列模型，工作內容主要為：

將大缺弦的 https://github.com/daquexian/faster-rwkv 倉庫中的 RWKV World模型tokenizer實現掛到 mlc-ai 的 tokenizers.cpp 中，作為一個 3rd 庫提供給MLC-LLM。合并的PR為：https://github.com/mlc-ai/tokenizers-cpp/pull/14。

在上面的基礎上，在MLC-LLM中支持 RWKV World系列模型的部署，對齊 World 系列模型的 Prompt ，獲得良好的對話效果。分別在 Apple M2和A800顯卡上進行了部署和測試。PR為：https://github.com/mlc-ai/mlc-llm/pull/848 ，這個pr還wip，如果你現在要使用的話可以直接切到這個pr對應的分支就可以了。

debug到1.5B RWKV World小模型會炸精度的bug，相當于踩了個大坑。

我要特別感謝 mlc-ai 官方的馮思遠在我部署過程中提供的支持以及幫我Review讓代碼合并到 mlc-ai 社區，以及感謝大缺弦的 RWKV World Tokenizer c++實現以及在編譯第三方庫時幫我解決的一個bug。

以下是MLC-LLM 部署RWKV World系列模型教程，盡量提供大家部署最不踩坑的實踐。

效果：

在這里插入圖片描述

0x1. 將RWKV-4-World-7B部署在A800上

準備工作

RWKV-4-World模型地址：https://huggingface.co/StarRing2022/RWKV-4-World-7B

下載這里：https://github.com/BBuf/rwkv-world-tokenizer/releases/tag/v1.0.0 的 tokenizer_model.zip并解壓為tokenizer_model文件，這是RWKV World系列模型的Tokenizer文件。

克隆好 https://github.com/mlc-ai/mlc-llm 和 https://github.com/mlc-ai/relax ，注意克隆的時候一定要加上 --recursive 參數，這樣才會把它們依賴的第三方庫也添加上。

編譯Relax

git clone --recursive git@github.com:mlc-ai/relax.git
cd relax
mkdir build
cd build
cp ../cmake/config.cmake ./

然后修改build目錄下的config.cmake文件，由于我這里是在A800上面編譯，我改了以下設置：

set(USE_CUDA ON)
set(USE_CUTLASS ON)
set(USE_CUBLAS ON)

即啟用了CUDA，并開啟了2個加速庫CUTLASS和CUBLAS。然后在build目錄下執行cmake .. && make -j32 即可。

最后可以考慮把Relax添加到PYTHONPATH環境變量里面使得全局可見，在~/.bashrc上輸入以下內容：

export TVM_HOME=/bbuf/relax
export PYTHONPATH=$TVM_HOME/python:${PYTHONPATH}

然后source ~/.bashrc即可。

編譯和安裝MLC-LLM

git clone --recursive git@github.com:mlc-ai/mlc-llm.git
cd mlc-llm/cmake
python3 gen_cmake_config.py

執行python3 gen_cmake_config.py 可以按需選擇需要打開的編譯選項，比如我這里就選擇打開CUDA，CUBLAS，CUTLASS，另外需要注意的是這里的 TVM_HOME 路徑需要設置為上面編譯的Relax路徑。

然后執行下面的操作編譯：

cd ..
mkdir build
cp cmake/config.cmake build
cd build
cmake ..
make -j32

這里編譯時還需要安裝一下rust，按照建議的命令安裝即可，編譯完成之后即安裝上了mlc-llm提供的聊天程序mlc_chat_cli。然后為了做模型轉換和量化，我們還需要在mlc-llm目錄下執行一下pip install .安裝mlc_llm包。

模型轉換

模型轉換這里基本就是參考這個教程了：https://mlc.ai/mlc-llm/docs/compilation/compile_models.html 。

例如我們執行python3 -m mlc_llm.build --hf-path StarRing2022/RWKV-4-World-7B --target cuda --quantization q4f16_1 就可以將RWKV-4-World-7B模型權重量化為4個bit，然后activation還是以FP16的方式存儲。

target 則指定我們要在什么平臺上去運行，這里會將整個模型構成的圖編譯成一個動態鏈接庫（也就是TVM的IRModule）供后續的mlc_chat_cli程序（這個是在編譯mlc-llm時產生的）調用。

這里默認會在當前目錄下新建一個dist/models文件夾來存量化后模型和配置文件以及鏈接庫，轉換和量化好之后的模型會存儲在當前命令所在目錄的dist子目錄下（會自動創建），你也可以手動克隆huggingface模型到dist/models文件夾下。量化完之后的模型結構如下：

這里的mlc-chat-config.json指定來模型生成的一些超參數比如top_p，temperature等。

最后在推理之前，我們還需要把最開始準備的tokenizer_model文件拷貝到這個params文件夾中。

執行推理

我們在mlc-llm的上一層文件夾執行下面的命令：

./mlc-llm/build/mlc_chat_cli--modelRWKV-4-World-7B-q0f16

RWKV-4-World-7B-q0f16可以換成你量化模型時的名字，加載完并運行system prompt之后你就可以愉快的和RWKV-4-World模型聊天了。

程序有一些特殊的指令來退出，查看速度等等：

性能測試

這里給2組性能數據，大家感興趣的話可以測測其它配置。

逐層debug方法

在適配1.5B模型時出現了推理結果nan的現象，可以用mlc-llm/tests/debug/dump_intermediate.py這個文件來對齊輸入和tokenizer的結果之后進行debug，可以精準模擬模型推理并打印每一層的中間值，這樣我們就可以方便的看到模型是在哪一層出現了nan。

0x2. 將RWKV-4-World-3B部署在Apple M2上

在mac上部署和cuda上部署并沒有太大區別，主要是編譯relax和mlc-llm的時候編譯選項現在要選Metal而不是cuda了。我建議最好是在一個anconda環境里面處理編譯的問題，不要用系統自帶的python環境。

在編譯relax的時候需要同時打開使用Metal和LLVM選項，如果系統沒有LLVM可以先用Homebrew裝一下。

在mlc-llm中生成config.cmake時使用下面的選項：

編譯完并pip install .之后使用下面的命令量化模型：

python3 -m mlc_llm.build --hf-path StarRing2022/RWKV-4-World-3B --target metal --quantization q4f16_1

量化過程中日志如下：

(base) bbuf@MacBook-Pro RWKV % python3 -m mlc_llm.build --hf-path StarRing2022/RWKV-4-World-3B --target metal --quantization q4f16_1
Weights exist at dist/models/RWKV-4-World-3B, skipping download.
Using path "dist/models/RWKV-4-World-3B" for model "RWKV-4-World-3B"
[0908] /Users/bbuf/工作目錄/RWKV/relax/src/runtime/metal/metal_device_api.mm Intializing Metal device 0, name=Apple M2
Host CPU dection:
  Target triple: arm64-apple-darwin22.3.0
  Process triple: arm64-apple-darwin22.3.0
  Host CPU: apple-m1
Target configured: metal -keys=metal,gpu -max_function_args=31 -max_num_threads=256 -max_shared_memory_per_block=32768 -max_threads_per_block=1024 -thread_warp_size=32
Host CPU dection:
  Target triple: arm64-apple-darwin22.3.0
  Process triple: arm64-apple-darwin22.3.0
  Host CPU: apple-m1
Automatically using target for weight quantization: metal -keys=metal,gpu -max_function_args=31 -max_num_threads=256 -max_shared_memory_per_block=32768 -max_threads_per_block=1024 -thread_warp_size=32
Start computing and quantizing weights... This may take a while.
Finish computing and quantizing weights.
Total param size: 1.6060066223144531 GB
Start storing to cache dist/RWKV-4-World-3B-q4f16_1/params
[0808/0808] saving param_807
All finished, 51 total shards committed, record saved to dist/RWKV-4-World-3B-q4f16_1/params/ndarray-cache.json
Finish exporting chat config to dist/RWKV-4-World-3B-q4f16_1/params/mlc-chat-config.json
[0940] /Users/bbuf/工作目錄/RWKV/relax/include/tvm/topi/transform.h Warning: Fast mode segfaults when there are out-of-bounds indices. Make sure input indices are in bound
[0941] /Users/bbuf/工作目錄/RWKV/relax/include/tvm/topi/transform.h Warning: Fast mode segfaults when there are out-of-bounds indices. Make sure input indices are in bound
Save a cached module to dist/RWKV-4-World-3B-q4f16_1/mod_cache_before_build.pkl.
Finish exporting to dist/RWKV-4-World-3B-q4f16_1/RWKV-4-World-3B-q4f16_1-metal.so

同樣也需要把tokenizer_model文件拷貝到量化后模型文件夾的params目錄下，然后執行下面的命令啟動聊天程序：

./mlc-llm/build/mlc_chat_cli --model RWKV-4-World-3B-q0f16

最后也來一個Mac M2的速度測試：

建議使用q4f16的配置，這樣回復會快一些。

0x3. 總結

這篇文章介紹了一下筆者最近給mlc-llm做適配的工作，歡迎大家體驗MLC-LLM和RWKV-World模型。

審核編輯：彭菁