大規模言語モデルRWKVをローカルPCで試す

清水亮さんが、RWKVを試してたので俺も試してみた。
この時の清水さんのマシンは、A6000積んでる(VRAM48GB)逸般のご家庭用モンスターマシンなので、そのまま動く。そらそうよ。

npakaさんは、Google Colabで動かしている。

というわけで、一般のご家庭用ゲーミングPCでも動くように試してみた。

ちなみに試した環境はWindowsのWSLのubuntu上。
マシンのメモリは64GB。ubuntuには48GB割当。
GPUはRTX3090(VRAM 24GB)。


まずはChatRWKVをgitでcloneする。

git clone https://github.com/BlinkDL/ChatRWKV

次にモデルをダウンロードする。
14Bを使うことにした。

BlinkDL/rwkv-4-pile-14b · Hugging Face

ChatRWKV/v2フォルダにダウンロードしてきたモデルファイルを突っ込む。
今回は最新版の

RWKV-4-Pile-14B-20230313-ctx8192-test1050.pthを使う。

ChatRWKV/v2フォルダ内のchat.pyを編集する。

args.strategy = 'cpu fp32'
これをコメントアウトして
args.strategy = 'cuda fp16i8'
のコメントアウトを外す。

args.MODEL_NAME = '/fsx/BlinkDL/HF-MODEL/rwkv-4-pile-14b/RWKV-4-Pile-14B-20230228-ctx4096-test663'
を
args.MODEL_NAME = 'RWKV-4-Pile-14B-20230313-ctx8192-test1050'
に変更する

14Bのモデルは28GBあるが、fp16のINT8にすることで、なんとVRAM16GBで済んでしまう！(その分、精度などは落ちます)
ちなみに3BのモデルならINT8で2.8GBで乗りますので、多分これがVRAM3GBでOKの話なのかなと思います。
モデル別の必要VRAMはここ見るとわかりますよ。

rwkvstic

7BならINT8で8GBなので、普通のご家庭のゲーミングPCでVRAM10GBくらいあれば実現可能ですね。

あとはChatRWKV/v2フォルダ内のchat.pyを起動するだけ。

python3 chat.py

これで一般的なご家庭でもローカルで大規模言語モデルで遊べる。
prompt_toolkitがないとかいわれたら、

pip install prompt_toolkit

でインストールする。
もしかすると、tokenizersもないって言われるかもしれない

pip install tokenizers

で、インストールする。

もしかすると、rwkvsticもインストールしておいたほうがいいのかもしれない。

pip install rwkvstic

ちなみにGPUのメモリ使用量(14Bでfp16i8)

ただ、モデルがオープンソースのものを学習してるので、ChatGPTみたいな精度は出せない。
ただ、逆に考えるとモデルをファインチューニングしたものを作ればいいだけなので（だけって言っても、それが大変なんだけど）こういうローカルでも動かせるもの(flexgenとかalpacaとか)がどんどん出てきてるのがすごいなって感じですね。

小さいモデルのリンクも貼っとく。
7B

BlinkDL/rwkv-4-pile-7b · Hugging Face

ちなみに、7Bのこのモデル

RWKV-4-Pile-7B-Instruct-test4-20230326.pth · BlinkDL/rwkv-4-pile-7b at main

は、alpacaでファインチューニングしたやつなので、通常の7Bのよりは賢いはず。

3B

BlinkDL/rwkv-4-pile-3b · Hugging Face

１ヶ月で何回未来がくるんだよ！
みんなも未来を体感しようぜ！