GPT-neoxの学習用にマルチノード並列学習環境を整えた with DeepSpeed

1. はじめに

こんにちは。ABEJAの村主です。社内で「ABEJAで作った大規模GPTモデルとその道のり」なプロジェクトが立ち上がりました。面白そうだったので首を突っ込んでみました。そしていくつかハマりながら実装していったので、忘備録と共に記録していきます。今回はマルチノード並列学習部分の話になりますので、実際のGPT関連の中身や創意工夫の話等は ABEJAで作った大規模GPTモデルとその道のり や第三弾の記事を参照ください。今回の取り組みでは、GCPのクラウド環境とNVIDIA社のGPUを利活用し環境構築を行いました。

今回のプロジェクトは御多分に漏れず期限が決まっており「期限までに学習を終わらせられるように」というのが一つの要件でした。学習速度を上げるための並列学習の方法としては、1つのノード（サーバ、インスタンスと同義）で複数のGPUを扱う マルチGPUと、複数のノードで複数のGPUを扱う マルチノード という方法があります。今回は1つのノードで複数のGPUを使うだけだと間に合わず、マルチノードの並列学習を行いたい ということで実装方法を調べてみました。1つのノードで複数のGPUを扱うこと自体は難しくないのですが、複数のノードで学習を行うにはマルチノード自体の環境構築やノード間のネットワークなどボトルネックになるポイントが多々あります。GCP環境下でどう実現できるのか調べ始めました。

2. 並列学習環境を調べる

並列学習方法を調べる

今回利用しているのは gpt-neox というものです。詳しい話は ABEJAで作った大規模GPTモデルとその道のり」 を参照ください。

github.com

その中で DeepSpeed というOSSを使ってマルチノードの並列学習を実現できるみたいです。

DeepSpeed configures multi-node compute resources with hostfiles that are compatible with OpenMPI  and Horovod. A hostfile is a list of hostnames (or SSH aliases), which are machines accessible via passwordless SSH, and slot counts, which specify the number of GPUs available on the system.

https://www.deepspeed.ai/getting-started/#resource-configuration-multi-node

なるほど。OpenMPI等と同様にSSHを介して通信するみたいですね。以前horovodも実装したことあるし従来の手法なのでなんとかなる。ドキュメント読んでる限りインフラ環境は ノード間をパスワードなしのSSHが出来る ようにして hostfiles を配置する だけで出来るそうな。そうだとしたらちょっとハマったとしてもすぐに実装できそうな感触を得ました。

ネットワーク、コンピューティング周りを調べる

マルチGPUなどの並列学習時は自サーバ内で完結するためCPU・メモリ・GPUの通信は高速にやりとり可能です。しかしノード（サーバ・インスタンス）間の通信はCPU・メモリ・GPUと比べると非常に非常に遅く、マルチノードの並列学習はノード間の通信がボトルネックになりがちです。今回GCP環境を使うのですが、少しでもボトルネックにならない方法が無いか調べてみました。

cloud.google.com

今回はじめに実装したのは以下になります。

1. Google Virtual NIC（gVNIC）の有効化

2. Fast Socket でより高速のネットワーク帯域幅を使用する

   Fast Socket は NCCL 用の Google 独自のネットワーク トランスポートで、Compute Engine で Fast Socket を使用すると、複数の TCP 接続間の競合が減り、100 Gbps ネットワーク上の NCCL パフォーマンスが向上します。

   *NCCL（NVIDIA Collective Communications Library）は、TensorFlow、PyTorch、Holoods などのディープ ラーニング フレームワークでマルチ GPU やマルチノード トレーニングを行う場合に使用されます。

3. コンパクト プレースメント ポリシーを使用する

   プレースメント ポリシーでは、データセンターでの仮想マシン（VM）の配置を制御できます。コンパクト プレースメント ポリシーでは、単一のアベイラビリティ ゾーンに VM を配置するために、低レイテンシのトポロジが用意されています

   なお、これは後述しますが最後に外すことになりました

4. Deep Learning VM Image の使用

   Fast Socket を設定するには、Deep Learning VM を使用します。Deep Learning VM イメージには、GPU ドライバ、ML ソフトウェア、Fast Socket、gVNIC がプリインストールされています。

   これは各種プリセットされているので便利ですね。ありがたく使わせてもらいます。

3. インフラ環境を構築する

なんとなくイメージがついてきたので、とりあえず作って動かすことにします。

コンパクトプレースメントポリシーの作成

まず初めにコンパクトプレースメントポリシーを作成します。コンパクトプレースメントポリシーとは以下の通りです。

コンパクト プレースメント ポリシーは、単一のアベイラビリティ ゾーンに VM を配置する低レイテンシ トポロジを提供し、ネットワーク内の最寄りのノードでインスタンスをホストするのが特長です。

要約すると「アベイラビリティゾーンの中で、ネットワーク的に近いノードにVMを配置してくれる機能」です。基本的には近い方がネットワークレイテンシーが小さくなりますしね。

cloud.google.com

リージョンは使用したいGPUがある場所ならどこでも良いですが、基本的には us-central1 が GPU リソースが潤沢かなと思います。ゾーンの中にも潤沢なゾーンがありますので詳細はサポートに問い合わせください。作成する際のコマンドは以下になります。

$ PROJECT="xxx"
$ REGION="us-central1"

$ gcloud compute resource-policies create group-placement a100-${REGION} \
    --collocation COLLOCATED \
    --region ${REGION} \
    --project ${PROJECT}

Compute Engine を起動する (Fast Socket と gVNIC を利用する)

ここのサンプル を参考にコマンドを作成します。

基本は見ての通りですが、 Deep Learning VM Image を指定したり、gVNIC を指定したりしています。

Deep Learning VM Image を利用することでプリインストールされていて、実装がすごく簡単でありがたいですね。

コマンドはこちらになります。A2インスタンスを指定しているのでめちゃお金かかります。気をつけてください。

$ PROJECT="xxx"
$ REGION="us-central1"
$ ZONE="us-central1-f"
$ TYPE="a2-highgpu-8g"

$ gcloud compute instances create a100-1 \
    --project=${PROJECT} \
    --zone=${ZONE} \
    --machine-type=${TYPE} \
    --resource-policies a100-${REGION} \
    --maintenance-policy=TERMINATE \
    --no-restart-on-failure \
    --image-family=tf-latest-gpu-ubuntu-2004 \
    --image-project=deeplearning-platform-release \
    --boot-disk-size=2TiB \
    --network-interface=nic-type=GVNIC \
    --metadata="install-nvidia-driver=True,proxy-mode=project_editors" \
    --scopes=https://www.googleapis.com/auth/cloud-platform

4. まずはシングルノードで動かす

インストールは簡単ですね

$ sudo apt install -y libaio1 libaio-dev
$ pip install deepspeed

動作チェック

$ ds_report
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
sparse_attn ............ [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
async_io ............... [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/opt/conda/lib/python3.7/site-packages/torch']
torch version .................... 1.11.0+cu102
torch cuda version ............... 10.2
nvcc version ..................... 11.0
deepspeed install path ........... ['/opt/conda/lib/python3.7/site-packages/deepspeed']
deepspeed info ................... 0.5.10, unknown, unknown
deepspeed wheel compiled w. ...... torch 0.0, cuda 0.0

社内のコードを動かしてみる。とりあえず動作確認のために docker に入って叩く

$ git clone https://github.com/SO0529/gpt-neox.git
$ cd gpt-neox

$ docker build -t gpt-neox -f Dockerfile .
$ docker compose -f docker-compose.yml -p gpt-neox up -d
$ docker exec -it gpt-neox-gpt-neox-1 bash

# Docker内の操作
$ python /gpt-neox/megatron/fused_kernels/setup.py install
$ python prepare_data.py -d ./data

$ python ./deepy.py train.py -d configs small.yml local_setup.yml

横で使用状況を確認する。よしよし。ちゃんとフルで使えてるな

$ nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.103.01   Driver Version: 470.103.01   CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| 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 A100-SXM...  Off  | 00000000:00:04.0 Off |                    0 |
| N/A   58C    P0   379W / 400W |  10482MiB / 40536MiB |     97%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA A100-SXM...  Off  | 00000000:00:05.0 Off |                    0 |
| N/A   54C    P0   284W / 400W |  10914MiB / 40536MiB |     92%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   2  NVIDIA A100-SXM...  Off  | 00000000:00:06.0 Off |                    0 |
| N/A   54C    P0   309W / 400W |  10914MiB / 40536MiB |     94%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   3  NVIDIA A100-SXM...  Off  | 00000000:00:07.0 Off |                    0 |
| N/A   57C    P0   273W / 400W |  10914MiB / 40536MiB |     97%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   4  NVIDIA A100-SXM...  Off  | 00000000:80:00.0 Off |                    0 |
| N/A   53C    P0   180W / 400W |  10914MiB / 40536MiB |     96%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   5  NVIDIA A100-SXM...  Off  | 00000000:80:01.0 Off |                    0 |
| N/A   59C    P0   268W / 400W |  10914MiB / 40536MiB |     93%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   6  NVIDIA A100-SXM...  Off  | 00000000:80:02.0 Off |                    0 |
| N/A   57C    P0   316W / 400W |  10914MiB / 40536MiB |     92%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   7  NVIDIA A100-SXM...  Off  | 00000000:80:03.0 Off |                    0 |
| N/A   58C    P0   277W / 400W |  10482MiB / 40536MiB |     91%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A     70202      C   /usr/bin/python                 10479MiB |
|    1   N/A  N/A     70203      C   /usr/bin/python                 10911MiB |
|    2   N/A  N/A     70204      C   /usr/bin/python                 10911MiB |
|    3   N/A  N/A     70205      C   /usr/bin/python                 10911MiB |
|    4   N/A  N/A     70206      C   /usr/bin/python                 10911MiB |
|    5   N/A  N/A     70207      C   /usr/bin/python                 10911MiB |
|    6   N/A  N/A     70208      C   /usr/bin/python                 10911MiB |
|    7   N/A  N/A     70210      C   /usr/bin/python                 10479MiB |
+-----------------------------------------------------------------------------+

5. 次はマルチ環境で動かす w/ Docker

もう一台立ち上げます

$ PROJECT="xxx"
$ REGION="us-central1"
$ ZONE="us-central1-f"
$ TYPE="a2-highgpu-8g"

$ gcloud compute instances create a100-2 \
    --project=${PROJECT} \
    --zone=${ZONE} \
    --machine-type=${TYPE} \
    --resource-policies a100-${REGION} \
    --maintenance-policy=TERMINATE \
    --no-restart-on-failure \
    --image-family=tf-latest-gpu-ubuntu-2004 \
    --image-project=deeplearning-platform-release \
    --boot-disk-size=2TiB \
    --network-interface=nic-type=GVNIC \
    --metadata="install-nvidia-driver=True,proxy-mode=project_editors" \
    --scopes=https://www.googleapis.com/auth/cloud-platform

マルチ環境を実現するには以下を行います

1. .ssh/config を作成
2. sshの鍵認証でパスワード無しでログインできるようにする
3. job/hostfileの作成

リポジトリをクローン

$ git clone https://github.com/SO0529/gpt-neox.git
$ cd gpt-neox

ssh/config を作成

コツは port を 2222 にしているところです。docker 上の ssh を 2222 ポートで Listen して DeepSpeed から通信できるようにしています。

$ vim ~/.ssh/config

Host a100x8-para1
    HostName xx.xx.xx.xx
    Port 2222
    StrictHostKeyChecking no
    IdentityFile ~/.ssh/id_rsa
    user root
Host a100x8-para2
    HostName xx.xx.xx.xx
    Port 2222
    StrictHostKeyChecking no
    IdentityFile ~/.ssh/id_rsa
    user root

authorized_keys を作成

sshの鍵認証でパスワード無しでログインできるようにするために、authorized_keys を作成します。

$ ssh-keygen -t rsa
$ cat ~/.ssh/id_rsa.pub >> authorized_keys

hostfile を作成

DeepSpeed 用に hostfile を作成します。

$ vim hostfile

a100x8-para1 slots=8
a100x8-para2 slots=8

Docker を build

• Dockerfile

    FROM nvidia/cuda:11.1.1-devel-ubuntu20.04
  
    # 説明する部分以外は割愛します    
  
    # SSH
    COPY .ssh/id_rsa /root/.ssh/id_rsa
    RUN chmod 600 /root/.ssh/id_rsa
    COPY authorized_keys /root/.ssh/authorized_keys
    RUN chmod 644 /root/.ssh/authorized_keys
    COPY hostfile /job/hostfile
    COPY .ssh/config /root/.ssh/config
  
    RUN echo 'Port 2222' >> /etc/ssh/sshd_config
    RUN echo 'PasswordAuthentication no' >> /etc/ssh/sshd_config
    EXPOSE 2222
  

インフラ環境部分のキモは # ssh以下の部分です。 なお、本番利用する場合は id_rsa を Docker に埋め込むとかは辞めましょう。検証環境なのでヨシとしています。

1. id_rsa と authorized_keys をコピーしてパスワード無しで鍵認証を行うための仕込み。
2. hostfile をコピーして、DeepSpeed からホストを認識するための仕込み。
3. .ssh/config をコピーして、ホストを認識する時の名前解決と鍵やポート番号を指定。
4. 最後に sshd_config を設定して sshd を 2222 port で起動するように指定。

これで Docker の SSH ポートを 2222 で Listen し、パスワード無しでホスト間を ssh 通信できるようになりました。

• docker-compose.yml

    version: "3"
    services:
      gpt-neox:
        image: gpt-neox:latest
        volumes:
          - ./:/gpt-neox
          - /share:/share
        deploy:
          resources:
            reservations:
              devices:
                - driver: nvidia
                  count: 8
                  capabilities: [gpu]
        tty: true
        shm_size: '16gb'
        ulimits:
          memlock: -1
        network_mode: "host"
        command: /bin/bash -c "sudo /etc/init.d/ssh restart && tail -f /dev/null"
  

あとは、docker compose 時にネットワークモードを host に指定することで、ホストのネットワークに対してポート番号をexposeし、 command で sshd を常時起動するようにしています。この辺りは他の手法もあるので参考情報としてください。

あとはビルドして起動すればOKです。

$ docker build -t gpt-neox -f Dockerfile .
$ docker compose -f docker-compose.yml -p gpt-neox up -d

6. つまずいたポイント

学習途中に出力したファイルを再利用するのでNFSが必要に

どうも1系で生成された中間生成物が無いよ。って怒られます。共用ストレージを使いますか…。GCPはNFSのみだし遅くならないかなぁ...

a100-b: FileNotFoundError: [Errno 2] No such file or directory: 'data/jawiki-cc100_bpe/jawiki-cc100_bpe_text_document_train_indexmap_2560000ns_2048sl_1234s_doc_idx.npy'
a100-b: FileNotFoundError: [Errno 2] No such file or directory: 'data/jawiki-cc100_bpe/jawiki-cc100_bpe_text_document_train_indexmap_2560000ns_2048sl_1234s_doc_idx.npy'

速度を上げるために PREMIUM ティアを選択しました。それ以外は大きな指定無し

PROJECT="xxx"
LOCATION="us-central1-f"

gcloud filestore instances create neox \
    --project=${PROJECT} \
    --location=${LOCATION} \
    --tier=PREMIUM \
    --file-share=name="neox",capacity=10TiB \
    --network=name="default"

NFSのリージョンを間違えて速度が出なかった

実はNFSをデプロイした時にリージョン間違えてた。VMはus-central1なのに、NFSはasia-southeast1。通信出来ちゃってるので気付くのが遅れた…。ご迷惑をおかけしました。

そしてNFSとVMはゾーンレベルでも合わせた方がいい。

大量のGPUの調達はリソースを確保できないかもしれないので要サポート確認

今回A100を使用したのですが、リージョンによっては A100 が無いリージョンがあったり、ゾーンによっても利用できる量が異なります。一定量使う前提であれば、サポートに確認してどこのリージョン・ゾーンを使うのが良いのか事前に聞いておくとリソース不足問題が発生しにくいかもしれません。

コンパクトプレースメントポリシーは邪魔になりそうだった

「アベイラビリティゾーンの中で、ネットワーク的に近いノードにVMを配置してくれる機能」ですが、そもそもA100を確保時に近いノードにA100が無ければ調達できない可能性があるとのことで、調達を優先してコンパクトプレースメントポリシーは外すことにしました。

7. 結果

マルチノード環境（a100x8-para1-0）がシングルノード（a100x8-1-0）環境に対して概ね倍の速さで収束できていることが確認できました。

8. まとめ

DeepSpeed が優秀過ぎてちょっと環境整えるだけでした。GCP側もgVNICやFast Socket等のプリセット・環境が整備されていたので、ちょっと整えるだけでした。オンプレだったら大量のGPUをキッティングしまくったり100 Gbpsのネットワークを作ったりと「本題に行くまでに時間と費用いくらかかるねん」って感じですが、クラウドって便利だなと改めて実感しました。まる。

9. 最後に (採用情報)

ABEJAでは様々な技術で産業構造変革のチャレンジをしています！画像、自然言語等のドメインはもちろん、データサイエンスの価値をお客様にデリバリーする為の周辺技術にも力を入れいています。ABEJAに興味が出た方はぜひお話しましょう。ご連絡お待ちしております。 下記HRMOS経由で気軽にお問合せください。

hrmos.co