おうちの8GB VRAM GPUでChatRWKVと会話する

ChatGPTが話題ですが、そういうのをおうちで動かしたいですよね。
ということで、おうちで動かしやすくて割と会話ができるChatRWKVを試してみます。

ChatGPTは実装は公開されておらず手元で動かすことはできません。けど、サービスがたまに落ちてたりするので手元で動かせるなにかがあると安心ですね。何より、こんな変な技術を手元で動かしたい。
そこで手元で動かせるLLMとしてFlexGenが話題になりましたが、それでも強力なGPUが必要です。
https://github.com/FMInference/FlexGen

そして低リソースで動かせる実装としてAlpaca.cppなんかも出ましたが、そこまで賢くない感じ。
https://github.com/antimatter15/alpaca.cpp

そこに現れたのがChatRWKVです。一般のご家庭の8GB VRAMのGPUで動かせつつ、結構ちゃんと会話ができます。
https://github.com/BlinkDL/ChatRWKV

TransformerじゃなくてRNNであることで計算量を押さえれてるらしいです。Transformerが出てきてRNNは用なしになってしまったかと思ったら、そんなことはなかったですね。

ChatRWKVのインストール

いろいろいじりたいのにで、ちゃんとローカルに入れます。
インストールといっても特別な手順は不要で、上記GitHub リポジトリをCloneするかダウンロードします。

動かすのはv2/chat.pyです。

tokenizesとprompt_toolkitあたりが必要なので、インストールしておきます。

> pip install tokenizers prompt_toolkit

あと、PyTorchがCUDA対応じゃなくてGPUで動かせなかったのでインストールします。
ここでpipコマンドを生成。
https://pytorch.org/get-started/locally/

強制インストールするために-Iを付けておきます。大文字のアイ。

pip3 install -I torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

モデルのダウンロード

言語モデルは別でダウンロードが必要です。大きさが数段階あって、動かし方も選べるのだけど、目安のメモリ量はこんな感じ。

https://pypi.org/project/rwkvstic/

モデルはここからダウンロードします。まあ3B、7B、14Bくらいになるんではないかと。
https://huggingface.co/BlinkDL

小さいモデルで動かしてみる

3Bを動かしてみます。 RWKV-4-Pile-3B-20221110-ctx4096.pthをダウンロードして適当な場所に置いておきます。
https://huggingface.co/BlinkDL/rwkv-4-pile-3b

v2/chat.pyをいじって、次のような設定をします。

args.strategy = 'cuda fp16'
args.MODEL_NAME = 'RWKV-4-Pile-3B-20221110-ctx4096.pth'

fp16i8にすれば、4GB VRAMでも動くかもしれません。動かなかったら後述の7Bで試しているよう、CPUと組み合わせればいけそう。

MODEL_NAMEにはダウンロードしたモデルを指定します。ウィンドウズでもバックスラッシュじゃなくスラッシュを使えって書いてある。

さて起動するとこんな出力になります。

これで6GBくらいVRAM使います。

コマンドはこんなのが用意されてます。

Commands:
say something --> chat with bot. use \n for new line.
+ --> alternate chat reply
+reset --> reset chat

+gen YOUR PROMPT --> free generation with any prompt. use \n for new line.
+qa YOUR QUESTION --> free generation - ask any question (just ask the question). use \n for new line.
+++ --> continue last free generation (only for +gen / +qa)
++ --> retry last free generation (only for +gen / +qa)

レスポンス、かなり速いです。

日本語も理解してくれてるっぽい。英語で返ってくるけど。 pic.twitter.com/aQO3mwisc2
— きしだൠ(K1S) (@kis) 2023年3月26日

Javaコードもなんか生成してくれます。動かなそうだけど。

少し大きいモデルで動かす

それでは、少し大きい7Bモデルを動かしてみます。ここからダウンロード。
https://huggingface.co/BlinkDL/rwkv-4-pile-7b

RWKV-4-Pile-7B-20230109-ctx4096.pthがlikely bestらしいのでこれを。 8GB VRAMで7Bを動かすにはGPUを8bitで動かしつつ後段をCPUで動かすことが必要です。

args.strategy = 'cuda fp16i8 -> cpu fp32 *10'
args.MODEL_NAME = 'RWKV-4-Pile-7B-20230109-ctx4096.pth'

8K token用にはRWKV-4-Pile-7B-20230313-ctx8192-test380.pthがいいと書いてあります。長くない場合には4096より性能が劣るっぽい。

メモリは5.9GB使ってます。

動作速度としてはこんな感じ。

おー動いた！ちょっと重いけど、半分CPUで動いてるにしてはかなり速い。そして日本語もなんかあやしいけど受け答えできてる。そしてSF作ってって言ったら死んだｗ(ターミナルごと落ちた)
日本語は落ちがちかな。 pic.twitter.com/5vd24jMuaf
— きしだൠ(K1S) (@kis) 2023年3月26日