自從2023年3月左右，chatgpt火熱起來之后，我把關(guān)注的一些知乎帖子都記錄到了這個markdown里面，：https://github.com/BBuf/how-to-optim-algorithm-in-cuda/tree/master/large-language-model-note ，從2023年3月左右到現(xiàn)在保持了持續(xù)動態(tài)更新整理，有關(guān)于LLM基礎(chǔ)知識，LLM訓(xùn)練，LLM推理等各個方面的知乎文章鏈接，感興趣的讀者可以看一下。

0x0. 前言

繼續(xù)填 使用MLC-LLM將RWKV 3B模型跑在Android手機上（redmi k50每s可解碼8個token 這篇文章留下的坑。由于上面這篇文章支持的是RWKV4模型，不支持最近RWKV社區(qū)正在訓(xùn)練的RWKV5模型，所以利用業(yè)余時間在MLC-LLM里面支持了最新的RWKV5模型的推理，同時也可以帶大家看一下RWKV5的3B模型表現(xiàn)是否有驚艷之處。目前我跑通了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的時候也發(fā)現(xiàn)了RWKV4的表現(xiàn)和HuggingFace版本的表現(xiàn)相差比較多，也修復(fù)了這個bug。總的來說，在MLC-LLM里面適配一個新的RWKV5模型是比較麻煩的，我前后肝了幾個周末，并且在Hzfengsy的熱心幫助下解決了一個關(guān)鍵的TIR實現(xiàn)問題后。這篇文章我會分享一下適配過程中的主要問題是什么，給想使用MLC-LLM適配其它不支持的模型的讀者一個踩坑經(jīng)驗。

關(guān)于RWKV模型的更多信息大家可以關(guān)注bo的兩篇博客：

RWKV-5 的訓(xùn)練進(jìn)展，與 SOTA GPT 模型的性能對比：https://zhuanlan.zhihu.com/p/659872347

RWKV-5 的訓(xùn)練進(jìn)展（之二），與 SotA GPT 模型的性能對比：https://zhuanlan.zhihu.com/p/664079347

再次感謝@Hzfengsy 在適配RWKV-5過程中的指導(dǎo)。

本文涉及到的工程代碼體現(xiàn)在下面的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修復(fù)，也提升了RWKV-4-World系列模型的效果）

另外，目前MLC-LLM支持RWKV-5在Metal和Android的推理，但是在nvidia gpu上因為一個已知的tvm bug導(dǎo)致編譯失敗，如果要在Nvidia GPU上部署RWKV-5-World模型需要等官方完成這個bug fix，具體請關(guān)注 https://github.com/mlc-ai/mlc-llm/pull/1275 進(jìn)展。

0x1. 筆者為何關(guān)注RWKV

對LLM的理解比較有限，從代碼實現(xiàn)的角度來說，RWKV的狀態(tài)和KV Cache不同，不依賴序列長度，這讓RWKV模型在各種長度下運行內(nèi)存和運行速度都是趨于穩(wěn)定的，所以我感覺工程價值是比基于Transformer架構(gòu)比如Llama更好的，部署的性價比會天然更優(yōu)。這個特點讓他在更長的序列比如100K長度下的推理也更有前景吧。但是，RWKV是否可以取得和Transformer主流架構(gòu)相同的效果呢？我個人感覺還是需要等待時間的檢驗，目前最新的RWKV5模型最多scale up到7B，并且數(shù)據(jù)也是很有限只有1.12TB，這個信息我是從HuggingFace的項目看到的，如下圖所示。（這里的v2就是最新的RWKV5架構(gòu)，內(nèi)部小版本命名稍顯混亂，這一點也可以從ChatRWKV的model.py看出）。

所以如果RWKV架構(gòu)真的可以取得和Transformer開源SOTA架構(gòu)一樣的效果，前景是很好的。RWKV-5 的訓(xùn)練進(jìn)展（之二），與 SotA GPT 模型的性能對比：https://zhuanlan.zhihu.com/p/664079347 這里已經(jīng)貼出一些BenchMark結(jié)果：

從作者這里選取的一些數(shù)據(jù)集來看，RWKV-5-World 7B目前僅訓(xùn)練30%的checkpoint的效果已經(jīng)和Baichuan2-7B-Base非常接近了，還是值得期待一下的。

不過，這里存在的問題是這里的這些測試的數(shù)據(jù)集可能需要使用一些更加有說服力的，比如MMLU/CMMLU/HummanEval/MBPP/CMRC2018等等。這個屬于開源大模型評測的知識，大家應(yīng)該能找到很多榜單，RWKV官方是否考慮去opencompass打一下榜，更全面的做個對比。

因為這里有個明顯的疑問就是，按照官方的說法，為什么使用1.12T數(shù)據(jù)訓(xùn)練30%之后在上面的任務(wù)里面就可以幾乎持平使用2.6T數(shù)據(jù)進(jìn)行全量預(yù)訓(xùn)練的Baichuan2-7B-Base模型的效果呢？所以我個人感覺這里需要更多的榜單數(shù)據(jù)來看效果。

在這里插入圖片描述

0x2. RWKV-5-3B模型在Mac上的一些文創(chuàng)和代碼生成效果演示

我個人感覺7B模型和3B模型就是為了手機上離線運行而生的尺寸，所以我這里使用上面編譯的Apk來演示一下使用MLC-LLM推理的RWKV-5-3B模型的一些文創(chuàng)效果和代碼生成效果。下面演示的文創(chuàng)問題大多數(shù)來自昆侖天工的Skywork-13B例子（https://github.com/SkyworkAI/Skywork），感謝。下面的User是我問的問題，Assistant是RWKV-5-3B模型的回答，運行環(huán)境為Mac M2 FP16模式。由于這個模型是基礎(chǔ)模型，所以對話效果會受到上下文多輪對話干擾，所以在測試不同種類的問題時，可以使用/reset來重置對話。

概念介紹

在這里插入圖片描述

廣告文案

在這里插入圖片描述

作文生成

在這里插入圖片描述

演講稿生成

在這里插入圖片描述

心得體會

在這里插入圖片描述

科技文稿

記錄文

在這里插入圖片描述

評論評語

在這里插入圖片描述

問題生成

在這里插入圖片描述

起名字

在這里插入圖片描述

簡單代碼

總的來說，對于大多數(shù)文學(xué)創(chuàng)作問題，RWKV-5-3B的回答還算像那回事，不過也可以明顯感覺到一些瑕疵以及指令跟隨的能力很有限，比如對數(shù)字非常不敏感，讓他說5個字他似乎不明白意思。此外，3b模型擁有了一定的代碼能力，可以寫有限的簡單代碼。

最后，我比較期待7b最終訓(xùn)練完之后的效果，希望RWKV可以在opencompass榜單上證明自己。

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

這一節(jié)可能會寫得流水賬一點。模型實現(xiàn)文件：https://github.com/mlc-ai/mlc-llm/pull/1275 里的 rwkv5.py

首先，由于MLC-LLM已經(jīng)支持了RWKV4架構(gòu)，所以我們大體上是可以使用RWKV4的實現(xiàn)的，然后把RWKV5的改動加上去。

我們可以從ChatRWKV的rwkv4/rwkv5模型實現(xiàn)（https://github.com/BlinkDL/ChatRWKV/blob/main/rwkv_pip_package/src/rwkv/model.py）看出rwkv4和rwkv5的不同之處主要在于RWKV5引入了多頭的線性Attention，代碼上體現(xiàn)為對Attention部分的重寫，包括state的個數(shù)也從5個變成了3個。從MLC-LLM的模型實現(xiàn)代碼上來看，如果要在同一個實現(xiàn)中進(jìn)行兼容會相當(dāng)麻煩，所以我使用了一個新的文件來實現(xiàn)RWKV5，接下來就是對著ChatRWKV修改代碼把RWKV5的初版本改上去。在RWKV5的prefill階段，會調(diào)用一個新的CUDA Kernel：https://github.com/BlinkDL/ChatRWKV/blob/main/rwkv_pip_package/src/rwkv/model.py#L465-L497 。而這個Kernel的原始實現(xiàn)則對應(yīng)這里的Python公式：https://github.com/BlinkDL/RWKV-CUDA/blob/main/wkv5/run.py#L67-L87

在這里插入圖片描述

但需要注意的是，在真正的模型實現(xiàn)中，這里的state是需要更新的全局變量而非local的。由于這個函數(shù)有一個循環(huán)會在T的維度上進(jìn)行迭代，而T是序列長度是可變的，所以這里需要類似于RWKV4的實現(xiàn)寫一個TIR來模擬這個python程序的邏輯，在馮博的幫助下得到了一版初始的TIR實現(xiàn)：

這個實現(xiàn)過程中也幫助發(fā)現(xiàn)一個DLight的bug，由@Hzfengsy在tvm里面進(jìn)行了修復(fù)。https://github.com/apache/tvm/pull/16124

解決了上面的TIR問題之后就可以在MLC-LLM里面編譯RWKV5模型了，然后使用TVM的dump ir工具和ChatRWKV來對比精度，這里需要固定輸入的Tensor才行，為了方便我將輸入固定為一個全1的十個元素的ids。然后在對比精度的實現(xiàn)發(fā)現(xiàn)，上面實現(xiàn)的TIR的輸入的所有值都是可以對上的，但是TIR的輸出out卻是錯誤的。仍舊是馮博幫我解決了這個bug，原因是因為上面的版本中對于state來說T不應(yīng)該是spatial的而是reduction。修復(fù)后的正確版本長這樣：

在這里插入圖片描述

接著又從dump的結(jié)果觀察到attention部分的groupnorm的結(jié)果無法對上，但輸入都是可以對上的，然后我手動實現(xiàn)了一下groupnorm的過程（下面的237-247行）發(fā)現(xiàn)結(jié)果竟然是可以對上的。

后面經(jīng)Hzfengsy提醒確認(rèn)是開始的groupnorm調(diào)用參數(shù)寫錯了，修復(fù)之后繼續(xù)下一步。這一下attention和ffn的結(jié)果是可以對上了。

然后開始使用mlc chat程序嘗試進(jìn)行對話，發(fā)現(xiàn)輸出會亂碼。又懷疑中間某個地方精度沒對齊，所以繼續(xù)完整模擬了一遍prefill+decode，發(fā)現(xiàn)prefill+第一輪decode的結(jié)果完全能對上，想擺爛了。。

然后我使用相同的問題問了一下ChatRWKV，發(fā)現(xiàn)ChatRWKV的結(jié)果也是亂碼。。。直覺告訴我一定是烏龍了，由于我這里對比的ChatRWKV是我自己fork的，可能不小心改了bug。我重新拉官方的ChatRWKV一一對比，找到了問題所在。是因為我的代碼里錯誤的去掉一個transpose op，我也忘記了為什么要這么做，但是這個transpose op去transpose的兩個維度的大小是相同的，所以輸出shape也是相同的，導(dǎo)致了對精度浪費了很多時間。

解決這個問題之后，發(fā)現(xiàn)輸出就是正常的了。但，真的正常嗎？

我在嘗試一些問題時發(fā)現(xiàn)輸出非常奇怪：

感覺這里一定還有bug，既然模型精度方面沒有bug，要么就是prompt技巧，tokenizer，sampling。sampling是比較正常并且經(jīng)過眾多模型檢驗的，應(yīng)該問題不大。然后恰好想起daquexian的faster-rwkv里面更新過tokenzier，之前的實現(xiàn)應(yīng)該有bug：

接下來就是更新tokenzier的代碼修復(fù)bug，最后在review 初始化prompt的時候也發(fā)現(xiàn)了一個bug，將其修復(fù)。

最終獲得的代碼效果就是0x2節(jié)展示的了，這些prompt的輸出和ChatRWKV相差不大，理論上來說應(yīng)該是完成了正確的適配。

0x4. 總結(jié)

本文記錄了筆者使用 MLC-LLM 支持RWKV-5推理的過程以及對RWKV-5的一些思考，謝謝。

編輯：黃飛