"RWKV_CUDA_ON"を1に設定してRWKVを高速化する

ローカル環境でも高速で動作するLLMとして話題のRWKVですが、ドキュメントを見ていると環境変数の"RWKV_CUDA_ON"を1に設定することでより高速で動くということなので試してみました。

設定を変更するときに１箇所エラーで躓いたところがあったのでそちらも記録しておきます。

RWKVとは

RWKVとはTransformerレベルの性能を持つRNNです。他のLLMと同様に高性能でありつつ推論が高速でVRAMを節約しているのが特徴です。

RWKVをチャットで使えるようにしたものとしてChatRWKVがあります。
また、RWKVをAlpacaデータセットなどを使用してファインチューニングしたものとしてRavenがあります。

ChatRWKVの使い方は次の記事で分かりやすく解説されています。

またRWKVのファインチューニングの方法については以下の記事で解説されています。

RWKV_CUDA_ONを1に設定する

本題になりますが、環境変数RWKV_CUDA_ONを1に設定してRWKVを高速化していきます。

環境変数の設定ですがChatRWKVをcloneしてv2/chat.pyの部分の50行目あたりにある以下の部分を0から1に変更するだけです。

変更前

os.environ["RWKV_JIT_ON"] = '1' # '1' or '0', please use torch 1.13+ and benchmark speed
os.environ["RWKV_CUDA_ON"] = '0' # '1' to compile CUDA kernel (10x faster), requires c++ compiler & cuda libraries

変更後

os.environ["RWKV_JIT_ON"] = '1' # '1' or '0', please use torch 1.13+ and benchmark speed
+ os.environ["RWKV_CUDA_ON"] = '1' # '1' to compile CUDA kernel (10x faster), requires c++ compiler & cuda libraries

上記の変更を行いv2/Chat.pyを実行したところ筆者の環境では次のようなエラーがでました。

エラー

nvcc fatal   : Unknown option '-t'
ninja: build stopped: subcommand failed.
Traceback (most recent call last):
  File "/home/toshi_456/.pyenv/versions/3.8.10/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1717, in _run_ninja_build
    subprocess.run(
  File "/home/toshi_456/.pyenv/versions/3.8.10/lib/python3.8/subprocess.py", line 516, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['ninja', '-v']' returned non-zero exit status 1.

エラー文を見てみると "nvcc -t"(nvcc --threads) というコマンドが使えないことが原因でエラーが起きているようです。

その原因について調べてみると下記のissueでnvcc -tという機能は cuda toolkit 11.2 から追加されたものだということが分かりました。筆者の環境がcuda toolkit 11.1だったことが原因でこのエラーが発生したようです。

原因が分かったのでcuda toolkit 11.2をインストールして再トライしてみたら無事動作しました。

念のためUbuntu20.04でのcuda toolkit 11.2のインストール手順を残しておきます。

インストール

wget https://developer.download.nvidia.com/compute/cuda/11.2.0/local_installers/cuda_11.2.0_460.27.04_linux.run
sudo sh cuda_11.2.0_460.27.04_linux.run

環境変数の設定

export PATH=/usr/local/cuda-11.2/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-11.2/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

速度比較

無事設定ができたところRWKV_CUDA_ONが0の場合と1の場合の速度を簡単に比較していきます。

筆者の動作環境と使用したモデルやパラメータは以下のとおりです。CHUNK_LENはRWKV_CUDA_ONが0の場合デフォルトの256だとメモリに乗り切らないので128にしています。その他の表に書いていないパラメータはChatRWKVのデフォルト設定のままです。

項目	内容
CPU	Intel Core i9-11950H
メモリ	32GB
GPU	NVIDIA RTX3080(16GB)
モデル	RWKV-4-Raven-14B-v6-EngChnJpn-20230401-ctx4096.pth(int8化しています)
CHUNK_LEN	128

RWKV_CUDA_ONが0の場合

RWKV_CUDA_ONが1の場合

RWKV_CUDA_ONを1に設定した方が明らかに速いですね。

また、VRAMの効率も良くなっているような感じがして、CHUNK_LENを512に設定しても問題なく動作することを確認しています。(RWKV_CUDA_ONが0の場合はメモリ不足で落ちます)

まとめ

GPUメモリが16GBあればRWKVの14bモデルもそこそこのスピードで動くことが分かりました。
RWKVの今後に期待ですね。