自從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的一些思考，謝謝。

編輯：黃飛