NCCL Fast Socket で分散型 ML のトレーニング時間を短縮

※この投稿は米国時間 2023 年 7 月 14 日に、Google Cloud blog に投稿されたものの抄訳です。

大規模言語モデル、ジェネレーティブ AI、ビジョンモデルなどの大規模な ML モデルは、トレーニング可能なパラメータを劇的に増加させ、最先端の結果を達成しています。しかし、パラメータの数が増えると、モデルが大きすぎて単一の VM インスタンスで対応できなくなるため、モデルを複数のノードに分散させる分散コンピューティングが必要になります。Google Kubernetes Engine（GKE）は、NCCL Fast Socket を組み込みでサポートしているため、分散されたマルチノード クラスタを使用して大規模な ML モデルをトレーニングする時間を短縮できます。

企業は、ML モデルをトレーニングするために、費用を抑えながら高速なパフォーマンスを得る方法を模索しています。分散トレーニングでは、ノード間の勾配通信がパフォーマンスのボトルネックになるため、ノード間のレイテンシを最適化することが、トレーニング時間と費用の削減において重要といえます。分散トレーニングは、複数のホスト間でのネットワークのトランスポート層として集団通信を使用します。ML の分散トレーニングでは、all-gather、all-reduce、broadcast、reduce、reduce-scatter、point-to-point send and receive などの集団通信プリミティブが使用されます。

NVIDIA Collective Communication Library（NCCL）は、TensorFlow や PyTorch などの一般的な ML フレームワークでよく使用されています。NVIDIA GPU 間の高帯域幅と低レイテンシに対応する、高度に最適化された実装です。Google は、Google Cloud でのディープ ラーニングのパフォーマンスを最適化するために、NCCL Fast Socket と呼ばれる NCCL の独自バージョンを開発しました。

NCCL Fast Socket は、さまざまな技術を使用して、より一貫性のある優れた NCCL パフォーマンスを実現しています。

• 最大スループットを達成するための複数のネットワーク フローの利用。NCCL Fast Socket では、複数の通信要求のオーバーラップの改善など、NCCL の組み込みマルチストリームのサポートが最適化されています。

• 複数のネットワーク フローの動的なロード バランシング。NCCL は、変化するネットワークやホストの状況にも適応します。この最適化によって、ストラグラー ネットワーク フローが NCCL 全体の集団オペレーションを大幅に遅らせることはありません。

• Google Cloud の Andromeda 仮想ネットワーク スタックとのインテグレーション。これにより、仮想マシン（VM）の競合を回避し、ネットワーク全体のスループットが向上します。

私たちは、2 ノードの GKE クラスタを使用するさまざまなマシンシェイプで、NCCL Fast Socket と NCCL のパフォーマンスの比較テスト（NVIDIA NCCL テスト）を行いました。

次のグラフは、その結果を示しています。各マシンシェイプにおいて、Fast Socket を使用しない NCCL のパフォーマンスを 1 として正規化しています。いずれのマシンシェイプも、NCCL Fast Socket を使用することで、インターネットワークでの通信速度が 1.3～2.6 倍の範囲でパフォーマンスが向上することが実証されました。

組み込みの機能により、GKE のユーザーはアプリケーションを変更または再コンパイルすることなく、NCCL Fast Socket を使用できます。これは、ML フレームワーク（TensorFlow や PyTorch など）でも、NCCL ライブラリ自体でも同様です。NCCL Fast Socket を使い始めるには、--enable-fast-socket と --enable-gvnic フラグを指定して、プラグインを使用するノードプールを作成します。また、gcloud container node-pools update を使用して、既存のノードプールを更新することもできます。

NCCL でさらに高いネットワーク スループットを実現するには、VM インスタンスの作成時に Google Virtual NIC（gVNIC）を有効にする必要があります。gVNIC の使用方法について詳しくは、gVNIC ガイドをご覧ください。

NCCL Fast Socket が有効になっていることを確認するには、kube-system Pod を表示します。

出力は次のようになります。

詳細については、GKE NCCL Fast Socket ドキュメントをご覧ください。NCCL Fast Socket により、GKE での ML トレーニングがどのように向上したかについて、フィードバックをお待ちしております。