【手順】RWKV-LMがローカルPCで動作するまで
最近話題のRWKV-LMのRWKC-v4neo 14Bモデルを手元のPCで動かしましたので手順を書き残しておきます。とても簡単に書いてありますが、だいたい私が実行した通りです。
スペック
- CPU Intel Core i7 12700K
- 64GB Memory(16GB x 4)
- NVIDIA GeForce RTX 3070(8GB) ※ただしCPUモードで動かしています
- Windows 11
- Python 3.9.12
- PyTorch 1.13.1+cu117
なお、17Bモデルは64GBメモリでは本来不足なので、swap領域を増やしています。Windows 11の場合、「設定」⇒「バージョン情報」⇒「システムの詳細設定」⇒「詳細設定」タブ⇒「パフォーマンス」の「設定」ボタン⇒「詳細設定」タブ⇒「仮想メモリ」の「変更」とたどって、「カスタムサイズ」で十分大きなswap領域を設定します(私の場合、初期サイズは「推奨」の値だった9055MB、最大サイズは102400MBとしました)。
なお、swapの変更後は再起動が必要です。これで、メモリが足りない場合はディスク上にデータを逃がしてある程度耐えてくれます(万能ではなくブルースクリーンになることもあるかも)。
GitHubからclone
レポジトリをcloneします。
git clone https://github.com/BlinkDL/RWKV-LM
HuggingFaceからモデルをダウンロード
- (14Bモデル) rwkv-4-pile-14b ⇒ Files and versionsタブ ⇒ RWKV-4-Pile-14B-20230213-8019.pth ⇒ download
- (7Bモデル) は同じようにたどってRWKV-4-Pile-7B-20230109-ctx4096.pth
- (3Bモデルもダウンロードしておくとテストによいかも)
ファイルがたくさん並んでいて新しいのがいいのかな?と思ってしまいますが、まずは後述のchat.pyをみて、そこにあるファイル名と同じものを入手しましょう。
chat.pyの調整
RWKV-LM / RWKV-v4neo / chat.pyを編集します。ファイルの冒頭2行に以下を作成します。
import os, copy, types, gc, sys # この行を最初に持ってくる
os.environ["RWKV_JIT_ON"] = "1" # 環境変数を追加
print('Loading...')
from src.model_run import RWKV_RNN
import numpy as np
少し下の以下の行を自分の環境に合わせてフルパスにします。
フォルダをたどるときは¥マークではなくスラッシュで。chat.pyと同じフォルダにある20B_tokenizer.jsonを指定します。CHAT_LANGはEnglishのままでよいです。日本語もある程度応答してくれます。
CHAT_LANG = 'English' # English Chinese
WORD_NAME = [
"D:/Python_vscode/RWKV/RWKV-LM/RWKV-v4neo/20B_tokenizer.json",
"D:/Python_vscode/RWKV/RWKV-LM/RWKV-v4neo/20B_tokenizer.json",
] # [vocab, vocab] for Pile model
その下のargsを編集します。私はCPUモードで動かしているので、RUN_DEVICEはcpu、FLOAT_MODEはfp32です。強いGPUを持っている方はcuda / fp16でよいのだと思います。
n_layer、n_embd、ctx_lenは決まった値が必要のようなので、4行ひと組で有効/無効にします。以下は14Bの時の例です。ここで、MODEL_NAMEが実際にあるファイルと同じになっているか確認します。違う場合は、そのモデルに合わせてパラメータも適切に変更する必要があるのかも(よく調べてないですが動かない)しれません。
args = types.SimpleNamespace()
args.RUN_DEVICE = "cpu" # 'cpu' (already very fast) // 'cuda'
args.FLOAT_MODE = "fp32" # fp32 (good for CPU) // fp16 (recommended for GPU) // bf16 (less accurate)
args.vocab_size = 50277
args.head_qk = 0
args.pre_ffn = 0
args.grad_cp = 0
args.my_pos_emb = 0
args.MODEL_NAME = 'D:/Python_vscode/RWKV/RWKV-4-Pile-14B-20230213-8019'
args.n_layer = 40
args.n_embd = 5120
args.ctx_len = 1024
#args.MODEL_NAME = 'D:/Python_vscode/RWKV/RWKV-4-Pile-7B-20230109-ctx4096'
#args.n_layer = 32
#args.n_embd = 4096
#args.ctx_len = 4096
# args.MODEL_NAME = 'D:/Python_vscode/RWKV/RWKV-4-Pile-3B-20221008-8023'
#args.n_layer = 32
#args.n_embd = 2560
#args.ctx_len = 1024
chat.pyの実行
> python .\chat.py
あとはそっと待つだけです。実行前に、タスクマネージャを起動してメモリのグラフを表示しておくと、いつWindowsがパンクするか、とスリルが味わえます。
プロンプトが起動したら
プロンプトに出てきますが、基本は "+gen (生成したい文章の冒頭部分)"と書いて実行するとそこから続きを書いてくれるようです。途中で生成が終わったら、"+more"で続きを要求します。別のことを生成するときは"+reset"を実行して状態をリセットしてから次の+genを入力します。
Discussion