Google Colab で RWKV のファインチューニングを試す

npaka

2023年4月6日 20:59

「Google Colab」で「RWKV」のファインチューニングを試したので、まとめました。

前回

1. RWKV

「RWKV」は、TransformerレベルのLLM性能を備えたRNNです。高性能、高速推論、VRAMの節約、高速学習、長い文脈長、自由な埋め込みを実現しています。

2. データセットの準備

今回は、練習用に「あかねと〜くデータセット」の「dataset.txt」を使います。

テキスト終端をスペシャルトークン<|endoftext|> で記述しています。
どのようなスペシャルトークンがあるかは、RWKVのトークナイザーの設定ファイル「20B_tokenizer.json」で確認できます。

3. Colabでの学習

Colabでの学習手順は、次のとおりです。

(1) 以下のの学習用のColabノートブックを開き、メニュー「ファイル→ドライブにコピーを保存」。

・RWKV-4 PileのファインチューニングのためのColab

(2) 「Setup」の4セルを実行。
4セル目の「pytorch-lightning」は、1.7を指定します。指定しないと2.0がインストールされ動きませんでした。

!pip install transformers pytorch-lightning==1.7 deepspeed wandb ninja

(3) 「Load Base Model」で、モデルを選択してセルを実行。
今回は、練習のため一番小さい「RWKV-v4-Pile-169M」を選択します。Pileで事前学習してあるファインチューニング用のモデルになります。

・RWKV-v4-Pile-169M  : VRAM 10.2GB
・RWKV-v4-Pile-490M  : VRAM 21.7GB、プレミアム必要
・RWKV-v4-Pile-1B5  : プレミアムでもOutOfMemory (batch_size減らせば学習できるかも)

「HuggingFace」で7Bと14Bのモデルも提供されています。

(4) 「dataset.txt」をGoogleドライブにアップロードし、「Training Data Options」でそのパスを指定してセルを実行。

(5) 「Training」の2セルを実行。
今回は、学習パラメータはデフォルトのままとしました。

・n_epoc : 何エポックで tunesX.pth を保存するか
・epoc_save_frequency : 何エポック毎に tunesX.pth を保存するか
・batch_size : バッチサイズ
・ctl_len : コンテキスト長
・T_MAX : ???
・RWKV_FLOAT_MODE : floatモード (fp16 / bf16 / bf32)
・RWKV_DEEPSPEED : DEEPSPEED (0 / 1)

1エポックにかかる時間は、無料版Colabは5分、プレミアムは1分ほどです。Googleドライブの「rwvk-v4-rnn-training」に「tunedX.pth」(Xは1、26、51、76、101)の名前で生成されます。n_epoch過ぎても止まらないので、十分学習できたら自分で停止します。

4. Colabでの推論

前回と同じです。モデルのパスとして生成した「tunedX」を指定します。

次回

この記事が気に入ったらサポートをしてみませんか？