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.gitcd relaxmkdir buildcd buildcp ../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/relaxexport pythonpath=$tvm_home/python:${pythonpath}
然后source ~/.bashrc即可。
编译和安装mlc-llm
git clone --recursive git@github.com:mlc-ai/mlc-llm.gitcd mlc-llm/cmakepython3 gen_cmake_config.py
执行python3 gen_cmake_config.py 可以按需选择需要打开的编译选项，比如我这里就选择打开cuda，cublas，cutlass，另外需要注意的是这里的 tvm_home 路径需要设置为上面编译的relax路径。
然后执行下面的操作编译：
cd ..mkdir buildcp cmake/config.cmake buildcd buildcmake ..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 --model rwkv-4-world-7b-q0f16
rwkv-4-world-7b-q0f16可以换成你量化模型时的名字，加载完并运行system prompt之后你就可以愉快的和rwkv-4-world模型聊天了。
程序有一些特殊的指令来退出，查看速度等等：
性能测试
硬件 量化方法 速度
a800 q0f16 prefill: 362.7 tok/s, decode: 72.4 tok/s
a800 q4f16_1 prefill: 1104.7 tok/s, decode: 122.6 tok/s
这里给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_1weights 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 m2host cpu dection: target triple: arm64-apple-darwin22.3.0 process triple: arm64-apple-darwin22.3.0 host cpu: apple-m1target 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=32host cpu dection: target triple: arm64-apple-darwin22.3.0 process triple: arm64-apple-darwin22.3.0 host cpu: apple-m1automatically 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=32start computing and quantizing weights... this may take a while.finish computing and quantizing weights.total param size: 1.6060066223144531 gbstart storing to cache dist/rwkv-4-world-3b-q4f16_1/params[0808/0808] saving param_807all finished, 51 total shards committed, record saved to dist/rwkv-4-world-3b-q4f16_1/params/ndarray-cache.jsonfinish 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 boundsave 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的速度测试：
硬件 量化方法 速度
apple m2 q0f16 204.9 tok/s, decode: 12.1 tok/s
apple m2 q4f16_1 prefill: 201.6 tok/s, decode: 26.3 tok/s
建议使用q4f16的配置，这样回复会快一些。
0x3. 总结
这篇文章介绍了一下笔者最近给mlc-llm做适配的工作，欢迎大家体验mlc-llm和rwkv-world模型。
在中、美、德关键听证会最快速度促使苹果与高通达成和解
2019年机器人趋势预测 扶持减少 回归实用
比亚迪发布国内首款车规级共晶LED
2023年射频前端市场预测
微软发布全新卸载工具，Windows和macOS平台已发布
MLC-LLM的编译部署流程
零线故障现象你关注了吗？
细品AMD的3D缓存技术
Linux电源管理实例分析
国泰证券：风起粤港澳大湾区
网络互联,什么是网络互联
激光焊缝跟踪系统在冷凝器自动焊接的应用
就在本周六！RT-Thread线下入门培训深圳站！
什么是IOT？
负氧离子传感器的分类介绍
高通与苹果的关系是敌还是友
法拉第未来与美国混合动力公司合作开发新能源产品
停车场里的人工智能 五大服务迎来新变局
小米POCO M3海外发布:支持18W有线快充
夏普扩大提前退休制度至普通员工 最多发放12个月薪资补贴
我的 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.gitcd relaxmkdir buildcd buildcp ../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/relaxexport pythonpath=$tvm_home/python:${pythonpath}
然后source ~/.bashrc即可。
编译和安装mlc-llm
git clone --recursive git@github.com:mlc-ai/mlc-llm.gitcd mlc-llm/cmakepython3 gen_cmake_config.py
执行python3 gen_cmake_config.py 可以按需选择需要打开的编译选项，比如我这里就选择打开cuda，cublas，cutlass，另外需要注意的是这里的 tvm_home 路径需要设置为上面编译的relax路径。
然后执行下面的操作编译：
cd ..mkdir buildcp cmake/config.cmake buildcd buildcmake ..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 --model rwkv-4-world-7b-q0f16
rwkv-4-world-7b-q0f16可以换成你量化模型时的名字，加载完并运行system prompt之后你就可以愉快的和rwkv-4-world模型聊天了。
程序有一些特殊的指令来退出，查看速度等等：
性能测试
硬件 量化方法 速度
a800 q0f16 prefill: 362.7 tok/s, decode: 72.4 tok/s
a800 q4f16_1 prefill: 1104.7 tok/s, decode: 122.6 tok/s
这里给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_1weights 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 m2host cpu dection: target triple: arm64-apple-darwin22.3.0 process triple: arm64-apple-darwin22.3.0 host cpu: apple-m1target 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=32host cpu dection: target triple: arm64-apple-darwin22.3.0 process triple: arm64-apple-darwin22.3.0 host cpu: apple-m1automatically 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=32start computing and quantizing weights... this may take a while.finish computing and quantizing weights.total param size: 1.6060066223144531 gbstart storing to cache dist/rwkv-4-world-3b-q4f16_1/params[0808/0808] saving param_807all finished, 51 total shards committed, record saved to dist/rwkv-4-world-3b-q4f16_1/params/ndarray-cache.jsonfinish 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 boundsave 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的速度测试：
硬件 量化方法 速度
apple m2 q0f16 204.9 tok/s, decode: 12.1 tok/s
apple m2 q4f16_1 prefill: 201.6 tok/s, decode: 26.3 tok/s
建议使用q4f16的配置，这样回复会快一些。
0x3. 总结
这篇文章介绍了一下笔者最近给mlc-llm做适配的工作，欢迎大家体验mlc-llm和rwkv-world模型。
在中、美、德关键听证会最快速度促使苹果与高通达成和解
2019年机器人趋势预测 扶持减少 回归实用
比亚迪发布国内首款车规级共晶LED
2023年射频前端市场预测
微软发布全新卸载工具，Windows和macOS平台已发布
MLC-LLM的编译部署流程
零线故障现象你关注了吗？
细品AMD的3D缓存技术
Linux电源管理实例分析
国泰证券：风起粤港澳大湾区
网络互联,什么是网络互联
激光焊缝跟踪系统在冷凝器自动焊接的应用
就在本周六！RT-Thread线下入门培训深圳站！
什么是IOT？
负氧离子传感器的分类介绍
高通与苹果的关系是敌还是友
法拉第未来与美国混合动力公司合作开发新能源产品
停车场里的人工智能 五大服务迎来新变局
小米POCO M3海外发布:支持18W有线快充
夏普扩大提前退休制度至普通员工 最多发放12个月薪资补贴