CPU后端和CUDA后端的執行代碼和效果
0x0. 前言
RWKV社區在Huggingface上放了rwkv-4-world和rwkv-5-world相關的一系列模型，見：https://huggingface.co/BlinkDL/rwkv-4-world & https://huggingface.co/BlinkDL/rwkv-5-world ，然而這些模型的格式是以PyTorch的格式進行保存的即*.pt文件，并沒有將其轉換為標準的Huggingface模型。后來了解到這里還有一個問題是RWKV的世界模型系列的tokenizer是自定義的，在Huggingface里面并沒有與之對應的Tokenizer。沒有標準的Huggingface模型就沒辦法使用TGI進行部署，也不利于模型的傳播以及和其它模型一起做評測等等。
讓RWKV world系列模型登陸Huggingface社區是必要的，這篇文章介紹了筆者為了達成這個目標所做的一些努力，最后成功讓rwkv-4-world和rwkv-5-world系列模型登陸Huggingface。大家可以在 https://huggingface.co/RWKV 這個空間找到目前所有登陸的RWKV模型：
在這里插入圖片描述本系列的工作都整理開源在 https://github.com/BBuf/RWKV-World-HF-Tokenizer ，包含將 RWKV world tokenizer 實現為 Huggingface 版本，實現 RWKV 5.0 的模型，提供模型轉換腳本，Lambda數據集ppl正確性檢查工具 等等。
0x1. 效果
以 RWKV/rwkv-4-world-3b 為例，下面分別展示一下CPU后端和CUDA后端的執行代碼和效果。
CPU
fromtransformersimportAutoModelForCausalLM,AutoTokenizer
model=AutoModelForCausalLM.from_pretrained("RWKV/rwkv-4-world-3b")
tokenizer=AutoTokenizer.from_pretrained("RWKV/rwkv-4-world-3b",trust_remote_code=True)
text="
Inashockingfinding,scientistdiscoveredaherdofdragonslivinginaremote,previouslyunexploredvalley,inTibet.EvenmoresurprisingtotheresearcherswasthefactthatthedragonsspokeperfectChinese."
prompt=f'Question:{text.strip()}
Answer:'
inputs=tokenizer(prompt,return_tensors="pt")
output=model.generate(inputs["input_ids"],max_new_tokens=256)
print(tokenizer.decode(output[0].tolist(),skip_special_tokens=True))
輸出：
Question: In a shocking finding, scientist discovered a herd of dragons living in a remote, previously unexplored valley, in Tibet. Even more surprising to the researchers was the fact that the dragons spoke perfect Chinese.
Answer: The dragons in the valley spoke perfect Chinese, according to the scientist who discovered them.
GPU
importtorch
fromtransformersimportAutoModelForCausalLM,AutoTokenizer
model=AutoModelForCausalLM.from_pretrained("RWKV/rwkv-4-world-3b",torch_dtype=torch.float16).to(0)
tokenizer=AutoTokenizer.from_pretrained("RWKV/rwkv-4-world-3b",trust_remote_code=True)
text="你叫什么名字？"
prompt=f'Question:{text.strip()}
Answer:'
inputs=tokenizer(prompt,return_tensors="pt").to(0)
output=model.generate(inputs["input_ids"],max_new_tokens=40)
print(tokenizer.decode(output[0].tolist(),skip_special_tokens=True))
輸出：
Question: 你叫什么名字？
Answer: 我是一個人工智能語言模型，沒有名字。
我們可以在本地通過上述代碼分別運行CPU/GPU上的wkv-4-world-3b模型，當然這需要安裝transformers和torch庫。
0x2. 教程
下面展示一下在 https://github.com/BBuf/RWKV-World-HF-Tokenizer 做的自定義實現的RWKV world tokenizer的測試，RWKV world模型轉換，檢查lambda數據集正確性等的教程。
使用此倉庫（https://github.com/BBuf/RWKV-World-HF-Tokenizer）的Huggingface項目
上傳轉換后的模型到Huggingface上時，如果bin文件太大需要使用這個指令
transformers-cli lfs-enable-largefiles解除大小限制.
- RWKV/rwkv-5-world-169m
- RWKV/rwkv-4-world-169m
- RWKV/rwkv-4-world-430m
- RWKV/rwkv-4-world-1b5
- RWKV/rwkv-4-world-3b
- RWKV/rwkv-4-world-7b
RWKV World模型的HuggingFace版本的Tokenizer
下面的參考程序比較了原始tokenizer和HuggingFace版本的tokenizer對不同句子的編碼和解碼結果。
fromtransformersimportAutoModelForCausalLM,AutoTokenizer
fromrwkv_tokenizerimportTRIE_TOKENIZER
token_path="/Users/bbuf/工作目錄/RWKV/RWKV-World-HF-Tokenizer/rwkv_world_tokenizer"
origin_tokenizer=TRIE_TOKENIZER('/Users/bbuf/工作目錄/RWKV/RWKV-World-HF-Tokenizer/rwkv_vocab_v20230424.txt')
fromtransformersimportAutoTokenizer
hf_tokenizer=AutoTokenizer.from_pretrained(token_path,trust_remote_code=True)
#測試編碼器
asserthf_tokenizer("Hello")['input_ids']==origin_tokenizer.encode('Hello')
asserthf_tokenizer("S:2")['input_ids']==origin_tokenizer.encode('S:2')
asserthf_tokenizer("MadeinChina")['input_ids']==origin_tokenizer.encode('MadeinChina')
asserthf_tokenizer("今天天氣不錯")['input_ids']==origin_tokenizer.encode('今天天氣不錯')
asserthf_tokenizer("男：聽說你們公司要派你去南方工作?")['input_ids']==origin_tokenizer.encode('男：聽說你們公司要派你去南方工作?')
#測試解碼器
asserthf_tokenizer.decode(hf_tokenizer("Hello")['input_ids'])=='Hello'
asserthf_tokenizer.decode(hf_tokenizer("S:2")['input_ids'])=='S:2'
asserthf_tokenizer.decode(hf_tokenizer("MadeinChina")['input_ids'])=='MadeinChina'
asserthf_tokenizer.decode(hf_tokenizer("今天天氣不錯")['input_ids'])=='今天天氣不錯'
asserthf_tokenizer.decode(hf_tokenizer("男：聽說你們公司要派你去南方工作?")['input_ids'])=='男：聽說你們公司要派你去南方工作?'
Huggingface RWKV World模型轉換
使用腳本scripts/convert_rwkv_world_model_to_hf.sh，將huggingface的BlinkDL/rwkv-4-world項目中的PyTorch格式模型轉換為Huggingface格式。這里，我們以0.1B為例。
#!/bin/bash
set-x
cdscripts
pythonconvert_rwkv_checkpoint_to_hf.py--repo_idBlinkDL/rwkv-4-world
--checkpoint_fileRWKV-4-World-0.1B-v1-20230520-ctx4096.pth
--output_dir../rwkv4-world4-0.1b-model/
--tokenizer_file/Users/bbuf/工作目錄/RWKV/RWKV-World-HF-Tokenizer/rwkv_world_tokenizer
--size169M
--is_world_tokenizerTrue
cp/Users/bbuf/工作目錄/RWKV/RWKV-World-HF-Tokenizer/rwkv_world_tokenizer/rwkv_vocab_v20230424.json../rwkv4-world4-0.1b-model/
cp/Users/bbuf/工作目錄/RWKV/RWKV-World-HF-Tokenizer/rwkv_world_tokenizer/tokenization_rwkv_world.py../rwkv4-world4-0.1b-model/
cp/Users/bbuf/工作目錄/RWKV/RWKV-World-HF-Tokenizer/rwkv_world_tokenizer/tokenizer_config.json../rwkv4-world4-0.1b-model/
使用腳本 scripts/convert_rwkv5_world_model_to_hf.sh，將來自 huggingface BlinkDL/rwkv-5-world 項目的 PyTorch 格式模型轉換為 Huggingface 格式。在這里，我們以 0.1B 為例。
#!/bin/bash
set-x
cdscripts
pythonconvert_rwkv5_checkpoint_to_hf.py--repo_idBlinkDL/rwkv-5-world
--checkpoint_fileRWKV-5-World-0.1B-v1-20230803-ctx4096.pth
--output_dir../rwkv5-world-169m-model/
--tokenizer_file/Users/bbuf/工作目錄/RWKV/RWKV-World-HF-Tokenizer/rwkv_world_tokenizer
--size169M
--is_world_tokenizerTrue
cp/Users/bbuf/工作目錄/RWKV/RWKV-World-HF-Tokenizer/rwkv_world_v5.0_model/configuration_rwkv5.py../rwkv5-world-169m-model/
cp/Users/bbuf/工作目錄/RWKV/RWKV-World-HF-Tokenizer/rwkv_world_v5.0_model/modeling_rwkv5.py../rwkv5-world-169m-model/
cp/Users/bbuf/工作目錄/RWKV/RWKV-World-HF-Tokenizer/rwkv_world_tokenizer/rwkv_vocab_v20230424.json../rwkv5-world-169m-model/
cp/Users/bbuf/工作目錄/RWKV/RWKV-World-HF-Tokenizer/rwkv_world_tokenizer/tokenization_rwkv_world.py../rwkv5-world-169m-model/
cp/Users/bbuf/工作目錄/RWKV/RWKV-World-HF-Tokenizer/rwkv_world_tokenizer/tokenizer_config.json../rwkv5-world-169m-model/
另外，您需要在生成文件夾中的 config.json 文件開頭添加以下幾行：
"architectures":[
"RwkvForCausalLM"
],
"auto_map":{
"AutoConfig":"configuration_rwkv5.Rwkv5Config",
"AutoModelForCausalLM":"modeling_rwkv5.RwkvForCausalLM"
},
運行Huggingface的RWKV World模型
run_hf_world_model_xxx.py演示了如何使用Huggingface的AutoModelForCausalLM加載轉換后的模型，以及如何使用通過AutoTokenizer加載的自定義RWKV World模型的HuggingFace版本的Tokenizer進行模型推斷。請看0x1節，這里就不贅述了。
檢查Lambda
如果你想運行這兩個腳本，首先需要下載一下 https://github.com/BlinkDL/ChatRWKV ，然后cd到rwkv_pip_package 目錄做pip install -e . ，然后再回到ChatRWKV的v2目錄下面，運行這里提供的lambda_pt.py和lambda_hf.py。
check_lambda文件夾下的lambda_pt.py和lambda_hf.py文件分別使用RWKV4 World 169M的原始PyTorch模型和HuggingFace模型對lambda數據集進行評估。從日志中可以看出，他們得到的評估結果基本上是一樣的。
lambda_pt.py lambda評估日志
#CheckLAMBADA...
#100ppl42.41acc34.0
#200ppl29.33acc37.0
#300ppl25.95acc39.0
#400ppl27.29acc36.75
#500ppl28.3acc35.4
#600ppl27.04acc35.83
...
#5000ppl26.19acc35.84
#5100ppl26.17acc35.88
#5153ppl26.16acc35.88
lambda_hf.py lambda評估日志
#CheckLAMBADA...
#100ppl42.4acc34.0
#200ppl29.3acc37.0
#300ppl25.94acc39.0
#400ppl27.27acc36.75
#500ppl28.28acc35.4
#600ppl27.02acc35.83
...
#5000ppl26.17acc35.82
#5100ppl26.15acc35.86
#5153ppl26.14acc35.86
從lambda的輸出結果可以驗證原始基于ChatRWKV系統運行的模型和轉換后的Huggingface模型是否精度是等價的。
0x3. 實現
Tokenizer的實現
Tokenizer的實現分為兩步。
因為目前社區的RWKV world模型tokenizer文件是一個txt文件：https://github.com/BBuf/RWKV-World-HF-Tokenizer/blob/main/rwkv_vocab_v20230424.txt 。我們需要將其轉換為Huggingface的AutoTokenizer可以讀取的json文件。這一步是通過 https://github.com/BBuf/RWKV-World-HF-Tokenizer/blob/main/scripts/convert_vocab_json.py 這個腳本實現的，我們對比一下執行前后的效果。
在這里插入圖片描述轉換為json文件后：
在這里插入圖片描述這里存在一個轉義的關系，讓gpt4解釋一下u0000和x00的關系：
在這里插入圖片描述有了這個json文件之后，我們就可以寫一個繼承PreTrainedTokenizer類的RWKVWorldTokenizer了，由于RWKV world tokenzier的原始實現是基于Trie樹（https://github.com/BBuf/RWKV-World-HF-Tokenizer/blob/main/rwkv_tokenizer.py#L5），所以我們實現 RWKVWorldTokenizer 的時候也要使用 Trie 樹這個數據結構。具體的代碼實現在 https://github.com/BBuf/RWKV-World-HF-Tokenizer/blob/main/rwkv_world_tokenizer/tokenization_rwkv_world.py 這個文件。
需要注意 https://github.com/BBuf/RWKV-World-HF-Tokenizer/blob/main/rwkv_world_tokenizer/tokenization_rwkv_world.py#L211 這一行，當解碼的token id是0時表示句子的結束（eos），這個時候就停止解碼。
RWKV World 5.0模型實現
實現了rwkv world tokenizer之后就可以完成所有rwkv4 world模型轉換到HuggingFace模型格式了，因為rwkv4 world模型的模型結構和目前transformers里面支持的rwkv模型的代碼完全一樣，只是tokenzier有區別而已。
但是如果想把 https://huggingface.co/BlinkDL/rwkv-5-world 這里的模型也轉換成 HuggingFace模型格式，那么我們就需要重新實現一下模型了。下面紅色部分的這個模型就是rwkv5.0版本的模型，剩下的都是5.2版本的模型。
在這里插入圖片描述我對照著ChatRWKV里面rwkv world 5.0的模型實現完成了HuggingFace版本的rwkv world 5.0版本模型代碼實現，具體見：https://github.com/BBuf/RWKV-World-HF-Tokenizer/blob/main/rwkv_world_v5.0_model/modeling_rwkv5.py ，是在transformers官方提供的實現上進一步修改得來。
實現了這個模型之后就可以完成rwkv world 5.0的模型轉換為HuggingFace結果了，教程請看上一節或者github倉庫。成品：https://huggingface.co/RWKV/rwkv-5-world-169m
0x4. 踩坑
在實現tokenizer的時候，由于RWKV world tokenzier的eos是0，我們沒辦法常規的去插入eos token，所以直接在Tokenzier的解碼位置特判0。
RWKVWorldTokenizer的初始化函數在 super().__init__ 的時機應該是在構造self.encoder之后，即：
在這里插入圖片描述否則會在當前最新的transformers==4.34版本中造成下面的錯誤：
在這里插入圖片描述在模型實現部分踩了一個坑，在embedding的時候第一次需要需要對embedding權重做一個pre layernorm的操作：
在這里插入圖片描述這個操作只能在第一次輸入的時候做一般就是prefill，我沒注意到這一點導致后面decoding的時候也錯誤做了這個操作導致解碼的時候亂碼，后面排查了不少時間才定位到這個問題。
另外，在做check lambda的時候，如果不開啟torch.no_grad上下文管理器禁用梯度計算，顯存會一直漲直到oom：https://github.com/BBuf/RWKV-World-HF-Tokenizer/blob/main/check_lambda/lambda_hf.py#L48 。
0x5. 總結
這件事情大概花了國慶的一半時間加一個完整的周六才搞定，所以這篇文章記錄一下也可以幫助有相關需求的小伙伴們少踩點坑。
-
人工智能
+關注
關注
1780文章
44671瀏覽量
231532 -
語言模型
+關注
關注
0文章
463瀏覽量
10127 -
數據集
+關注
關注
4文章
1182瀏覽量
24446
原文標題：0x5. 總結
文章出處：【微信號：GiantPandaCV，微信公眾號：GiantPandaCV】歡迎添加關注！文章轉載請注明出處。
發布評論請先 登錄
相關推薦
基于Serverless的前后端一體化框架
IC前端和后端設計的區別
數字IC后端設計介紹，寫給哪些想轉IC后端的人！
什么是后端?
基于層次法實現EOS芯片的后端設計
后端系統,后端系統是什么意思
芯片后端報告之A53后端顯示結果的解讀
工商網監
評論