OpenAI 大規模ニューラルネットワークのトレーニング手法

大規模なニューラルネットワークのトレーニングには、単一の同期計算を実行するために GPU クラスタを調整する必要があります。モデルとクラスタの規模が拡大したため、実務者はさまざまな並列化手法を用いて計算負荷とメモリ要件を複数のハードウェアアクセラレータに分散させています。

データ並列

データ並列は、同一のモデルパラメータを各ワーカーにコピーし、トレーニングバッチの異なるサブセットを同時に処理させることで、複数の GPU の使用を可能にします。これにより多数の GPU の計算能力を活用できますが、モデルは依然として単一の GPU のメモリに収まる必要があります。

ワーカー間の一貫性を保つために、システムはパラメータ更新を調整しなければなりません。標準的な同期アプローチは以下の3ステップで構成されます：

  1. 各ワーカーで勾配を独立に計算する。
  2. 全ワーカー間で勾配を平均化する（ブロッキング通信ステップ）。
  3. 各ワーカーで新しいパラメータを独立に計算する。

ブロッキング平均のオーバーヘッドを削減するための非同期同期方式も存在しますが、これらは学習効率を低下させることが多く、ほとんどの実務者は同期方式を好む傾向にあります。

パイプライン並列

パイプライン並列は、モデルの連続したチャンク（連続レイヤー）を異なる GPU に分割し、デバイスごとのメモリ使用量を削減します。

「バブル」問題への対処

単純なパイプライン実装では「バブル」―前のマシンの出力を待つワーカーのアイドル時間―が発生します。これらのバブルを最小化するために、バッチはより小さなマイクロバッチに分割されます。各ワーカーは次のマイクロバッチが利用可能になり次第処理を開始し、計算と待機時間を重ね合わせます。勾配はマイクロバッチ間で平均化され、パラメータ更新はすべてのマイクロバッチが完了した後にのみ行われます。

スケジューリング戦略

順方向と逆方向のパスを管理するために、主に2つのスケジューリング方式が使用されます：

  • GPipe: ワーカーは順方向と逆方向のパスを連続して処理し、最後に複数のマイクロバッチからの勾配を同期的に集約します。
  • PipeDream: ワーカーは順方向と逆方向のパスを交互に処理し、効率が向上する可能性がありますが、一部の計算が古いパラメータを使用することがあります。

テンソル並列

テンソル並列は、レイヤー内の個々の演算を GPU 間で「水平」に分割します。Transformer のような最新アーキテクチャでは、主なボトルネックは活性化バッチ行列と大規模な重み行列の乗算です。テンソル並列は重み行列を均等なサイズのシャードに分割し、各シャードを別々の GPU に配置して全体の行列積の一部を計算し、結果を通信して結合します。

具体的な実装例は以下の通りです：

  • Megatron-LM: Transformer の MLP と自己注意層内の行列乗算を並列化します。
  • PTD-P: テンソル、データ、パイプライン並列を組み合わせ、非連続レイヤーを各デバイスに割り当ててバブルオーバーヘッドを削減します。
  • Sequence Parallelism: 入力シーケンスを時間軸で複数のサブ例に分割し、ピークメモリ使用量を削減します。

Mixture-of-Experts (MoE)

Mixture-of-Experts (MoE) は、計算コストを比例的に増やすことなくモデルのパラメータ数をスケールさせることを可能にします。これは、ゲーティング機構によりネットワークの重みの一部（「エキスパート」）だけを選択して任意の入力に対する出力を計算することで実現されます。異なるエキスパートを異なる GPU に配置できるため、MoE はモデルに使用する GPU 数をスケーラブルに増やす手段を提供します。

メモリ節約設計

並列化に加えて、トレーニングに必要なデバイス RAM を削減するためにいくつかの戦略が採用されています：

  • Checkpointing (Activation Recomputation): 勾配計算のためにすべての元の活性化を保存する代わりに、チェックポイントは一部だけを保存し、逆伝播時に必要に応じて中間活性化を再計算します。選択的活性化再計算は、保存コストは高いが計算コストは低い活性化のみをチェックポイントすることでさらに最適化します。
  • Mixed Precision Training: 低精度（通常は FP16）の数値でトレーニングすることで、FLOP 数を増やしつつデバイス RAM を節約し、精度の低下は最小限に抑えられます。
  • Offloading: 未使用データを一時的に CPU や他のデバイスへ移動します。ZeRO 実装はパラメータ、勾配、オプティマイザ状態をハードウェア間で分割し、必要に応じて実体化します。
  • Memory Efficient Optimizers: Adafactor のようなオプティマイザを使用して、実行状態のメモリフットプリントを削減します。
  • Compression: 中間結果を圧縮します。例えば Gist は逆伝播の活性化を圧縮し、DALL·E は同期前に勾配を圧縮します。

Sources