Distributed Training

分散学習では、通常は1台のマシン内、またはクラウドやオンプレミスのクラスター上にある複数のGPUやTPUなどのプロセッサーに、1つのモデルの学習処理を分散させます。これにより、学習をより早く完了したり、1台のデバイスのメモリには収まらない大きさのモデルを学習したりできます。デバイス間では、NVIDIA NCCLなどの通信ライブラリを使い、高速インターコネクト経由で勾配、活性化値、パラメーターのシャードを交換します。これは、最新の基盤モデルや大規模言語モデルの学習に使われており、マルチGPUワークステーションでの日常的な物体検出の学習も短縮します。

分散トレーニングの仕組み#

データ並列化は最も一般的な手法です。各GPUにモデル全体のレプリカを保持し、それぞれが学習データの各バッチから割り当てられた部分を処理します。その後、all-reduce演算に参加して勾配を平均し、すべてのレプリカが同じ勾配降下法の更新を適用します。PyTorchのDistributedDataParallel（DDP）は、この通信をバックワードパスと並行して実行します。グローバルなバッチサイズは、GPUごとのバッチサイズにGPU数を掛けた値です。そのため、通常は学習率もそれに応じてスケーリングする必要があります。Goyal et al.は、線形スケーリング則とウォームアップフェーズを使用して、8,192枚の画像を1バッチとして256基のGPUでImageNetのResNet-50を1時間で学習し、8基から256基のGPUへのスケーリング効率約90%を達成しました。

シャーディングを使ったデータ並列化では、レプリカに必要なメモリを削減します。ZeROとPyTorchのFully Sharded Data Parallel（FSDP）は、パラメーター、勾配、オプティマイザーの状態をGPU間で分割し、各レイヤーの計算中に限ってそのレイヤーの重み全体を集約します。そのため、GPUを追加するほど各GPUに必要なメモリが減少します。

モデル並列化では、ニューラルネットワーク自体を分割します。テンソル並列化では個々の行列乗算をGPU間で分割し、パイプライン並列化では連続するレイヤーを異なるGPUに割り当ててマイクロバッチを順次処理します。コンテキスト並列化では長いシーケンスを分割し、エキスパート並列化ではMixture of Expertsモデルのエキスパートを異なるGPUに配置します。Megatron-LMなどのフレームワークは、これらをデータ並列化と組み合わせ、数千億個のパラメーターを持つモデルを学習します。

マルチGPU学習の実践#

主なオーバーヘッドは通信です。all-reduceのトラフィックはモデルサイズとともに増加するため、ノード内でNVLink接続されたGPUは、低速ネットワークで接続されたノードよりもスケールしやすくなります。層を同期バッチ正規化に変換しない限り、バッチ正規化の統計量はGPUごとに計算されます。各GPUに割り当てられるバッチが小さい場合、この点が重要です。Ultralytics YOLO26はDDPを使ったマルチGPU学習に対応しています。複数のデバイスIDを渡すと、torch.distributed.runを使ってGPUごとに1つのプロセスを起動し、CUDAではNCCLバックエンドを使用します。また、batchの値を各GPUに均等に分割します。

from ultralytics import YOLO

# 事前学習済みYOLO26モデルを読み込みます
model = YOLO("yolo26n.pt")

# DDPでGPU 0と1を使って学習します。batch=32の場合、各GPUは1ステップあたり16枚の画像を処理します
results = model.train(data="coco8.yaml", epochs=100, imgsz=640, batch=32, device=[0, 1])

実際のアプリケーション#

• 大規模言語モデル: Metaは最大16K基のH100 GPUでLlama 3 405Bを学習し、テンソル、パイプライン、コンテキスト、データの各並列化を組み合わせた手法を、著者らは4D並列化と呼んでいます。
• 自動運転車: 認識チームは、車両フリートから収集したカメラ、ライダー、レーダーの大規模な記録アーカイブを使って検出モデルやセグメンテーションモデルを学習します。これは、大規模なGPUクラスターがなければ現実的ではありません。
• 医用画像解析: 高解像度の3D CTやMRIボリュームは単一GPUのメモリを圧迫します。バッチやモデルのシャードをGPU間で分散させることで、より大きなボリュームやモデルを学習できるようになります。

分散学習と関連概念の比較#

• 連合学習: 連合学習も多数のデバイスにまたがって学習しますが、プライバシー保護のため、データはスマートフォン、病院、その他のクライアントに保持され、モデルの更新情報だけが中央サーバーに送られます。分散学習では、速度を最大化するためにデータを1つのクラスターに集約します。
• 分割学習: 分割学習では、クライアントとサーバーの間でモデルを分割し、生データがクライアントから外部に出ないようにします。その代わりに、通信量が増加します。
• 勾配累積: 勾配累積では、複数の小さなバッチの勾配を合計してから1回の更新を行うことで、時間はかかりますが、1台のGPUで大きなバッチと同等の計算を実現します。一方、データ並列化では複数のGPUでこれらのバッチを同時に計算します。