Zenn
🐍

【手順】RWKV-LMがローカルPCで動作するまで

2023/03/28に公開
Python
PyTorch
#
LLM
#
rwkv
tech

最近話題の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します。

RWKV-LM

git clone https://github.com/BlinkDL/RWKV-LM

HuggingFaceからモデルをダウンロード

BlinkDL

  • (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