LoginSignup
2
2

RWKVで量子化を試す

Last updated at Posted at 2023-07-23

SS 192.png こんどはうまく行った。

やること

前回の記事で、RTX-3090とCUDAを使って、フリーの大規模言語モデルであるChatRWKVが使えるようになりました。

今回はその性能をUPしていく作業に入りたいと思います。

作業環境

  • Windows10 Home 22H2 (64bit x68ベース)
  • RTX3090 24GB
  • CPU i7-4770 CPU @ 3.40GHz
  • メモリ 16GB

最大サイズの14Bをいきなり試す

HugginfaceからRWKV-4-Raven-14B-v12-Eng98%25-Other2%25-20230523-ctx8192.pthをDLして、ローカルのChatRWKV/v2に置きます。
chat.pyの該当箇所を書き換えて起動してみます。
ひょっとしたらなんとかなるかもと思いましたが、14Bはそもそものモデルのファイルサイズが28GBあるので、RTX3090の24GBをもってしても無理でした。
なにか方法があるかもしれませんが、本体のメモリも少ないですし一旦諦めます。

データの量子化とは?

量子化ってなんでしょう。とにかく響きはかっこいいですよね。

ここで大規模言語モデルを非常にざっくり説明します。

気が遠くなるほど大量のテキストデータを、何億という数の 単語に分類し、その単語の一つ一つに何十億、何千億という個数のパラメータをつけ、そのパラメータのそれぞれが度合いで表されるようにしたものです。
パラメーターは人智を超えた単語の関係性を示しており、その関係性に基づいて単語を並べていくと意味が通じる文章になる。というものです。(違ってたらすみません。)

この量子化という作業は、「度合い」の表し方を、16bitの細かさから8bitまでの度合い表現の解像度を荒くすることです。たぶん。
そうすると単純にデータが半分になり、軽量化します。

データの量子化を実行する

ここで「RWKV-4-Raven-14B-v12-Eng98%-Other2%-20230523-ctx8192.pth」というファイルを8ビットに量子化し、Raven-14B98Eng-cudafp16i8.pthという名前で保存するとします。

まず、chat.pyの中身も、変更後のモデルを使うように指定しておきます。
(変換ツールの最後にchat.pyを確認しているような気がなんとなくしています)

(中略)
# args.strategy = 'cpu fp32'
# args.strategy = 'cuda fp16'
# args.strategy = 'cuda:0 fp16 -> cuda:1 fp16'
# args.strategy = 'cuda fp16i8 *10 -> cuda fp16'
# args.strategy = 'cuda fp16i8 *10 -> cuda fp16'
args.strategy = 'cuda fp16i8'
# args.strategy = 'cuda fp16i8 -> cpu fp32 *10'
# args.strategy = 'cuda fp16i8 *10+'
(中略)
CHAT_LANG = 'Japanese' # English // Chinese // more to come
(中略)
elif CHAT_LANG == 'Japanese':
    args.MODEL_NAME = 'Raven-14B98Eng-cudafp16i8'

変換を実行します。
ChatRWKV/v2のディレクトリに入り、

python convert_model.py --in RWKV-4-Raven-14B-v12-Eng98%-Other2%-20230523-ctx8192.pth --out Raven-14B98Eng-cudafp16i8.pth --strategy "cuda fp16i8"

と入力して実行します。
"cuda fp16i8"の指定については、ダブルクォーテーションのほうがうまく動いてくれました。ダブルクォーテーションもシングルクォーテーションも、コピペ元が形が似ている別の文字を使っている場合があります。その場合はエラーになります。
シングルクォーテーションでも問題なく動く方はそのままでよいと思います。

8bit化した14Bファイルを実行する

ファイルを指定してchat.pyで動きます。

7Bに比べてもかなり賢くなったような感じがありますし、なにより質問文を入力してから1秒程度で返信開始してくれるクイックレスポンスが小気味よいです。

WS000150.JPG

すぐ返答してくれるので、いろいろな会話を実験できます。

RTX3090の24Gに乗るサイズにする

RWKVでは、データの量子化について単純に半分にする以外にも、何層目までの計算を8bitにするかを指定することができるようです。
ためしに10層まで量子化するという--strategy "cuda fp16i8 *10 -> cuda fp16"を試してみました。

結果として、無事に動いたものの、返答開始まで3~4秒ぐらいかかるようになり、文章の文字が出てくるスピードも格段に遅くなりました。手軽に楽しむには全層量子化のまま使うのがよさそうです。

次は実験しやすくしたい

chat.pyの起動に少し時間がかかるので、chat.pyをAPIのエンドポイントっぽく使えるように改造しておけると便利かな、と思いました。

では楽しいRWKVライフを。

参考

前回の記事

2
2
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
2
2