🎋

大規模言語モデル(LLM)の作り方 GPT-NeoX編 Part 1

2023/07/19に公開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について知りたい方は以下の記事をご覧ください。

https://zenn.dev/turing_motors/articles/04c1328bf6095a

環境構築

今回実験に使用する環境は A100(40GB) ✕ 16 と、A100(40GB) ✕ 8 です。
1 node 8GPU環境下での実験と、2 node 16 GPU環境下におけるマルチノード実験を行えるように解説を行います。

CUDA11.8環境でのセットアップについて解説しますが、CUDA11.7環境の場合はどのようにすれば良いか適時触れますのでご安心ください。

以下で用いるコードをGitHub上にて公開しています。適時参照ください。

https://github.com/okoge-kaz/gpt-neox-abci-llm-hackathon

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を修正します。修正済みのものを以下のリンク先に用意していますのでお使いください。

https://github.com/okoge-kaz/gpt-neox-abci-llm-hackathon/blob/hpc/fujii/wandb-logging/requirements/requirements.txt

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は、

\text{DP-SIZE} = \text{WORLD-SIZE} / (\text{TP-SIZE} \times \text{PP-SIZE})

にて計算されます。

上記で出てきた用語の理解が怪しい方は、以下の記事を参照ください。

https://zenn.dev/turing_motors/articles/0e6e2baf72ebbc

モデル 設定

  "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を作成し、以下のようにします。

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を作成します。

中身の内容は以下のようにしてください。

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を用います。

該当コード
https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/launcher/runner.py#L120-L124

そのため、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を作成しましょう。以下のような形で作成します。

~/.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としています)

scripts/2.7B/hostfile
10.xx.xx.xxx slots=8
10.xx.xx.yyy slots=8

学習

準備が整ったので、2 node 16 GPUでも学習させてみましょう。

scripts/2.7B/2node-dp16-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 \
  --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ファイルを作成します。

configs/multi-node/2.7B/dp16-tp1-pp1.yml
{
  "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に以下のようなスクリプトファイルを作成します。

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 からでもお気軽にご連絡ください

Tech Blog - Turing

Discussion