使用MLC-LLM支持RWKV-5推理的过程思考

自从2023年3月左右，chatgpt火热起来之后，我把关注的一些知乎帖子都记录到了这个markdown里面，：https://github.com/BBuf/how-to-optim-algorithm-in-cuda/tree/master/large-language-model-note ，从2023年3月左右到现在保持了持续动态更新整理，有关于LLM基础知识，LLM训练，LLM推理等各个方面的知乎文章链接，感兴趣的读者可以看一下。

0x0. 前言

继续填 使用MLC-LLM将RWKV 3B模型跑在Android手机上（redmi k50每s可解码8个token 这篇文章留下的坑。由于上面这篇文章支持的是RWKV4模型，不支持最近RWKV社区正在训练的RWKV5模型，所以利用业余时间在MLC-LLM里面支持了最新的RWKV5模型的推理，同时也可以带大家看一下RWKV5的3B模型表现是否有惊艳之处。目前我跑通了Metal和Android平台的RWKV5推理（包含1.5B和3B），并且也编译出了一个3B int8模式的apk提供给android用户使用，地址为：https://github.com/BBuf/run-rwkv-world-4-in-mlc-llm/releases/download/v1.0.0/rwkv5-3b-int8.apk 。大家可以下载这个apk来体验最新的RWKV-5-3B模型。

另外，我在测试RWKV-5-3B的时候也发现了RWKV4的表现和HuggingFace版本的表现相差比较多，也修复了这个bug。总的来说，在MLC-LLM里面适配一个新的RWKV5模型是比较麻烦的，我前后肝了几个周末，并且在Hzfengsy的热心帮助下解决了一个关键的TIR实现问题后。这篇文章我会分享一下适配过程中的主要问题是什么，给想使用MLC-LLM适配其它不支持的模型的读者一个踩坑经验。

关于RWKV模型的更多信息大家可以关注bo的两篇博客：

RWKV-5 的训练进展，与 SOTA GPT 模型的性能对比：https://zhuanlan.zhihu.com/p/659872347

RWKV-5 的训练进展（之二），与 SotA GPT 模型的性能对比：https://zhuanlan.zhihu.com/p/664079347

再次感谢@Hzfengsy 在适配RWKV-5过程中的指导。

本文涉及到的工程代码体现在下面的2个PR：

https://github.com/mlc-ai/mlc-llm/pull/1275 （MLC-LLM中支持RWKV5）

https://github.com/mlc-ai/tokenizers-cpp/pull/19 （对RWKV World Tokenzier的bug修复，也提升了RWKV-4-World系列模型的效果）

另外，目前MLC-LLM支持RWKV-5在Metal和Android的推理，但是在nvidia gpu上因为一个已知的tvm bug导致编译失败，如果要在Nvidia GPU上部署RWKV-5-World模型需要等官方完成这个bug fix，具体请关注 https://github.com/mlc-ai/mlc-llm/pull/1275 进展。

0x1. 笔者为何关注RWKV

对LLM的理解比较有限，从代码实现的角度来说，RWKV的状态和KV Cache不同，不依赖序列长度，这让RWKV模型在各种长度下运行内存和运行速度都是趋于稳定的，所以我感觉工程价值是比基于Transformer架构比如Llama更好的，部署的性价比会天然更优。这个特点让他在更长的序列比如100K长度下的推理也更有前景吧。但是，RWKV是否可以取得和Transformer主流架构相同的效果呢？我个人感觉还是需要等待时间的检验，目前最新的RWKV5模型最多scale up到7B，并且数据也是很有限只有1.12TB，这个信息我是从HuggingFace的项目看到的，如下图所示。（这里的v2就是最新的RWKV5架构，内部小版本命名稍显混乱，这一点也可以从ChatRWKV的model.py看出）。

所以如果RWKV架构真的可以取得和Transformer开源SOTA架构一样的效果，前景是很好的。RWKV-5 的训练进展（之二），与 SotA GPT 模型的性能对比：https://zhuanlan.zhihu.com/p/664079347 这里已经贴出一些BenchMark结果：

从作者这里选取的一些数据集来看，RWKV-5-World 7B目前仅训练30%的checkpoint的效果已经和Baichuan2-7B-Base非常接近了，还是值得期待一下的。

不过，这里存在的问题是这里的这些测试的数据集可能需要使用一些更加有说服力的，比如MMLU/CMMLU/HummanEval/MBPP/CMRC2018等等。这个属于开源大模型评测的知识，大家应该能找到很多榜单，RWKV官方是否考虑去opencompass打一下榜，更全面的做个对比。

因为这里有个明显的疑问就是，按照官方的说法，为什么使用1.12T数据训练30%之后在上面的任务里面就可以几乎持平使用2.6T数据进行全量预训练的Baichuan2-7B-Base模型的效果呢？所以我个人感觉这里需要更多的榜单数据来看效果。

在这里插入图片描述

0x2. RWKV-5-3B模型在Mac上的一些文创和代码生成效果演示

我个人感觉7B模型和3B模型就是为了手机上离线运行而生的尺寸，所以我这里使用上面编译的Apk来演示一下使用MLC-LLM推理的RWKV-5-3B模型的一些文创效果和代码生成效果。下面演示的文创问题大多数来自昆仑天工的Skywork-13B例子（https://github.com/SkyworkAI/Skywork），感谢。下面的User是我问的问题，Assistant是RWKV-5-3B模型的回答，运行环境为Mac M2 FP16模式。由于这个模型是基础模型，所以对话效果会受到上下文多轮对话干扰，所以在测试不同种类的问题时，可以使用/reset来重置对话。

概念介绍

在这里插入图片描述

广告文案

在这里插入图片描述

作文生成

在这里插入图片描述

演讲稿生成

在这里插入图片描述

心得体会

在这里插入图片描述

科技文稿

记录文

在这里插入图片描述

评论评语

在这里插入图片描述

问题生成

在这里插入图片描述

起名字

在这里插入图片描述

简单代码

总的来说，对于大多数文学创作问题，RWKV-5-3B的回答还算像那回事，不过也可以明显感觉到一些瑕疵以及指令跟随的能力很有限，比如对数字非常不敏感，让他说5个字他似乎不明白意思。此外，3b模型拥有了一定的代码能力，可以写有限的简单代码。

最后，我比较期待7b最终训练完之后的效果，希望RWKV可以在opencompass榜单上证明自己。

0x3. MLC-LLM支持RWKV-5步骤

这一节可能会写得流水账一点。模型实现文件：https://github.com/mlc-ai/mlc-llm/pull/1275 里的 rwkv5.py

首先，由于MLC-LLM已经支持了RWKV4架构，所以我们大体上是可以使用RWKV4的实现的，然后把RWKV5的改动加上去。

我们可以从ChatRWKV的rwkv4/rwkv5模型实现（https://github.com/BlinkDL/ChatRWKV/blob/main/rwkv_pip_package/src/rwkv/model.py）看出rwkv4和rwkv5的不同之处主要在于RWKV5引入了多头的线性Attention，代码上体现为对Attention部分的重写，包括state的个数也从5个变成了3个。从MLC-LLM的模型实现代码上来看，如果要在同一个实现中进行兼容会相当麻烦，所以我使用了一个新的文件来实现RWKV5，接下来就是对着ChatRWKV修改代码把RWKV5的初版本改上去。在RWKV5的prefill阶段，会调用一个新的CUDA Kernel：https://github.com/BlinkDL/ChatRWKV/blob/main/rwkv_pip_package/src/rwkv/model.py#L465-L497 。而这个Kernel的原始实现则对应这里的Python公式：https://github.com/BlinkDL/RWKV-CUDA/blob/main/wkv5/run.py#L67-L87

在这里插入图片描述

但需要注意的是，在真正的模型实现中，这里的state是需要更新的全局变量而非local的。由于这个函数有一个循环会在T的维度上进行迭代，而T是序列长度是可变的，所以这里需要类似于RWKV4的实现写一个TIR来模拟这个python程序的逻辑，在冯博的帮助下得到了一版初始的TIR实现：

这个实现过程中也帮助发现一个DLight的bug，由@Hzfengsy在tvm里面进行了修复。https://github.com/apache/tvm/pull/16124

解决了上面的TIR问题之后就可以在MLC-LLM里面编译RWKV5模型了，然后使用TVM的dump ir工具和ChatRWKV来对比精度，这里需要固定输入的Tensor才行，为了方便我将输入固定为一个全1的十个元素的ids。然后在对比精度的实现发现，上面实现的TIR的输入的所有值都是可以对上的，但是TIR的输出out却是错误的。仍旧是冯博帮我解决了这个bug，原因是因为上面的版本中对于state来说T不应该是spatial的而是reduction。修复后的正确版本长这样：

在这里插入图片描述

接着又从dump的结果观察到attention部分的groupnorm的结果无法对上，但输入都是可以对上的，然后我手动实现了一下groupnorm的过程（下面的237-247行）发现结果竟然是可以对上的。

后面经Hzfengsy提醒确认是开始的groupnorm调用参数写错了，修复之后继续下一步。这一下attention和ffn的结果是可以对上了。

然后开始使用mlc chat程序尝试进行对话，发现输出会乱码。又怀疑中间某个地方精度没对齐，所以继续完整模拟了一遍prefill+decode，发现prefill+第一轮decode的结果完全能对上，想摆烂了。。

然后我使用相同的问题问了一下ChatRWKV，发现ChatRWKV的结果也是乱码。。。直觉告诉我一定是乌龙了，由于我这里对比的ChatRWKV是我自己fork的，可能不小心改了bug。我重新拉官方的ChatRWKV一一对比，找到了问题所在。是因为我的代码里错误的去掉一个transpose op，我也忘记了为什么要这么做，但是这个transpose op去transpose的两个维度的大小是相同的，所以输出shape也是相同的，导致了对精度浪费了很多时间。

解决这个问题之后，发现输出就是正常的了。但，真的正常吗？

我在尝试一些问题时发现输出非常奇怪：

感觉这里一定还有bug，既然模型精度方面没有bug，要么就是prompt技巧，tokenizer，sampling。sampling是比较正常并且经过众多模型检验的，应该问题不大。然后恰好想起daquexian的faster-rwkv里面更新过tokenzier，之前的实现应该有bug：

接下来就是更新tokenzier的代码修复bug，最后在review 初始化prompt的时候也发现了一个bug，将其修复。

最终获得的代码效果就是0x2节展示的了，这些prompt的输出和ChatRWKV相差不大，理论上来说应该是完成了正确的适配。

0x4. 总结

本文记录了笔者使用 MLC-LLM 支持RWKV-5推理的过程以及对RWKV-5的一些思考，谢谢。

编辑：黄飞