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

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

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

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

RWKVをチャットで使えるようにしたものとして[ChatRWKV](https://github.com/BlinkDL/ChatRWKV)があります。
また、RWKVをAlpacaデータセットなどを使用してファインチューニングしたものとして[Raven](https://huggingface.co/BlinkDL/rwkv-4-raven)があります。

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

https://zenn.dev/karaage0703/articles/d58d79d8e77ab8

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

https://note.com/npaka/n/n8f3c2c491901

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

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

```Python:変更前
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
```

```diff_Python:変更後
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だったことが原因でこのエラーが発生したようです。

https://github.com/NVIDIA/cuda-samples/issues/57

原因が分かったので[cuda toolkit 11.2](https://developer.nvidia.com/cuda-11.2.0-download-archive)をインストールして再トライしてみたら無事動作しました。

念のため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](https://huggingface.co/BlinkDL/rwkv-4-raven/blob/main/RWKV-4-Raven-14B-v6-EngChnJpn-20230401-ctx4096.pth)(int8化しています)  |
| CHUNK_LEN | 128 |

## RWKV_CUDA_ONが0の場合
![RWKV_CUDA_ON0.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/571964/8a89fe2e-7940-0975-400b-00bdd43ac04f.gif)

## RWKV_CUDA_ONが1の場合
![RWKV_CUDA_ON1.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/571964/8c93ab46-1c36-cdea-d954-50d9a66a1912.gif)

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

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

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