Google Colab で RWKV のファインチューニングを試す
「Google Colab」で「RWKV」のファインチューニングを試したので、まとめました。
前回
1. RWKV
「RWKV」は、TransformerレベルのLLM性能を備えたRNNです。高性能、高速推論、VRAMの節約、高速学習、長い文脈長、自由な埋め込みを実現しています。
2. データセットの準備
今回は、練習用に「あかねと〜くデータセット」の「dataset.txt」を使います。
テキスト終端をスペシャルトークン<|endoftext|> で記述しています。
どのようなスペシャルトークンがあるかは、RWKVのトークナイザーの設定ファイル「20B_tokenizer.json」で確認できます。
3. Colabでの学習
Colabでの学習手順は、次のとおりです。
(1) 以下のの学習用の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」を指定します。
