データ並列(Data Parallelism)とは、マルチGPU学習において最も基本的な並列化手法の一つで、モデルの完全なコピーを各GPUに配置し、学習データを分割して各GPUに異なるバッチを処理させる方式です。各GPUで計算された勾配を集約して同期することで、効率的な並列学習を実現します。
データ並列の仕組み
データ並列では、全GPUが同じモデルパラメータを持ちます。各GPUは異なるミニバッチのデータで順伝播・逆伝播を行い、勾配を計算します。その後、AllReduce通信により全GPUの勾配を平均化し、全GPUのパラメータを同じ値に更新します。N台のGPUを使えば、理想的には実効バッチサイズがN倍になり、学習速度もN倍に近づきます。
| ひとことで言うと | 同じモデルを各GPUに複製し、データを分けて同時に学習する方式。 |
|---|---|
| 何がうれしいか | 実装が簡単で、GPUを増やすほど学習が速くなる。最も一般的な方式。 |
| 注意点 | モデル全体が1枚のGPUに収まる必要がある。大きすぎると使えない。 |
通信の最適化
データ並列のボトルネックは勾配の集約(AllReduce)通信です。Ring AllReduceアルゴリズムにより、GPU数に対して効率的な通信が可能です。また、勾配の計算と通信をオーバーラップさせる(逆伝播の途中から通信を開始する)ことで、通信待ち時間を隠蔽する最適化も行われます。
1モデルを複製全GPUに同じものを配置
▶
2データを分割GPUごとに別の分を担当
▶
3並行して計算それぞれ勾配を求める
▶
4勾配を集約全GPUで平均を取る
▶
5同じ更新を適用モデルの同期を保つ
↺ バッチごとに1〜5をくり返す
ZeROとFSDP
DeepSpeedのZeRO(Zero Redundancy Optimizer)やPyTorchのFSDP(Fully Sharded Data Parallel)は、データ並列の発展形です。パラメータ、勾配、オプティマイザ状態をGPU間で分割して保持することで、各GPUのメモリ使用量を大幅に削減し、より大規模なモデルの学習を可能にします。