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)
レスポンス、かなり速いです。
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使ってます。

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

全体としてこんな感じのコードが生成されました。
最初に脈絡なく出てるint count;
をクラスに含めるだけで、とりあえず実行できる。

なんか小さいモデルでちゃんと動くコードが生成されてすごい。
14Bを動かしたいけど8GB VRAMだと無理そう。なのでフルCPUでやる必要があるけど、3時間たっても読み込みが終わらずで試せてません。
7Bでかなり賢いので、結構期待しています。
(6時間まったら読み込み終わったけど、そのあと入力しても反応なし・・・おそらく4時間くらい待つと1単語くらい出力されると思う)

こうやって使えるモデルが手元で動くと、いろいろ試せて楽しそうです。
LLMを動かすのに一番コスパがいいのはRTX3060で、12GBメモリ載っていますね。入手しておくか。