分散学習とは

深層学習において、必要な計算を多数の計算機に分散させることで高速化を図る手法である。

大きくは以下2つの方法に分けることができる。

Data Parallel

異なるデータで同じモデルを並列処理し、実行効率を上げる方法。

手順

各デバイスに異なるデータ（バッチ）を割り当てる
学習・逆伝播を行う
勾配を共有し、モデルを更新する

更新方式

同期更新：パラメータサーバ
- 1つ以上のプロセスモデルでモデルを管理
- 各GPUからすべて勾配を受け取り次第更新
非同期更新：Peer-to-Peer
- GPU間で直接勾配情報を通信
- 1つ以上の勾配計算が完了し次第、その時点での情報で更新

Model Parallel

モデルを分散し、1つのGPUに収まらない計算を実行可能にする。

手順

モデルを分割し、それぞれを別のGPUに展開
必要に応じてGPU間通信をしながら学習

分散学習の課題

分散学習は、一般的に処理時間を短くできることや大規模なモデルを学習できることなどメリットばかりに思えるが、課題もあるらしい。

Large Batch Size Problem

大規模化に伴う精度劣化

小さいバッチサイズでは、1ステップあたりのデータが少ないため、データの一部しか反映できない
⇒局所解に捕らわれにくくなる。
バッチサイズを大きくすると、1ステップ当たりのデータが多くなり、データ全体に適応しやすくなってしまう
⇒局所解にとらわれやすい
バッチサイズが大きいと、学習率が過剰に大きくなってしまう
⇒LARS、LAMBといった学習率をスケールアップする手法が有用らしい

Staleness

非同期更新時の精度劣化

速く終了した勾配が最新のモデルに対して古くなってしまい、劣化を招く

通信コストの増大

学習の過程で必ずどこかで情報を集約・通信する必要があるため、コストが生じる

基本概念

Host

分散学習ノード（マシン）のIPアドレスやホスト名

Port

ノード間通信に使用するネットワークポート番号

Rank

分散環境におけるノードやプロセスの識別番号

World Size

分散環境全体のプロセス数（全GPU数）

分散学習の種類

Data Parallel方式

Data Parallel (DP)

1つのノードがメインとなって、GPUを管理する。オーバーヘッドが大きいため、現在は非推奨。

それぞれのGPUにmodelをコピー
バッチを分割して、各GPUで処理
情報を集約して勾配の計算

Distributed Data Parallel (DDP)

それぞれのGPU上のpipelineを別々のprocessが持つ。GPU間の通信は勾配の集約・分散のみとDDより少ない。

それぞれのGPUで学習を進める
勾配を計算
それぞれのGPU同士で勾配の集約・分散を行う

Model Parallel方式

Tensor Parallel (TP)

行列積は行方向・列方向ともに分割できるという特性から、TensorをそれぞれのGPUに分割する。

Pipeline Parallel (PP)

モデルを分割し、レイヤー毎にGPUへ振り分けることで、１つのGPUに乗りきらない大規模なModelでも学習できる。

Fully Sharded Data Parallel (FSDP)

モデルパラメータ、オプティマイザの状態、および勾配を分割して、各GPUに分散する手法。これによりメモリを効率的に使用し、大規模なモデルを学習できる。

PyTorchに統合されており、比較的実装が簡単らしい。

ZeRO (Zero Redundancy Optimizer)

メモリ効率を最適化するため、段階的にモデルパラメータ、勾配、オプティマイザの状態を分散して管理する仕組み。

DeepSpeedというMicrosoftのフレームワークに実装されており、複雑な通信パターンにより実装が少し難しいらしい。

使い分け

モデルが1つのGPUに載る場合 ⇒ DDP
モデルが1つのGPUに載らない場合
- 演算ごとに細かく分割したい ⇒ TP
- モデルの階層ごとに細かく分割したい ⇒ PP
- PyTorchに分割はお任せしたい ⇒ FSDP

参考文献

第111回お試しアカウント付き並列プログラミング講習会「第1回ディープラーニング分散学習ハッカソン」分散学習基礎講座
 大規模言語モデルを支える分散並列学習の仕組み Part1
Distributed and Parallel Training for PyTorch