完全フリーで3GBのVRAMでも超高速に動く14B大規模言語モデルRWKVを試す

shi3z

2023年3月26日 11:20

Transformerは分散できる代償として計算量が爆発的に多いという不利がある。

一度みんなが忘れていたリカレントニューラルネットワーク(RNN)もボケーっとしている場合ではなかった。

なんと、GPT3並の性能を持つ、しかも完全にオープンな大規模言語モデルが公開されていた。

そのなもRWKV(RuwaKuvと発音しろと書いてある。ルワクフ?)
RWKVはRNNなのでGPUメモリをそれほど大量に必要としない。
3GBのVRAMでも動くという。

時間がない方はビデオをご覧ください

それに何より完全にフリーである。Apacheライセンス2.0
これはまさに民主化。

ただ、何日か前から挑戦していたのだが、モデルのダウンロードが全然終わらなくて試すことができないでいた。

とりあえず試したいだけならhuggingfaceですぐ試せる

で、モデルデータはここにある

僕の失敗は、何も考えずにgit lfs installでディレクトリごとコピーしようとしたこと。

このディレクトリには過去のモデルデータが全部あるので、ひとつ30GBのモデルデータを何十個もダウンロードしようとしていて終わらなかったのだ。

モデルデータは一とつあれば十分なのでひとつだけにする。

次に、chatのリポジトリをクローンする

$ git clone https://github.com/BlinkDL/ChatRWKV/tree/main/v2

ここのchat.pyを修正
まずはファイルの場所

次に、CPUのみで実行する場合は、「ストラテジー」を書き換える。

これで

$ python3 chat.py

だけで動く
が、CPUだけだと当たり前だが鬼のようにおそい


ChatRWKV v2 https://github.com/BlinkDL/ChatRWKV

English - cuda fp16 - /Users/shi3z/git/ChatRWKV/v2/prompt/default/English-2.py
Loading model - RWKV-4-Pile-14B-20230313-ctx8192-test1050.pth
RWKV_JIT_ON 1 RWKV_CUDA_ON 0 RESCALE_LAYER 6

Loading RWKV-4-Pile-14B-20230313-ctx8192-test1050.pth ...
Strategy: (total 40+1=41 layers)
* cuda [float16, float16], store 41 layers
0-cuda-float16-float16 1-cuda-float16-float16 2-cuda-float16-float16 3-cuda-float16-float16 4-cuda-float16-float16 5-cuda-float16-float16 6-cuda-float16-float16 7-cuda-float16-float16 8-cuda-float16-float16 9-cuda-float16-float16 10-cuda-float16-float16 11-cuda-float16-float16 12-cuda-float16-float16 13-cuda-float16-float16 14-cuda-float16-float16 15-cuda-float16-float16 16-cuda-float16-float16 17-cuda-float16-float16 18-cuda-float16-float16 19-cuda-float16-float16 20-cuda-float16-float16 21-cuda-float16-float16 22-cuda-float16-float16 23-cuda-float16-float16 24-cuda-float16-float16 25-cuda-float16-float16 26-cuda-float16-float16 27-cuda-float16-float16 28-cuda-float16-float16 29-cuda-float16-float16 30-cuda-float16-float16 31-cuda-float16-float16 32-cuda-float16-float16 33-cuda-float16-float16 34-cuda-float16-float16 35-cuda-float16-float16 36-cuda-float16-float16 37-cuda-float16-float16 38-cuda-float16-float16 39-cuda-float16-float16 40-cuda-float16-float16 
emb.weight                        f16      cpu  50277  5120 
Traceback (most recent call last):
  File "/Users/shi3z/git/ChatRWKV/v2/chat.py", line 120, in <module>
    model = RWKV(model=args.MODEL_NAME, strategy=args.strategy)
  File "/Users/shi3z/Library/Python/3.9/lib/python/site-packages/torch/jit/_script.py", line 292, in init_then_script
    original_init(self, *args, **kwargs)
  File "/Users/shi3z/Library/Python/3.9/lib/python/site-packages/rwkv/model.py", line 273, in __init__
    w[x] = w[x].to(device=DEVICE).contiguous()
  File "/Users/shi3z/Library/Python/3.9/lib/python/site-packages/torch/cuda/__init__.py", line 239, in _lazy_init
    raise AssertionError("Torch not compiled with CUDA enabled")
AssertionError: Torch not compiled with CUDA enabled
shi3z@shi3zettonoMacBook-Pro v2 % python3 chat.py


ChatRWKV v2 https://github.com/BlinkDL/ChatRWKV

English - cpu fp32 - /Users/shi3z/git/ChatRWKV/v2/prompt/default/English-2.py
Loading model - RWKV-4-Pile-14B-20230313-ctx8192-test1050.pth
RWKV_JIT_ON 1 RWKV_CUDA_ON 0 RESCALE_LAYER 0

Loading RWKV-4-Pile-14B-20230313-ctx8192-test1050.pth ...
Strategy: (total 40+1=41 layers)
* cpu [float32, float32], store 41 layers
0-cpu-float32-float32 1-cpu-float32-float32 2-cpu-float32-float32 3-cpu-float32-float32 4-cpu-float32-float32 5-cpu-float32-float32 6-cpu-float32-float32 7-cpu-float32-float32 8-cpu-float32-float32 9-cpu-float32-float32 10-cpu-float32-float32 11-cpu-float32-float32 12-cpu-float32-float32 13-cpu-float32-float32 14-cpu-float32-float32 15-cpu-float32-float32 16-cpu-float32-float32 17-cpu-float32-float32 18-cpu-float32-float32 19-cpu-float32-float32 20-cpu-float32-float32 21-cpu-float32-float32 22-cpu-float32-float32 23-cpu-float32-float32 24-cpu-float32-float32 25-cpu-float32-float32 26-cpu-float32-float32 27-cpu-float32-float32 28-cpu-float32-float32 29-cpu-float32-float32 30-cpu-float32-float32 31-cpu-float32-float32 32-cpu-float32-float32 33-cpu-float32-float32 34-cpu-float32-float32 35-cpu-float32-float32 36-cpu-float32-float32 37-cpu-float32-float32 38-cpu-float32-float32 39-cpu-float32-float32 40-cpu-float32-float32 
emb.weight                        f32      cpu  50277  5120 
blocks.0.ln1.weight               f32      cpu   5120       
blocks.0.ln1.bias                 f32      cpu   5120       
blocks.0.ln2.weight               f32      cpu   5120       
blocks.0.ln2.bias                 f32      cpu   5120       
blocks.0.att.time_decay           f32      cpu   5120       
blocks.0.att.time_first           f32      cpu   5120       
blocks.0.att.time_mix_k           f32      cpu   5120       
blocks.0.att.time_mix_v           f32      cpu   5120       
blocks.0.att.time_mix_r           f32      cpu   5120       
blocks.0.att.key.weight           f32      cpu   5120  5120 
blocks.0.att.value.weight         f32      cpu   5120  5120 
blocks.0.att.receptance.weight    f32      cpu   5120  5120 
blocks.0.att.output.weight        f32      cpu   5120  5120 
blocks.0.ffn.time_mix_k           f32      cpu   5120       
blocks.0.ffn.time_mix_r           f32      cpu   5120       
blocks.0.ffn.key.weight           f32      cpu   5120 20480 
blocks.0.ffn.receptance.weight    f32      cpu   5120  5120 
blocks.0.ffn.value.weight         f32      cpu  20480  5120 
............................................................................................................................................................................................................................................................................................................................................................................................................................................................................................................................................................................................................................................................................................................
blocks.39.ln1.weight              f32      cpu   5120       
blocks.39.ln1.bias                f32      cpu   5120       
blocks.39.ln2.weight              f32      cpu   5120       
blocks.39.ln2.bias                f32      cpu   5120       
blocks.39.att.time_decay          f32      cpu   5120       
blocks.39.att.time_first          f32      cpu   5120       
blocks.39.att.time_mix_k          f32      cpu   5120       
blocks.39.att.time_mix_v          f32      cpu   5120       
blocks.39.att.time_mix_r          f32      cpu   5120       
blocks.39.att.key.weight          f32      cpu   5120  5120 
blocks.39.att.value.weight        f32      cpu   5120  5120 
blocks.39.att.receptance.weight   f32      cpu   5120  5120 
blocks.39.att.output.weight       f32      cpu   5120  5120 
blocks.39.ffn.time_mix_k          f32      cpu   5120       
blocks.39.ffn.time_mix_r          f32      cpu   5120       
blocks.39.ffn.key.weight          f32      cpu   5120 20480 
blocks.39.ffn.receptance.weight   f32      cpu   5120  5120 
blocks.39.ffn.value.weight        f32      cpu  20480  5120 
ln_out.weight                     f32      cpu   5120       
ln_out.bias                       f32      cpu   5120       
head.weight                       f32      cpu   5120 50277 

Run prompt...
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)

Now talk with the bot and enjoy. Remember to +reset periodically to clean up the bot's memory. Use RWKV-4 14B for best results.
This is not instruct-tuned for conversation yet, so don't expect good quality. Better use +gen for free generation.

Prompt is VERY important. Try all prompts on https://github.com/BlinkDL/ChatRWKV first.

English - RWKV-4-Pile-14B-20230313-ctx8192-test1050.pth - cpu fp32

The following is a coherent verbose detailed conversation between a girl named Alice and her friend Bob. Alice is very intelligent, creative and friendly. Alice is unlikely to disagree with Bob, and Alice doesn't like to ask Bob questions. Alice likes to tell Bob a lot about herself and her opinions. Alice usually gives Bob kind, helpful and informative advices.

Bob: Hello Alice, how are you doing?
Alice: Hi! Thanks, I'm fine. What about you?

Bob: I am fine. It's nice to see you. Look, here is a store selling tea and juice.
Alice: Sure. Let's go inside. I would like to have some Mocha latte, which is my favourite!

Bob: What is it?
Alice: Mocha latte is usually made with espresso, milk, chocolate, and frothed milk. Its flavors are frequently sweet.

Bob: Sounds tasty. I'll try it next time. Would you like to chat with me for a while?
Alice: Of course! I'm glad to answer your questions or give helpful advices. You know, I am confident with my expertise. So please go ahead!

Bob:

で、語りかける時は、「+gen」を頭につけると生成されるとか、「+++」だと続きが読めるとかあるのでまとめておく

+ --> 他の返答を生成させる
+gen YOUR PROMPT --> 自由な文章生成。 \nを使うと行を変更できる
+qa YOUR QUESTION --> 自由な質問
+++ --> 直前の自由生成を続ける( +gen / +qa の続きのみ)
++ --> 同じ生成を再度行う ( +gen / +qa の時のみ)
+reset --> 会話をリセットする

ちなみにMacBookPro M2 MAX(32GB)でやると、一単語を生成するのに2分くらいかかるのでGPUはやはりあった方が良さそう。
でもこれで本当に「フリーな」大規模言語モデルが出てきたくれたって感じ。

ちなみにすでにLoRAとかもある

あと、INT8版もある

huggingfaceで試したところ、かなりのスピードで動作する
騒ごうぜ!