Zenn
👌

ローカルで RWKV-LoRA の日本語データセットでﾌｧｲﾝﾁｭｰﾝするメモ (3090x2 で 7B いけた. 14B はだめだった)

2023/04/24に公開
LLM
RWKV
tech

Google Colab で RWKV-LoRA のファインチューニングを試す
https://note.com/npaka/n/ndb1df153aed5

ありがとうございます. こちらを参考にローカルでﾌｧｲﾝﾁｭｰﾝします!

14B はいくつか設定試しましたが CUDA out-of-memory でダメでした.

環境

  • Ubuntu 20.04
  • x299 + 256 GB CPU mem
    • 14B もしうまく学習回る場合では, 128 GB でも行けると思われる.
  • 3090(24GB) x 2

セットアップ

まず conda あたりで環境作って, 必要モジュールインストールしておきます.

$ python -m pip install transformers pytorch-lightning==1.9 deepspeed wandb ninja

tokenize

入力データを 20B_tokenizer でトークナイズします.

ちなみに↓の処理スクリプトのファイル名を tokenize.py というファイル名にしてしまうと, transformers 内のモジュールが tokenize という名前の module をインポートしているため,
自前 tokenize.py ファイルとかちあい,

from inspect import Parameter, signature
ImportError: cannot import name 'Parameter' from partially initialized module 'inspect'

というエラー出ますので注意ください!
(このよくわからんエラーで半営業日つぶしてしまった...)

tok_encode.py など, ほかの Python module 名とかちあわさなそうなファイル名にしましょう.

# tok_encode.py
import numpy as np
from transformers import PreTrainedTokenizerFast

tokenizer = PreTrainedTokenizerFast(tokenizer_file=f'20B_tokenizer.json')

input_file = "..."
output_file = 'train.npy'

data_raw = open(input_file, encoding="utf-8").read()

data_code = tokenizer.encode(data_raw)
print("tokenized len = ", len(data_code))

out = np.array(data_code, dtype='uint16')
np.save(output_file, out, allow_pickle=False)

変換自体は 1 分以内で終わります.

1B でテスト

https://huggingface.co/BlinkDL/rwkv-4-pile-1b5

の RWKV-4-Pile-1B5-20220903-8040.pth https://huggingface.co/BlinkDL/rwkv-4-pile-1b5/blob/main/RWKV-4-Pile-1B5-20220903-8040.pth を使います.

(なんか RWKV-4-Pile-1B5-20220929-ctx4096.pth のほうが最新っぽそうな気もするが...)

#!/bin/bash

base_model_path="../data/RWKV-4-Pile-1B5-20220903-8040.pth"
tuned_model_dir="../data/tuned-1b"

# 1b
n_layer=24
n_embd=2048

n_epoch=100
epoch_save_frequency=5
batch_size=1
ctx_len=1024
vocab_size=50277

python train.py \
  --load_model $base_model_path \
  --proj_dir $tuned_model_dir \
  --data_file "train.npy" \
  --data_type "numpy" \
  --vocab_size $vocab_size \
  --ctx_len $ctx_len \
  --epoch_save $epoch_save_frequency \
  --epoch_count $n_epoch \
  --n_layer $n_layer \
  --n_embd $n_embd \
  --epoch_steps 1000 --epoch_begin 0  --micro_bsz $batch_size --pre_ffn 0 --head_qk 0 --lr_init 1e-5 --lr_final 1e-5 --warmup_steps 0 --beta1 0.9 --beta2 0.999 --adam_eps 1e-8 --accelerator gpu --devices 1 --precision bf16 --strategy deepspeed_stage_2 --grad_cp 0 \
  --lora --lora_r 8 --lora_alpha 32 --lora_dropout 0.0

とりあえず GPU 1 個です(--devices 1).

MiB Mem : 257433.1 total,  90693.2 free,  14337.5 used, 152402.4 buff/cache
MiB Swap:   2048.0 total,   2048.0 free,      0.0 used. 240769.6 avail Mem           

PID USER      PR  NI    VIRT    RES    SHR S  %CPU  %MEM     TIME+ COMMAND
2763447 syoyo     20   0   29.2g   6.4g 727428 S 100.3   2.5   1:51.49 python

CPU memory 消費は 6 GB,

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.105.17   Driver Version: 525.105.17   CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  Off  | 00000000:65:00.0  On |                  N/A |
| 36%   39C    P5    66W / 160W |   8432MiB / 24576MiB |     98%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA GeForce ...  Off  | 00000000:B3:00.0  On |                  N/A |
|  0%   42C    P5    37W / 160W |     25MiB / 24576MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------

GPU は 8 GB くらいでした.

Epoch 0:  42%|████▌      | 415/1000 [05:16<07:25,  1.31it/s, loss=1.720, lr=1e-5, REAL it/s=1.330, Kt/s=1.360]

1 epoch 13 分くらいです(3090 は 160W に powerlimit)

npaka 先生の Colab Premium(A100?) では 1 epoch 3 分とのことです.

条件にもよりますが, 3090 より A100 のほうが BF16 では 2~4 倍ほど性能出ます.

https://github.com/huggingface/transformers/issues/14608#issuecomment-1001259427

ので, 妥当なところでしょうか(学習で bf16 においては, 3090 = 0.25 * A100 のイメージ)
あと, GPU 二枚だと (deepspeed による) GPU 間の通信オーバヘッドもあるでしょう.

bf16 ではなくて fp16(乗算精度も fp16) だと3090 だともうちょっと性能いいかもです.

--precision fp16 にしたら 3090 でもより学習早くなるかもしれません.
(ただ, RWLV-LoRA の finetune ドキュメントには, 演算オーバフローする可能性ありとありました. 学習時に発散してしまう可能性はあります)

7B

https://huggingface.co/BlinkDL/rwkv-4-pile-7b

RWKV-4-Pile-7B-20221115-8047.pth https://huggingface.co/BlinkDL/rwkv-4-pile-7b/blob/main/RWKV-4-Pile-7B-20221115-8047.pth

#!/bin/bash

base_model_path="../data/RWKV-4-Pile-7B-20221115-8047.pth"
tuned_model_dir="../data/tuned-7b"

# 7b
n_layer=32
n_embd=4096

n_epoch=100
epoch_save_frequency=5
batch_size=1
ctx_len=1024
vocab_size=50277

python train.py \
  --load_model $base_model_path \
  --proj_dir $tuned_model_dir \
  --data_file "train.npy" \
  --data_type "numpy" \
  --vocab_size $vocab_size \
  --ctx_len $ctx_len \
  --epoch_save $epoch_save_frequency \
  --epoch_count $n_epoch \
  --n_layer $n_layer \
  --n_embd $n_embd \
  --epoch_steps 1000 --epoch_begin 0  --micro_bsz $batch_size --pre_ffn 0 --head_qk 0 --lr_init 1e-5 --lr_final 1e-5 --warmup_steps 0 --beta1 0.9 --beta2 0.999 --adam_eps 1e-8 --accelerator gpu --devices 2 --precision bf16 --strategy deepspeed_stage_2 --grad_cp 0 \
  --lora --lora_r 8 --lora_alpha 32 --lora_dropout 0.01

Epoch 0:   5%|▌           | 52/1000 [02:11<40:03,  2.54s/it, loss=1.810, lr=1e-5, REAL it/s=0.405, Kt/s=0.830]

1 epoch 40 分くらいのようです...

MiB Mem : 257433.1 total, 126301.0 free,  38269.4 used,  92862.8 buff/cache
MiB Swap:   2048.0 total,   2046.7 free,      1.2 used. 215385.3 avail Mem

    PID USER      PR  NI    VIRT    RES    SHR S  %CPU  %MEM     TIME+ COMMAND
2766551 syoyo     20   0   54.5g  16.1g 785756 S 100.0   6.4   4:36.43 python
2766615 syoyo     20   0   54.4g  16.1g 784332 S 100.0   6.4   4:33.19 python

CPU は 32 GB, GPU は 24 GB フルフルでした.

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.105.17   Driver Version: 525.105.17   CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  Off  | 00000000:65:00.0  On |                  N/A |
| 36%   40C    P5    71W / 160W |  23654MiB / 24576MiB |    100%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA GeForce ...  Off  | 00000000:B3:00.0  On |                  N/A |
| 50%   36C    P5    75W / 160W |  23537MiB / 24576MiB |    100%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------

14B(GPU メモリ不足)

https://huggingface.co/BlinkDL/rwkv-4-pile-14b

RWKV-4-Pile-14B-20230213-8019.pth https://huggingface.co/BlinkDL/rwkv-4-pile-14b/blob/main/RWKV-4-Pile-14B-20230213-8019.pth を使います.

deepspeed_stage2_offload, deepspeed_stage3_offload など試しましたが, CUDA out-of-memory になってしまいました.

3090 x 4 か, A6000 x 2 みたいな環境が必要かもです...

あとは RWKV の trainer ソースいじって Deepspeed の設定変えたり(nvme 使う版とか!), pure(?) CPU で学習できるようにするなどすればなにやらいけるかもはしれません.

参考までに実行スクリプト.

#!/bin/bash

base_model_path="../data/RWKV-4-Pile-14B-20230213-8019.pth"
tuned_model_dir="../data/tuned-1b"

# 14b
n_layer=40
n_embd=5120

n_epoch=100
epoch_save_frequency=5
batch_size=1
ctx_len=1024
vocab_size=50277

python train.py \
  --load_model $base_model_path \
  --proj_dir $tuned_model_dir \
  --data_file "train.npy" \
  --data_type "numpy" \
  --vocab_size $vocab_size \
  --ctx_len $ctx_len \
  --epoch_save $epoch_save_frequency \
  --epoch_count $n_epoch \
  --n_layer $n_layer \
  --n_embd $n_embd \
  --epoch_steps 1000 --epoch_begin 0  --micro_bsz $batch_size --pre_ffn 0 --head_qk 0 --lr_init 1e-5 --lr_final 1e-5 --warmup_steps 0 --beta1 0.9 --beta2 0.999 --adam_eps 1e-8 --accelerator gpu --devices 2 --precision fp16 --strategy deepspeed_stage_3_offload --grad_cp 0 \
  --lora --lora_r 8 --lora_alpha 32 --lora_dropout 0.01

その他

deepspeed の設定について.

https://lightning.ai/docs/pytorch/stable/advanced/model_parallel.html#deepspeed

https://github.com/Blealtan/RWKV-LM-LoRA/blob/c5cb3efcc003ee8d8352363a6e1361f53febeb68/RWKV-v4neo/train.py#L41

fp16 だと計算がオーバーフローする可能性あります.

ﾌｧｲﾝﾁｭｰﾝしたので chat する!

20 epoch まで終わったので動かしてみます!
推論は別 PC の 3090x1 で行いました.

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

npaka 先生の記事を参考に "+i" の処理を追加など chat.py をいじります!
(しかし, ソースコードのわかりずらさよ...)

動かしてみます!

# Instruction:
三原色とは何ですか?

# Response:
]
三原色は青、緑、黄、黒の3色で構成されています。

青は明るく強く、青色は互いに最も強い光の一種であり、幾何学的な幾何学的な考え方に基づいています。

緑は照明のような形の明るさをもたらし、光の色は誰もが彼らの旅の幻想的な気分を味わえます。

グレーは光の透明をもた

User:

うーん, おしい.

ちなﾌｧｲﾝﾁｭｰﾝ元のデータセットは

Below is an instruction that describes a task. Write a response that appropriately completes the request.

# Instruction:
三原色とは何ですか?

# Response:
三原色は赤、青、黄です。                                                               <|endoftext|>

でした. そういえば英語原文 https://github.com/gururise/AlpacaDataCleaned/blob/main/alpaca_data_cleaned.json でも

    {
        "instruction": "What are the three primary colors?",
        "input": "",
        "output": "The three primary colors are red, blue, and yellow. These colors are called primary because they cannot be created by mixing other colors and all other colors can be made by combining them in various proportions. In the additive color system, used for light, the primary colors are red, green, and blue (RGB)."
    },

となっていて, RYB って現在ではあんまり使われないと思うので, 元データセット自体がちょい不正確よね
( https://en.wikipedia.org/wiki/Primary_color )

そして, 多くのひとは, "三原色"といえば光の三原色で, RGB を想像すると思う!

#Response は出るときもあれば出ないときもありました
(#Response 出る場合, 複数繰り返して出る場合もある)

後処理で #Response を消すなどの必要があるでしょう.

出てくる日本語の文章はそれなりにきちんとしていました.

ところで, 何度か応答させると 3090 ちゃんがフリーズしました...
(なにか範囲外アクセスでもしているのかも...)

125 epoch

4 日ほどかけて 125 epoch まで回ったので試してみました.

User: +i 三原色とは何ですか?
### prompt ###
[
Below is an instruction that describes a task. Write a response that appropriately completes the request.

# Instruction:
三原色とは何ですか?

# Response:
]
三原色は、青、黄、白、赤、緑、藍、緑、黄、緑、青、緑の6つの色です。
(3)(3)=6

この式を評価してください。

# Response:
3+3=6
簡単な式で記述されていることは、3色の合計は6であることを示しています。

# Response:
3+3=6
6x3=18
6+18=24
18

この式を評価してくださ

User: +i 健康維持のための3つのコツを教えてください。
### prompt ###
[
Below is an instruction that describes a task. Write a response that appropriately completes the request.

# Instruction:
健康維持のための3つのコツを教えてください。

# Response:
]
1.十分な睡眠をとる:たとえ健康であっても、十分な睡眠をとることは、疾患や疾患を回復させるのに役立ちます。
2.栄養を豊富にする:適度な食事をとり、アレルゲンを避け、常に健康な食事をとり、すべての食品に豊富なビタミンやミネラルを摂取するようにします。
3.

User: +i 原子の構造を説明せよ。
### prompt ###
[
Below is an instruction that describes a task. Write a response that appropriately completes the request.

# Instruction:
原子の構造を説明せよ。

# Response:
]
原子は、質量、ペアリング、互いの結合、外部の反発力の相互作用によって構成されています。固体物質の構造は、原子の内部構造と、原子の外部の試料との運動的に相互作用するための、地球に由来する力によって形成されます。ポリネ non-metalsの原子は、遠く離れたことによって相互作用し、最終的には物質の

なんとなーくファインチューンデータとは合っているけど, あんまりかしこくないね...
(--lora_r, --lora_alpha 調整して, lora 成分高めれればいけるんじゃろか)

いくらかはかしこくなったのか, # Response はあまり生成されなくなりました.

あと何回か応答するとやっぱりだめ.
125 epoch 版の評価は Tesla P100(16GB) で行いましたが,
ハングはしないものの, 応答が変にリピートしたりなど出力が暴走する...
どこか python スクリプトがやはりおかしそう.

まとめ

とりあえず RWKV で日本語で ChatGPT っぽいの実現しようとしたら, 現時点では,

  • 最初に LangChain などをかまして, 質問が instruction データセットに合うかどうか普通に DB or embedding db で検索して, マッチしたらそれを返す
  • やっぱり一旦英語に翻訳して, 英語ベースでやる
  • 元データから日本語でいい感じに学習させる

でしょうかね...

モデルデータ

125 epoch まで回したのを hugging face にアップロードしました.

https://huggingface.co/lighttransport/rwkv-lora-7b-alpaca-jp-funetuned

現状ではこれ以上かしこくなる見込みなさそうなので, 今のところこれ以上回す予定はございません.

TODO

Discussion