今回は、勉強がてら我々がよく使っているLLMをフルスクラッチで作ってみたいと思います。超今更感はありますがｗ
今回作るものは、実用的に使えるレベルまで達していません。というのも事前学習をさせるためには大量のデータセット（テラバイト単位）と大量のGPUマシンが必要になりますが、そんな大きなものを準備したりチューニングできないので、GoogleColabでもある程度は動くであろう”enwik8”という小さいデータセットを使います。
今回使うフレームワーク
- GPT-NeoX
- GPT-NeoX（ジーピーティー・ネオエックス）は、EleutherAIが開発した言語モデルのトレーニング用のライブラリであり、主にGPUで大規模な言語モデルをトレーニングするためのものです。このライブラリは、NVIDIAのMegatron言語モデルを基にしています。
データセット
- enwiki8
- 英語版Wikipediaから抽出したデータセットになります。軽量です。
GPT-NeoXを使って作ったLLMたち（日本語対応しているものだけ）
-
Rinna Co.'s bilingual-gpt-neox-4b (English / Japanese)
- rinna株式会社が作成したLLMで日本語英語がバランスよく入っています。
-
CyberAgent's Open-CLM (125M through 7B) (Japanese)
- こちらはアベマを手掛ける、CyberAgentが作った日本語データセットを作って学習させたモデルです。
-
The University of Tokyo's weblab-10b and weblab-10b-instruct (Japanese)
- 東京大学松尾研究室で開発された言語モデル
事前準備
GoogleColabの設定
GoogleColabにアクセス後、A100シリーズを選び、ハイメモリモデルを選びます。
これで、GoogleColabのランタイムの設定が完了しました。
wandbのアカウント登録
Weights & Biases - 機械学習開発者のためのコラボレーションプラットフォーム (wandb.jp)
今回はこちらを使って、学習の状況等を見ていきますので、登録します。できるだけGoogleアカウントを使って登録してください。
後々、GoogleColab上で接続するのに便利です。
ノートブックへの記載
GoogleDriveに接続する
下記で、GoogleDriveに接続します。こちらを実行すると、Googleからアクセスを許可するかダイアログで聞いてきますので、許可をクリックしてください。
# Google Colabに接続する
from google.colab import drive
drive.mount('/content/drive')
GPT-NeoXのリポジトリを取得する
下記で、GPT-NeoXのリポジトリを取得します。このリポジトリ内のtrain.pyなどを利用して学習させていきます。
その後、リポジトリ内にパスを遷移させます。
!git clone https://github.com/EleutherAI/gpt-neox.git
%cd "./gpt-neox"
必要ライブラリを導入し設定を同時に行います。
!pip install -r requirements/requirements.txt
!pip install -r requirements/requirements-wandb.txt # optional, if logging using WandB
!pip install -r requirements/requirements-tensorboard.txt # optional, if logging via tensorboard
!python ./megatron/fused_kernels/setup.py install # optional, if using fused kernels
学習させるデータを用意する
enwik8 データセットをダウンロードして GPT2 Tokenizer でトークン化し、dataフォルダ内に保存します。
!python prepare_data.py -d ./data
トークナイザーを変更する場合megatron/neox_arguments/neox_args.py内に
tokenizer_type: Literal[
"GPT2BPETokenizer",
"HFTokenizer",
"HFGPT2Tokenizer",
"SPMTokenizer",
"CharLevelTokenizer",
"TiktokenTokenizer",
] = "GPT2BPETokenizer"
のようになっています。デフォルトではこのトークナイザーが対応していることになります。もし変更する場合は、ここを変更するのと、
megatron/tokenizer/tokenizer.pyを書き換える必要があります。
また、データセットを取得する際、enwikiではないものを取得させたい場合は、prepare.pyの中にある
parser.add_argument(
"dataset",
nargs="?",
default="enwik8",
help="name of dataset to download.",
choices=DATASET_CHOICES,
)
のenwiki8の部分をhuggingfaceに転がっているフォーマットが合致しているデータセットに変更することで、トークナイズ可能です。
wandbライブラリを導入する
下記はを必ず先にインストールします。これがないとwandb接続時にライブラリ不足のエラーが出力されます。
!pip install 'urllib3<2'
肝心のwandbのインストール
!pip install wandb -qU
下記で、wandbにログインしておきます。
上記でGoogleアカウントで登録してほしいと記載していますが、下記のコマンドだけでログインするためです。
Googleアカウント以外だとアクセストークンとか必要だと思います・・・
import wandb
wandb.login()
必要な設定ファイルを作成する
パスはどこでもいいのですが、トレーニングするための設定ファイルを作成します。
今回は永続化できるようにGoogleDrive内に作ることにします。
下記ファイルを作成してください。
下記のように作成すると、あなたのGoogleDriveのフォルダにNEOXというフォルダが作成され、その中に設定ファイルが保存されます。
-
/content/drive/MyDrive/NEOX/config.yml
- 必ずzero_optimizationのstageの値を0に設定します。これをしないと会話モードでエラーが発生します。
# GPT-2 pretraining setup { # parallelism settings ( you will want to change these based on your cluster setup, ideally scheduling pipeline stages # across the node boundaries ) "pipe_parallel_size": 1, "model_parallel_size": 1, # model settings "num_layers": 12, "hidden_size": 768, "num_attention_heads": 12, "seq_length": 2048, "max_position_embeddings": 2048, "norm": "layernorm", "pos_emb": "rotary", "no_weight_tying": true, "gpt_j_residual": false, "output_layer_parallelism": "column", # these should provide some speedup but takes a while to build, set to true if desired "scaled_upper_triang_masked_softmax_fusion": false, "bias_gelu_fusion": false, # init methods "init_method": "small_init", "output_layer_init_method": "wang_init", # optimizer settings "optimizer": { "type": "Adam", "params": { "lr": 0.0006, "betas": [0.9, 0.95], "eps": 1.0e-8, } }, "min_lr": 0.00006, # for all zero_optimization options, see https://www.deepspeed.ai/docs/config-json/#zero-optimizations-for-fp16-training "zero_optimization": { "stage": 0, #⇠超大事 これを設定しないと会話するときにエラーが発生する "allgather_partitions": True, "allgather_bucket_size": 500000000, "overlap_comm": True, "reduce_scatter": True, "reduce_bucket_size": 500000000, "contiguous_gradients": True, }, # batch / data settings "train_micro_batch_size_per_gpu": 4, "data_impl": "mmap", # activation checkpointing "checkpoint_activations": true, "checkpoint_num_layers": 1, "partition_activations": true, "synchronize_each_layer": true, # regularization "gradient_clipping": 1.0, "weight_decay": 0.1, "hidden_dropout": 0.0, "attention_dropout": 0.0, # precision settings "fp16": { "enabled": true, "loss_scale": 0, "loss_scale_window": 1000, "hysteresis": 2, "min_loss_scale": 1 }, # misc. training settings "train_iters": 320000, "lr_decay_iters": 320000, "distributed_backend": "nccl", "lr_decay_style": "cosine", "warmup": 0.01, "checkpoint_factor": 10000, "eval_interval": 1000, "eval_iters": 10, # logging "log_interval": 100, "steps_per_print": 10, "keep_last_n_checkpoints": 4, "wall_clock_breakdown": true, # networking "hostfile": "/mock_path" } -
/content/drive/MyDrive/NEOX/1node.yml
- ここはモデルの保存先やwandbの接続情報が記載されています。
# Suggested data paths when using GPT-NeoX locally { "data_path": "data/enwik8/enwik8_text_document", "vocab_file": "data/gpt2-vocab.json", "merge_file": "data/gpt2-merges.txt", "save": "checkpoints", "load": "checkpoints", "checkpoint_validation_with_forward_pass": False, "tensorboard_dir": "tensorboard", "log_dir": "logs", "use_wandb": True, "wandb_host": "https://api.wandb.ai", "wandb_project": "neox" }
いざ、学習！！
deepspeedをラップしたランチャーがデフォルトで用意されていますので、それを使います。
!python ./deepy.py train.py /content/drive/MyDrive/NEOX/config.yml /content/drive/MyDrive/NEOX/1node.yml
👆こんな感じで、ずっと流れ始めます。
👆wandbの方もこのようにグラフが出てくればOKです！
ずっと流れるので、モデルの保存先であるcheckpointsフォルダになにかでき始めて入れえば、今回はとりあえず体験という感じなので止めてもらっても構いません。
解答精度をあげるんどえあれば、もう少し学習させてもいいかもしれません。👆の設定だと終わるのに数週間くらいかかるので、ある程度のところで・・・
私はとりあえず５時間ほど学習させてみました。
作成したモデルと対話する
では、ある程度学習させたところで、
checkpointフォルダにもファイルができていることが確認できましたので、このモデルデータを使って、モデルと会話してみたいと思います。
text_generation.ymlを編集する
変更するところとしては、temperatureの値と、会話モードであるinteractiveという値をtext_gen_typeに設定します。
- configs/text_generation.yml
# Parameters used for text generation
# Make sure `load` is specified somewhere else
{
# Text gen type: `input-file`, `unconditional` or `interactive`
"text_gen_type": "interactive", #会話モードを指定
# Params for all
"maximum_tokens": 102,
"prompt_end": "\n",
"temperature": 0.0, #より正確に
"top_p": 0.0,
"top_k": 0,
"recompute": false,
# `unconditional`: samples
"num_samples": 10,
# input/output file
"sample_input_file": "sample_input.txt",
"sample_output_file": "sample_output.txt",
}
トレーニング時に設定したファイル２つに今回はtext_generation.ymlを設定します。
!python ./deepy.py generate.py -d configs /content/drive/MyDrive/NEOX/config.yml /content/drive/MyDrive/NEOX/1node.yml text_generation
３０秒くらい待つとプロンプトが入力できるようになります。
今回は、データセットにある内容を聞いてみようと思います。
データセットにはロシア革命が1917年に起きた事が書いてあります。
下記を聞いてみようと思います。
What is the Russian Revolution?
下記のような感じで帰ってきました。
Generated Text:
==The Russian Revolution==
[[Image:Russian_Revolution_1917.jpg|thumb|200px|Russian Revolution of 1917. From left: the Russian Revolution, the Russian Revolution, the Russian Revolution, the Russian Revolution, the Russian Revolution, the Russian Revolution, the Russian Revolution, the Russian Revolution, the Russian Revolution, the Russian Revolution, the Russian Revolution, the Russian Revolution, the Russian Revolution, the Russian Civil War, the Russian Civil War, the Russian Civil War,
Revolution of 1917. まではいいように見えますが、その後が意味不明ですね。。。やっぱりモデルを一からチューニングするにはかなりの時間とコンピュータリソースを必要とするようです。
一旦は、そのリソースが確保できたときにチューンングするメモとして記載しました。
GoogleColabでも複数ノードを束ねて学習できたらいいですね。複数セッションは動かせるのですが、、、、ノードを束ねてというところが難しいのでしょうか。。。