大規模言語モデル(LLM)の作り方 GPT-NeoX編 Part 1
はじめに
Turing 株式会社のリサーチチームでインターンをしている東京工業大学 B4 の藤井(@okoge_kaz)です。
大規模モデルへの注目の高さを肌で感じる今日このごろですが、事前学習の知見については依然として十分に共有されているとは言い難いと個人的に感じています。
Turing株式会社では、次世代の自動運転技術を支える技術の1つとして大規模言語モデルに注目しており、独自に研究開発を行っています。今回は大規模言語モデルを学習する際、用いるライブラリ候補の1つに上がるであろうGPT-NeoXについて解説します。
以下で環境構築方法、学習を行う方法などについて詳しく解説します。
GPT-NeoXとは
EleutherAIが管理しているNIDIA/Megatron-LM ベースの大規模言語モデル(Large Language Model: LLM)を学習するためのライブラリです。
Microsoftが開発するDeepSpeedの技術が組み込まれており、Megatron-DeepSpeedの対抗馬と言える存在です。Megatron-DeepSpeedについて知りたい方は以下の記事をご覧ください。
環境構築
今回実験に使用する環境は A100(40GB) ✕ 16 と、A100(40GB) ✕ 8 です。
1 node 8GPU環境下での実験と、2 node 16 GPU環境下におけるマルチノード実験を行えるように解説を行います。
CUDA11.8環境でのセットアップについて解説しますが、CUDA11.7環境の場合はどのようにすれば良いか適時触れますのでご安心ください。
以下で用いるコードをGitHub上にて公開しています。適時参照ください。
Pythonのセットアップ
Python3.8系を用いましょう。
Python3.9系でも動作するようですが、開発元がPython 3.8で開発、検証していると言っているので、あえて異なるversionにする強い動機がなければPython3.8を用いましょう。
> python --version
Python 3.8.17
> python -m venv .env
> source .env/bin/activate
> which python
/home/kazuki/gpt-neox/.env/bin/python
> python --version
Python 3.8.17
pip install
公式のGPT-NeoXからpip installを行うと、2023年7月16日現在 CUDA11.7対応のPyTorchがinstallされます。
今回は、CUDA11.8環境向けにセットアップを行うのでrequirements/requirements.txt
を修正します。修正済みのものを以下のリンク先に用意していますのでお使いください。
--find-links https://download.pytorch.org/whl/torch_stable.html
torch==2.0.1+cu118
best_download
git+https://github.com/EleutherAI/DeeperSpeed.git#egg=deepspeed
ftfy>=6.0.1
git+https://github.com/EleutherAI/lm_dataformat.git@4eec05349977071bf67fc072290b95e31c8dd836
huggingface_hub>=0.11.0
lm_eval>=0.3.0
mpi4py>=3.0.3
numpy>=1.22.0
pybind11>=2.6.2
regex
sentencepiece
six
tiktoken>=0.1.2
tokenizers>=0.12.1
transformers>=4.24.0
nvcc による versionの確認
Lambda Cloudなど、自分でセットアップしたわけではない環境でCUDA versionを確認する方法
> nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Wed_Sep_21_10:33:58_PDT_2022
Cuda compilation tools, release 11.8, V11.8.89
Build cuda_11.8.r11.8/compiler.31833905_0
では、実際にpip install
していきましょう。
pip install -r requirements/requirements.txt
pip install -r requirements/requirements-wandb.txt
pip install -r requirements/requirements-tensorboard.txt
wandb, tensorboard を導入してログの収集を容易にします。
python ./megatron/fused_kernels/setup.py install
install できると以下のようなログが出るはずです。
Adding fused-kernels 0.0.1 to easy-install.pth file
Installed /home/kazuki/gpt-neox/.env/lib/python3.8/site-packages/fused_kernels-0.0.1-py3.8-linux-x86_64.egg
Processing dependencies for fused-kernels==0.0.1
Finished processing dependencies for fused-kernels==0.0.1
学習準備
データの用意
python prepare_data.py -d ./data
gpt2-vocab.json
, gpt2-merges.txt
, enwik8.zip
がダウンロードされます。
gpt2-vocab.json
, gpt2-merges.txt
は、GPT2BPETokenizerに用いられ、enwik8.zip
はデータセットとして用いられます。
Setting ds_accelerator to cuda (auto detect)
> building GPT2BPETokenizer tokenizer ...
> padded vocab (size: 50257) with 47 dummy tokens (new size: 50304)
Vocab size: 50257
Output prefix: ./data/enwik8/enwik8
> building GPT2BPETokenizer tokenizer ...
> padded vocab (size: 50257) with 47 dummy tokens (new size: 50304)
以上のように自動でdata/enwik8/enwik8_text_document.bin
, data/enwik8/enwik8_text_document.idx
の形式にTokenizeされます。
Tokenizeを行っているので、しばらくの間待つ必要があります。
気長に待ちましょう。
学習
実行方法
GPT-NeoXでは、DeepSpeedをラップしたランチャー(launcher)がデフォルトで用意されています。DeepSpeedコマンドや、mpirun
にて立ち上げることも可能ですが、今回はデフォルトで用意されているランチャーを用いた方法を紹介します。
学習を行う際の実行方法は以下の通りです。
python ./deepy.py train.py [path/to/config1.yml] [path/to/config2.yml] ...
deepy.py
がデフォルトで用意されているランチャーです。学習時には、このランチャーにtrain.py
を渡します。またGPT-NeoXでは、学習時のパラメーターを YAMLファイル形式で管理し、YAMLにて定義した設定ファイルを上記のコマンドのように [path/to/config.yml]
の形で渡すことで学習を行います。
YAMLファイルには、Tensor Parallel, Pipeline ParallelなどのParallelismの設定や、Modelの詳細を決定するための設定、Optimizerの設定などがあります。
具体例を見たほうが理解し易いと思うので、いくつか例を示します。
Parallelism(並列化)設定
"pipe_parallel_size": 1,
"model_parallel_size": 1,
Data Parallel Sizeは、
にて計算されます。
上記で出てきた用語の理解が怪しい方は、以下の記事を参照ください。
モデル 設定
"num_layers": 12,
"hidden_size": 768,
"num_attention_heads": 12,
"seq_length": 2048,
"max_position_embeddings": 2048,
モデルアーキテクチャを指定します。
これによりモデルパラメータ数が決定されます。
Optimizer 設定
"optimizer": {
"type": "Adam",
"params": {
"lr": 0.0006,
"max_grad_norm": 1.0,
"betas": [0.9, 0.95]
}
}
モデルサイズやモデルアーキテクチャに応じて、Learning Rate(学習率)を変更する必要がある場合や、Optimizerのtype自体を変更する場合は、この部分を変更すれば良いです。
その他の設定
DeepSpeed ZeRO Optimizationの設定や、Batch Sizeの設定など、設定できる項目は複数存在します。設定可能なすべての項目一覧はこちらにありますので参照ください。
また、GPT-NeoXでは
python ./deepy.py train.py [path/to/config1.yml] [path/to/config2.yml] ...
のようにconfigファイルを複数渡すことができます。そのため、config1には変更が少ない設定を書き、config2には実験ごとに変える可能性が高いdata_path
, save
(=checkpointの保存先), wandb, tensorboardのlogging設定などを記述するといった運用方法が可能です。
シングルノード
2.7Bのモデルを学習してみましょう。
A100(40GB) ✕ 8枚の環境で学習を行います。
まずscripts/2.7B/1node-dp8-tp1-pp1.sh
を作成し、以下のようにします。
#!/bin/bash
source .env/bin/activate
# nccl log
export NCCL_DEBUG=INFO
# nvlink
nvidia-smi nvlink --status
# log dir
mkdir -p outputs/logs/2.7B
now=$(date +"%Y-%m-%d-%H-%M-%S")
python ./deepy.py train.py \
-d configs 2-7B.yml 1node/2.7b-dp8-tp1-pp1.yml \
&> outputs/logs/2.7B/1node-dp8-tp1-pp1-$now.log
次に configs/1node/2.7b-dp8-tp1-pp1.yml
を作成します。
中身の内容は以下のようにしてください。
{
"data_path": "data/enwik8/enwik8_text_document",
"vocab_file": "data/gpt2-vocab.json",
"merge_file": "data/gpt2-merges.txt",
"save": "checkpoints/default-launcher/gpt-2.7b-dp8-tp1-pp1",
"load": "checkpoints/default-launcher/gpt2-2.7b-dp8-tp1-pp1",
"checkpoint_validation_with_forward_pass": False,
"tensorboard_dir": "tensorboard",
"log_dir": "logs",
"use_wandb": True,
"wandb_host": "https://api.wandb.ai",
"wandb_project": "gpt-neox-abci",
"wandb_team": <user-name>,
"wandb_group": "gpt-2.7b-dp8-tp1-pp1",
}
諸々の準備が整いましたので、学習を行いましょう。
gpt-neox/
にて
bash scripts/2.7B/1node-dp8-tp1-pp1.sh
で学習を行いましょう。
上手く学習ができると、wandbに以下のようなグラフが得られます。
また以下のような標準出力が得られます。
標準出力
[2023-07-07 19:37:06,263] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 573.87 | optimizer_gradients: 5.38 | optimizer_step: 9.64
[2023-07-07 19:37:10,624] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 578.54 | optimizer_gradients: 5.31 | optimizer_step: 9.60
[2023-07-07 19:37:14,496] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 577.69 | optimizer_gradients: 5.40 | optimizer_step: 9.65
[2023-07-07 19:37:18,364] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 574.50 | optimizer_gradients: 5.34 | optimizer_step: 9.60
[2023-07-07 19:37:22,235] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 577.48 | optimizer_gradients: 5.52 | optimizer_step: 9.65
[2023-07-07 19:37:26,113] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 577.73 | optimizer_gradients: 5.37 | optimizer_step: 9.67
[2023-07-07 19:37:29,978] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 575.30 | optimizer_gradients: 5.36 | optimizer_step: 9.72
[2023-07-07 19:37:33,855] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 575.87 | optimizer_gradients: 5.44 | optimizer_step: 9.62
[2023-07-07 19:37:37,733] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 576.62 | optimizer_gradients: 5.30 | optimizer_step: 9.96
[2023-07-07 19:37:41,604] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 577.08 | optimizer_gradients: 5.33 | optimizer_step: 9.61
[2023-07-07 19:37:41,604] [INFO] [logging.py:96:log_dist] [Rank 0] step=10, skipped=0, lr=[5e-07, 5e-07], mom=[[0.9, 0.95], [0.9, 0.95]]
[2023-07-07 19:37:41,606] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | batch_input: 523.88 | forward_microstep: 6684.31 | backward_microstep: 15795.31 | backward_inner_microstep: 15794.06 | backward_allreduce_microstep: 0.24 | step_microstep: 6070.90
[2023-07-07 19:37:41,607] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | forward: 6684.10 | backward: 15795.29 | backward_inner: 15793.91 | backward_allreduce: 0.28 | step: 6071.11
steps: 10 loss: 9.5540 iter time (s): 4.179 samples/sec: 7.658
[2023-07-07 19:37:41,608] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms)
[2023-07-07 19:37:45,490] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 577.16 | optimizer_gradients: 5.30 | optimizer_step: 9.64
[2023-07-07 19:37:49,359] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 574.45 | optimizer_gradients: 5.35 | optimizer_step: 9.66
[2023-07-07 19:37:53,229] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 573.38 | optimizer_gradients: 5.30 | optimizer_step: 9.85
[2023-07-07 19:37:57,108] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 578.81 | optimizer_gradients: 5.34 | optimizer_step: 9.61
[2023-07-07 19:38:00,988] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 575.95 | optimizer_gradients: 5.40 | optimizer_step: 9.60
[2023-07-07 19:38:04,863] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 572.32 | optimizer_gradients: 5.56 | optimizer_step: 9.63
[2023-07-07 19:38:08,744] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 577.78 | optimizer_gradients: 5.42 | optimizer_step: 9.68
[2023-07-07 19:38:12,661] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 575.15 | optimizer_gradients: 5.35 | optimizer_step: 9.63
[2023-07-07 19:38:16,552] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 578.07 | optimizer_gradients: 5.39 | optimizer_step: 9.62
[2023-07-07 19:38:20,430] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 577.15 | optimizer_gradients: 5.37 | optimizer_step: 9.65
[2023-07-07 19:38:20,431] [INFO] [logging.py:96:log_dist] [Rank 0] step=20, skipped=0, lr=[1e-06, 1e-06], mom=[[0.9, 0.95], [0.9, 0.95]]
[2023-07-07 19:38:20,433] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | batch_input: 12.32 | forward_microstep: 4644.11 | backward_microstep: 15769.43 | backward_inner_microstep: 15767.64 | backward_allreduce_microstep: 0.44 | step_microstep: 6064.11
[2023-07-07 19:38:20,434] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | forward: 4643.94 | backward: 15769.51 | backward_inner: 15767.57 | backward_allreduce: 0.48 | step: 6064.31
マルチノード
シングルノードで学習ができたので、次はマルチノードでの学習を行ってみましょう。
セットアップ
マルチノードで学習するためには、そのためのセットアップが必要です。
DeepSpeedをlauncherとして用いる場合、マルチノードで学習する際にはPDSHを用います。
該当コード
そのため、pdshをinstallする必要があります。インストール方法は以下の通りです。
sudo apt-get update -y
sudo apt-get install -y pdsh
pdsh が install できていないと以下のようなエラーが発生します。
2023-07-12 17:12:39,692] [INFO] [runner.py:455:main] Using IP address of 10.xx.xx.xxx for node 10.xx.xx.xxx
Traceback (most recent call last):
File "./deepy.py", line 41, in <module>
main()
File "./deepy.py", line 37, in main
deepspeed.launcher.runner.main(deepspeed_main_args)
File "/home/kazuki/gpt-neox/.env/lib/python3.8/site-packages/deepspeed/launcher/runner.py", line 527, in main
raise RuntimeError(f"launcher '{args.launcher}' not installed.")
RuntimeError: launcher 'pdsh' not installed.
次に、DeepSpeedでマルチノード学習を行うには、ノード間でパスフレーズなしSSHが可能な必要があります。(より正確にはPDSHの仕様)
ssh-keygen -t ed25519 -f <key-name>
でSSHキーを作成し、ノード間で接続できるようにします。まず、接続先のノードに何らかの方法で公開鍵を登録します。(webの管理画面等から公開鍵をuploadできるようなサービスもあれば、接続先の~/.ssh/authorized_keys
に何らかの方法で公開鍵を追加する必要がある場合もあります。)
次にssh configを作成しましょう。以下のような形で作成します。
Host 10.xx.xx.xxx
HostName 10.xx.xx.xxx
User kazuki
IdentityFile ~/.ssh/<key-name>
ServerAliveInterval 15
また、マルチノードで学習を行うには hostfileを作成する必要があります。
以下のような要領で、Host名、GPU数を記述してください。
(今回は A100(40GB)✕8GPU ✕ 2nodeなので、slots=8としています)
10.xx.xx.xxx slots=8
10.xx.xx.yyy slots=8
学習
準備が整ったので、2 node 16 GPUでも学習させてみましょう。
#!/bin/bash
source .env/bin/activate
# nccl log
export NCCL_DEBUG=INFO
# nvlink
nvidia-smi nvlink --status
# log dir
mkdir -p outputs/logs/2.7B
now=$(date +"%Y-%m-%d-%H-%M-%S")
python ./deepy.py train.py \
--hostfile scripts/2.7B/hostfile \
-d configs 2-7B.yml multi-node/2.7B/dp16-tp1-pp1.yml \
&> outputs/logs/2.7B/2node-dp16-tp1-pp1-$now.log
のような学習用のJob Scriptを作成します。
次に、configファイルを作成します。
{
"data_path": "data/enwik8/enwik8_text_document",
"vocab_file": "data/gpt2-vocab.json",
"merge_file": "data/gpt2-merges.txt",
"save": "checkpoints/gpt-2.7b-dp16-tp1-pp1",
"load": "checkpoints/gpt-2.7b-dp16-tp1-pp1",
"checkpoint_validation_with_forward_pass": False,
"tensorboard_dir": "tensorboard",
"log_dir": "logs",
"use_wandb": True,
"wandb_host": "https://api.wandb.ai",
"wandb_project": "gpt-neox-abci",
"wandb_team": <user-name>,
"wandb_group": "gpt-neox-2.7b-dp16",
}
実は、1 node の場合と変わっているのはcheckpointの保存先、ロード先、wandb_groupの名前だけです。どうして、このようなことが可能かというと、batch sizeについてconfigファイルにて指定しているのが "train_micro_batch_size_per_gpu": 4
であるため、global batch sizeをあわせて変更する必要がないからです。また、Data Parallel Size(データ並列数)はWorld SizeとTensor Parallel Size, Pipeline Parallel Sizeから自動計算されるため、こちら側でconfigファイルを書き換える必要はありません。
準備が整ったので実際に実行してみましょう。
bash scripts/2.7B/2node-dp16-tp1-pp1.sh
で実行します。
実行する際に、他のノードのGPUが確認できているか、確認するには、実行時の標準出力を見てください
[2023-07-17 13:24:36,090] [INFO] [launch.py:162:main] global_rank_mapping=defaultdict(<class 'list'>, {'10.xx.xx.xxx': [0, 1, 2, 3, 4, 5, 6, 7], '10.xx.xx.yyy': [8, 9, 10, 11, 12, 13, 14, 15]})
10.xx.xx.xxx: [2023-07-17 13:24:36,090] [INFO] [launch.py:163:main] dist_world_size=16
上手く実行できている場合は、wandbにLossの推移が以下のように記録されているはずです。
マルチノードでも学習を行うことができました。🎉
標準出力は以下のようになります。
標準出力
[2023-07-17 13:27:17,846] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 789.93 | optimizer_gradients: 2.91 | optimizer_step: 4.84
10.xx.xx.xxx: [2023-07-17 13:27:22,223] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 787.71 | optimizer_gradients: 2.91 | optimizer_step: 4.83
10.xx.xx.xxx: [2023-07-17 13:27:26,602] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 788.85 | optimizer_gradients: 2.87 | optimizer_step: 4.81
10.xx.xx.xxx: [2023-07-17 13:27:30,980] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 787.38 | optimizer_gradients: 2.97 | optimizer_step: 4.87
10.xx.xxx.xxx: [2023-07-17 13:27:35,351] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 786.62 | optimizer_gradients: 2.89 | optimizer_step: 4.88
10.xx.xx.xxx: [2023-07-17 13:27:39,732] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 791.11 | optimizer_gradients: 2.99 | optimizer_step: 4.84
10.xx.xx.xxx: [2023-07-17 13:27:44,110] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 790.53 | optimizer_gradients: 2.89 | optimizer_step: 4.85
10.xx.xx.xxx: [2023-07-17 13:27:48,484] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 787.16 | optimizer_gradients: 2.87 | optimizer_step: 4.85
10.xx.xx.xxx: [2023-07-17 13:27:52,859] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 787.22 | optimizer_gradients: 2.91 | optimizer_step: 4.84
10.xx.xx.xxx: [2023-07-17 13:27:52,860] [INFO] [logging.py:96:log_dist] [Rank 0] step=10, skipped=0, lr=[5e-07, 5e-07], mom=[[0.9, 0.95], [0.9, 0.95]]
10.xx.xx.xxx: [2023-07-17 13:27:52,861] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | batch_input: 1438.54 | forward_microstep: 13076.98 | backward_microstep: 16110.56 | backward_inner_microstep: 16109.84 | backward_allreduce_microstep: 0.08 | step_microstep: 8128.62
10.xx.xx.xxx: [2023-07-17 13:27:52,862] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | forward: 13076.75 | backward: 16110.41 | backward_inner: 16109.70 | backward_allreduce: 0.09 | step: 8128.77
10.xx.xx.xxx: steps: 10 loss: 9.5307 iter time (s): 5.483 samples/sec: 11.672
10.xx.xx.xxx: [2023-07-17 13:27:52,863] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms)
10.xx.xxx.xxx: [2023-07-17 13:27:57,236] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 788.70 | optimizer_gradients: 2.88 | optimizer_step: 4.84
10.xx.xx.xxx: [2023-07-17 13:28:01,611] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 787.81 | optimizer_gradients: 2.96 | optimizer_step: 4.84
10.xx.xx.xxx: [2023-07-17 13:28:05,988] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 790.09 | optimizer_gradients: 2.87 | optimizer_step: 4.85
10.xx.xx.xxx: [2023-07-17 13:28:10,362] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 788.21 | optimizer_gradients: 2.87 | optimizer_step: 4.84
10.xx.xx.xxx: [2023-07-17 13:28:14,741] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 790.06 | optimizer_gradients: 2.88 | optimizer_step: 4.83
10.xx.xx.xxx: [2023-07-17 13:28:19,116] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 789.71 | optimizer_gradients: 2.87 | optimizer_step: 4.83
10.xx.xx.xxx: [2023-07-17 13:28:23,492] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 789.44 | optimizer_gradients: 2.87 | optimizer_step: 4.82
10.xx.xx.xxx: [2023-07-17 13:28:27,873] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 790.56 | optimizer_gradients: 2.90 | optimizer_step: 4.84
10.xx.xx.xxx: [2023-07-17 13:28:32,247] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 789.19 | optimizer_gradients: 2.87 | optimizer_step: 4.83
10.xx.xx.xxx: [2023-07-17 13:28:36,626] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 787.22 | optimizer_gradients: 2.89 | optimizer_step: 4.84
10.xx.xx.xxx: [2023-07-17 13:28:36,627] [INFO] [logging.py:96:log_dist] [Rank 0] step=20, skipped=0, lr=[1e-06, 1e-06], mom=[[0.9, 0.95], [0.9, 0.95]]
10.xx.xx.xxx: [2023-07-17 13:28:36,628] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | batch_input: 12.21 | forward_microstep: 4588.99 | backward_microstep: 15634.87 | backward_inner_microstep: 15634.26 | backward_allreduce_microstep: 0.07 | step_microstep: 8043.15
10.xx.xx.xxx: [2023-07-17 13:28:36,629] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | forward: 4588.86 | backward: 15634.69 | backward_inner: 15634.14 | backward_allreduce: 0.07 | step: 8043.33
10.xx.xx.xxx: steps: 20 loss: 8.4760 iter time (s): 4.374 samples/sec: 14.633
ZeRO Data Parallel + Pipeline Parallel
Data Parallel単体だけでなく、Pipeline ParallelとData Parallelを組み合わせる手法についても試してみましょう。
scripts/2.7B/1node-dp4-tp1-pp2.sh
に以下のようなスクリプトファイルを作成します。
#!/bin/bash
source .env/bin/activate
# nccl log
export NCCL_DEBUG=INFO
# nvlink
nvidia-smi nvlink --status
# log dir
mkdir -p outputs/logs/2.7B
now=$(date +"%Y-%m-%d-%H-%M-%S")
python ./deepy.py train.py \
-d configs 1node/2.7b-dp4-tp1-pp2.yml \
&> outputs/logs/2.7B/1node-dp4-tp1-pp2-$now.log
加えて config ファイルを作成します。
今回は複数ファイルに分けず1つのファイルに集約してみましょう。
configs/1node/2.7b-dp4-tp1-pp2.yml
に以下のようなYAMLファイルを作成します。
{
# parallelism settings ( you will want to change these based on your cluster setup, ideally scheduling pipeline stages
# across the node boundaries )
"pipe_parallel_size": 2,
"model_parallel_size": 1,
# model settings
"num_layers": 32,
"hidden_size": 2560,
"num_attention_heads": 32,
"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.00016, "betas": [0.9, 0.95], "eps": 1.0e-8 },
},
"min_lr": 0.000016,
# for all zero_optimization options, see https://www.deepspeed.ai/docs/config-json/#zero-optimizations-for-fp16-training
"zero_optimization":
{
"stage": 1,
"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,
"attention_dropout": 0,
# precision settings
"fp16":
{
"fp16": true,
"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,
"save_iters": 5000,
# logging
"log_interval": 1,
"steps_per_print": 10,
"keep_last_n_checkpoints": 4,
"wall_clock_breakdown": true,
"data_path": "data/enwik8/enwik8_text_document",
"vocab_file": "data/gpt2-vocab.json",
"merge_file": "data/gpt2-merges.txt",
# dataset settings
"save": "checkpoints/default-launcher/gpt-2.7b-dp4-tp1-pp2",
"load": "checkpoints/default-launcher/gpt-2.7b-dp4-tp1-pp2",
"checkpoint_validation_with_forward_pass": False,
"tensorboard_dir": "tensorboard",
"log_dir": "logs",
"use_wandb": True,
"wandb_host": "https://api.wandb.ai",
"wandb_project": "gpt-neox-abci",
"wandb_team": <user-name>,
"wandb_group": "gpt-neox-2.7b-dp4-tp1-pp2",
}
wandb_team
の名前は置き換えてください。
では実際に動かしてみましょう。
wandbでLossの推移を確認すると次のようになるはずです。
またPP=2 としたことで1GPUあたりのAllocated Bytesが減っていることも確認できます。
また標準出力は以下のようになります。
標準出力
samples/sec: 2.598 | iteration 1/ 320000 | elapsed time per iteration (ms): 6159.7 | learning rate: 5.000E-08 | approx flops per GPU: 17.8TFLOPS | lm_loss: 1.107905E+01 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
after 1 iterations memory (MB) | allocated: 6687.5341796875 | max allocated: 13676.14990234375 | reserved: 16750.0 | max reserved: 16750.0
time (ms)
[2023-07-17 15:02:21,036] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 239.21 | optimizer_gradients: 5.31 | optimizer_step: 9.61
samples/sec: 5.694 | iteration 2/ 320000 | elapsed time per iteration (ms): 2809.9 | learning rate: 1.000E-07 | approx flops per GPU: 38.9TFLOPS | lm_loss: 1.107772E+01 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:02:23,831] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 237.42 | optimizer_gradients: 5.30 | optimizer_step: 9.67
samples/sec: 5.730 | iteration 3/ 320000 | elapsed time per iteration (ms): 2792.4 | learning rate: 1.500E-07 | approx flops per GPU: 39.2TFLOPS | lm_loss: 1.105253E+01 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:02:26,635] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 240.49 | optimizer_gradients: 5.32 | optimizer_step: 9.65
samples/sec: 5.705 | iteration 4/ 320000 | elapsed time per iteration (ms): 2804.5 | learning rate: 2.000E-07 | approx flops per GPU: 39.0TFLOPS | lm_loss: 1.104417E+01 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:02:29,430] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 242.47 | optimizer_gradients: 5.32 | optimizer_step: 9.66
samples/sec: 5.731 | iteration 5/ 320000 | elapsed time per iteration (ms): 2791.6 | learning rate: 2.500E-07 | approx flops per GPU: 39.2TFLOPS | lm_loss: 1.087381E+01 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:02:32,232] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 242.70 | optimizer_gradients: 5.31 | optimizer_step: 9.62
samples/sec: 5.708 | iteration 6/ 320000 | elapsed time per iteration (ms): 2803.1 | learning rate: 3.000E-07 | approx flops per GPU: 39.0TFLOPS | lm_loss: 1.079356E+01 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:02:35,023] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 239.13 | optimizer_gradients: 5.31 | optimizer_step: 9.67
samples/sec: 5.724 | iteration 7/ 320000 | elapsed time per iteration (ms): 2795.1 | learning rate: 3.500E-07 | approx flops per GPU: 39.2TFLOPS | lm_loss: 1.042607E+01 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:02:37,835] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 243.55 | optimizer_gradients: 5.32 | optimizer_step: 9.66
samples/sec: 5.701 | iteration 8/ 320000 | elapsed time per iteration (ms): 2806.7 | learning rate: 4.000E-07 | approx flops per GPU: 39.0TFLOPS | lm_loss: 1.014001E+01 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:02:40,622] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 237.42 | optimizer_gradients: 5.32 | optimizer_step: 9.63
samples/sec: 5.727 | iteration 9/ 320000 | elapsed time per iteration (ms): 2793.8 | learning rate: 4.500E-07 | approx flops per GPU: 39.2TFLOPS | lm_loss: 9.925351E+00 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:02:43,429] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 238.75 | optimizer_gradients: 5.31 | optimizer_step: 9.68
[2023-07-17 15:02:43,429] [INFO] [logging.py:96:log_dist] [Rank 0] step=10, skipped=0, lr=[5e-07, 5e-07], mom=[[0.9, 0.95], [0.9, 0.95]]
[2023-07-17 15:02:43,431] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | batch_input: 77.71 | forward_microstep: 3775.43 | backward_microstep: 7808.20 | backward_inner_microstep: 7808.06 | backward_allreduce_microstep: 0.01 | step_microstep: 2695.14
[2023-07-17 15:02:43,431] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | forward: 3775.37 | backward: 7808.12 | backward_inner: 7807.97 | backward_allreduce: 0.01 | step: 2695.28
steps: 10 loss: 9.7646 iter time (s): 3.134 samples/sec: 5.105
[2023-07-17 15:02:43,443] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | pipe_send_output: 36.25 | pipe_recv_grad: 11829.17
samples/sec: 5.688 | iteration 10/ 320000 | elapsed time per iteration (ms): 2812.9 | learning rate: 5.000E-07 | approx flops per GPU: 38.9TFLOPS | lm_loss: 9.764610E+00 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:02:46,234] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 243.03 | optimizer_gradients: 5.31 | optimizer_step: 9.68
samples/sec: 5.730 | iteration 11/ 320000 | elapsed time per iteration (ms): 2792.3 | learning rate: 5.500E-07 | approx flops per GPU: 39.2TFLOPS | lm_loss: 9.616756E+00 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:02:49,033] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 239.08 | optimizer_gradients: 5.32 | optimizer_step: 9.67
samples/sec: 5.707 | iteration 12/ 320000 | elapsed time per iteration (ms): 2803.4 | learning rate: 6.000E-07 | approx flops per GPU: 39.0TFLOPS | lm_loss: 9.390484E+00 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:02:51,828] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 243.08 | optimizer_gradients: 5.31 | optimizer_step: 9.64
samples/sec: 5.734 | iteration 13/ 320000 | elapsed time per iteration (ms): 2790.4 | learning rate: 6.500E-07 | approx flops per GPU: 39.2TFLOPS | lm_loss: 9.173844E+00 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:02:54,620] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 238.17 | optimizer_gradients: 5.33 | optimizer_step: 9.62
samples/sec: 5.719 | iteration 14/ 320000 | elapsed time per iteration (ms): 2797.9 | learning rate: 7.000E-07 | approx flops per GPU: 39.1TFLOPS | lm_loss: 9.301665E+00 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:02:57,413] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 238.70 | optimizer_gradients: 5.31 | optimizer_step: 9.65
samples/sec: 5.729 | iteration 15/ 320000 | elapsed time per iteration (ms): 2792.8 | learning rate: 7.500E-07 | approx flops per GPU: 39.2TFLOPS | lm_loss: 9.087448E+00 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:03:00,210] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 238.83 | optimizer_gradients: 5.32 | optimizer_step: 9.62
samples/sec: 5.721 | iteration 16/ 320000 | elapsed time per iteration (ms): 2796.6 | learning rate: 8.000E-07 | approx flops per GPU: 39.1TFLOPS | lm_loss: 8.934504E+00 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:03:03,006] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 242.85 | optimizer_gradients: 5.31 | optimizer_step: 9.64
samples/sec: 5.736 | iteration 17/ 320000 | elapsed time per iteration (ms): 2789.6 | learning rate: 8.500E-07 | approx flops per GPU: 39.2TFLOPS | lm_loss: 9.194756E+00 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:03:05,804] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 242.41 | optimizer_gradients: 5.31 | optimizer_step: 9.62
samples/sec: 5.716 | iteration 18/ 320000 | elapsed time per iteration (ms): 2799.3 | learning rate: 9.000E-07 | approx flops per GPU: 39.1TFLOPS | lm_loss: 8.938364E+00 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:03:08,594] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 242.51 | optimizer_gradients: 5.31 | optimizer_step: 9.63
samples/sec: 5.733 | iteration 19/ 320000 | elapsed time per iteration (ms): 2790.7 | learning rate: 9.500E-07 | approx flops per GPU: 39.2TFLOPS | lm_loss: 8.941507E+00 | loss scale: 65536.0 | number of skipped iterations: 0 | number of nan iterations: 0 |
time (ms)
[2023-07-17 15:03:11,395] [INFO] [logging.py:96:log_dist] [Rank 0] rank=0 time (ms) | optimizer_allgather: 242.53 | optimizer_gradients: 5.31 | optimizer_step: 9.62
つづく
この記事ではGPT-NeoXを用いて事前学習を行う具体的な方法について説明しました。
続編「大規模言語モデル(LLM)の作り方 GPT-NeoX編 Part 2」では、Flash Attentionを組み合わせた使用方法などGPT-NeoXについてさらに詳しく掘り下げます。
今回の記事で分散並列学習に興味が湧いた方は、「大規模モデルを支える分散並列学習のしくみ Part 1」 をご覧ください。
また今回の記事で紹介した学習にはZeRO Stage1というメモリ削減技術が使われています。この仕組みについては 「大規模モデルを支える分散並列学習のしくみ Part 2」 にて解説する予定です。
Turing では自動運転モデルの学習や、自動運転を支えるための基盤モデルの作成のために分散並列学習の知見を取り入れた研究開発を行っています。興味がある方は、Turing の公式 Web サイト、採用情報などをご覧ください。話を聞きたいという方は私や AI チームのディレクターの山口さんの Twitter DM からでもお気軽にご連絡ください
Discussion
Amazing!
大変勉強になりました!