自從2023年3月左右，chatgpt火熱起來之后，我把關(guān)注的一些知乎帖子都記錄到了這個(gè)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推理等各個(gè)方面的知乎文章鏈接，感興趣的讀者可以看一下。

0x0. 前言

繼續(xù)填 使用MLC-LLM將RWKV 3B模型跑在Android手機(jī)上（redmi k50每s可解碼8個(gè)token 這篇文章留下的坑。由于上面這篇文章支持的是RWKV4模型，不支持最近RWKV社區(qū)正在訓(xùn)練的RWKV5模型，所以利用業(yè)余時(shí)間在MLC-LLM里面支持了最新的RWKV5模型的推理，同時(shí)也可以帶大家看一下RWKV5的3B模型表現(xiàn)是否有驚艷之處。目前我跑通了Metal和Android平臺的RWKV5推理（包含1.5B和3B），并且也編譯出了一個(gè)3B int8模式的apk提供給android用戶使用，地址為：https://github.com/BBuf/run-rwkv-world-4-in-mlc-llm/releases/download/v1.0.0/rwkv5-3b-int8.apk 。大家可以下載這個(gè)apk來體驗(yàn)最新的RWKV-5-3B模型。

另外，我在測試RWKV-5-3B的時(shí)候也發(fā)現(xiàn)了RWKV4的表現(xiàn)和HuggingFace版本的表現(xiàn)相差比較多，也修復(fù)了這個(gè)bug。總的來說，在MLC-LLM里面適配一個(gè)新的RWKV5模型是比較麻煩的，我前后肝了幾個(gè)周末，并且在Hzfengsy的熱心幫助下解決了一個(gè)關(guān)鍵的TIR實(shí)現(xiàn)問題后。這篇文章我會分享一下適配過程中的主要問題是什么，給想使用MLC-LLM適配其它不支持的模型的讀者一個(gè)踩坑經(jīng)驗(yàn)。

關(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個(gè)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上因?yàn)橐粋€(gè)已知的tvm bug導(dǎo)致編譯失敗，如果要在Nvidia GPU上部署RWKV-5-World模型需要等官方完成這個(gè)bug fix，具體請關(guān)注 https://github.com/mlc-ai/mlc-llm/pull/1275 進(jìn)展。

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

對LLM的理解比較有限，從代碼實(shí)現(xiàn)的角度來說，RWKV的狀態(tài)和KV Cache不同，不依賴序列長度，這讓RWKV模型在各種長度下運(yùn)行內(nèi)存和運(yùn)行速度都是趨于穩(wěn)定的，所以我感覺工程價(jià)值是比基于Transformer架構(gòu)比如Llama更好的，部署的性價(jià)比會天然更優(yōu)。這個(gè)特點(diǎn)讓他在更長的序列比如100K長度下的推理也更有前景吧。但是，RWKV是否可以取得和Transformer主流架構(gòu)相同的效果呢？我個(gè)人感覺還是需要等待時(shí)間的檢驗(yàn)，目前最新的RWKV5模型最多scale up到7B，并且數(shù)據(jù)也是很有限只有1.12TB，這個(gè)信息我是從HuggingFace的項(xiàng)目看到的，如下圖所示。（這里的v2就是最新的RWKV5架構(gòu)，內(nèi)部小版本命名稍顯混亂，這一點(diǎn)也可以從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等等。這個(gè)屬于開源大模型評測的知識，大家應(yīng)該能找到很多榜單，RWKV官方是否考慮去opencompass打一下榜，更全面的做個(gè)對比。

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

在這里插入圖片描述

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

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

概念介紹

在這里插入圖片描述

廣告文案

在這里插入圖片描述

作文生成

在這里插入圖片描述

演講稿生成

在這里插入圖片描述

心得體會

在這里插入圖片描述

科技文稿

記錄文

在這里插入圖片描述

評論評語

在這里插入圖片描述

問題生成

在這里插入圖片描述

起名字

在這里插入圖片描述

簡單代碼

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

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

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

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

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

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

在這里插入圖片描述

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

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

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

在這里插入圖片描述

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

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

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

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

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

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

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

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

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

0x4. 總結(jié)

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

編輯：黃飛