AWS、EKSでの分散AIトレーニングを簡素化するRay Train DLCをリリース

Amazon Web Services (AWS) は、Amazon Elastic Kubernetes Service (EKS) 上での分散機械学習トレーニングを効率化するために設計された事前設定済みのイメージである Ray Train Deep Learning Container (DLC) をリリースしました。2026年10月に公開されたこのコンテナには、Ray、PyTorch、および高性能コンピューティングに必要なネットワークライブラリがバンドルされており、エンジニアは複雑な依存関係チェーンを管理せずにマルチノードのトレーニングジョブを展開できます。

何が起きたか

クラスター全体でPyTorchトレーニングをスケーリングする際の主要なボトルネックは、モデルコードそのものよりも、コンテナインフラストラクチャであることが一般的です。エンジニアは通常、CUDAランタイム、PyTorchビルド、NVIDIA Collective Communications Library (NCCL)、およびElastic Fabric Adapter (EFA) ユーザースペースライブラリの整合性を合わせるために多大な時間を費やします。これらのコンポーネント間の不一致により、NCCLが高速度のEFA接続を使用せず、遅いTCPソケットにフォールバックするなど、サイレントなパフォーマンス低下が発生することがあります。これに対応するため、AWSは統一されたイメージを組み立て、テストし、Amazon Elastic Container Registry (ECR) Publicで公開しました。

この新しいDLCはAWSによって構築・保守され、スタックへのパッチ適用や、マルチノードEFAおよびKubeRayトレーニングテストを通じた検証が行われます。リリースにはRay 2.58.0、PyTorch 2.13.0、CUDA 13.0.2、NCCL 2.29.7が含まれており、さらにNCCLをEFAに接続するAWS OFI NCCLプラグインも含まれています。これを公開イメージとして提供することで、AWSはチームがこれらの基盤ライブラリの更新ごとに独自のカスタムコンテナイメージを構築・維持する必要をなくしました。

付随する技術ガイドでは、このDLCを使用して、8つのNVIDIA L4 GPU上でPyTorch Fully Sharded Data Parallel (FSDP) を用いてQwen2.5-1.5Bモデルをファインチューニングする方法を示しています。セットアップでは、単一のAvailability Zone内の2台のg6.12xlargeインスタンスを使用し、KubeRayオペレーターによって管理されます。この構成により、DLCが既に必要なネットワークスタックを含んでいるため、ノード間集合通信をすぐにEFA経由で実行できます。

仕組み

Ray Train DLCは、Rayヘッドポッドとワーカーポッドの両方にサービスを提供する、役割非依存の単一イメージとして機能します。Kubernetes上では、KubeRayオペレーターがRayClusterマニフェストを使用して必要なポッドを作成し、調整のためにワーカーをヘッドノードへ指し示します。イメージにはすべての必要なドライバとライブラリが含まれているため、ユーザーはカスタムイメージを構築する必要がありません。代わりに、Kubernetesマニフェスト内の環境変数とリソース要求を通じてクラスターを設定します。

ネットワーク処理は、事前設定された環境変数を通じて自動的に管理されます。DLCはデフォルトで FI_PROVIDER=efa 設定によりEFAを通信に使用し、NCCL設定からループバックインターフェースを除外します。ジョブがベースイメージに含まれていない追加のPythonパッケージを必要とする場合、Rayの runtime_env 機能が送信時にそれらをインストールし、ベースイメージを変更しないままにします。この分離により、コンテナの安定性が保たれつつ、特定の実験的なニーズに対する柔軟性が確保されます。

コンテナ内部では、Ray Trainが分散プロセスグループの形成を管理します。エンジニアは標準的なPython関数としてトレーニングループを書き、Rayが利用可能なすべてのGPUでそれを実行します。FSDPワークロードの場合、Rayはモデル重み、勾配、オプティマイザー状態をGPU間で分割し、デバイスごとのメモリ要件を削減します。DLCは、FSDPのgatherおよびreduce-scatter操作によって生成される重い通信トラフィックが、標準的なネットワークソケットではなく、高帯域幅のEFAインターフェースを経由して移動することを保証します。

主要な詳細

• イメージタグ: サンプルでは public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1 を使用しており、これにはRay 2.58.0とPyTorch 2.13.0が含まれています。
• ハードウェアサポート: リファレンスアーキテクチャでは、NVIDIA L4 GPUとEFAインターフェースを備えたg6.12xlargeインスタンスを使用し、単一のAvailability Zoneに展開します。
• ネットワークスタック: CUDA 13.0.2、NCCL 2.29.7、およびAWS OFI NCCLプラグインを含み、TCPフォールバックではなくEFAの使用を確認するためにテストされています。
• 設定方法: クラスターの形状とトランスポート設定は、イメージの再構築ではなく、環境変数とKubernetesリソース要求によって定義されます。
• コスト見積もり: サンプルセットアップのコストは約時給$9.50であり、2台のGPUインスタンス、EKSコントロールプレーン、およびNAT Gatewayを含みます。
• ドライバ互換性: ホストのNVIDIAドライバがイメージのCUDAバージョンより古い場合、コンテナのエントリーポイントが自動的にCUDA前方互換性ライブラリを有効にします。

なぜ重要なのか

ソフトウェアエンジニアやMLチームにとって、このリリースは分散トレーニングの運用オーバーヘッドを大幅に削減します。CUDA、NCCL、EFAドライバの互換性のあるバージョンを管理することは脆いプロセスであり、しばしば微妙なパフォーマンス問題を引き起こします。テスト済みの公開イメージを提供することで、AWSはチームがインフラストラクチャのデバッグではなく、モデル開発に集中できるようにします。特定バージョンのピン留めやセキュリティ更新のための移動タグの使用が可能であり、再現可能な研究と本番環境の安定性の両方に対して柔軟性を提供します。

KubeRayおよびEKSとの統合もクラスター管理を簡素化します。エンジニアは単一のマニフェストでトレーニングインフラストラクチャ全体を定義し、GPU数やEFAインターフェースを標準的なKubernetesリソースとして指定できます。このアプローチは既存のKubernetesツールを活用し、Rayクラスターのライフサイクルイベントを管理するためのカスタムスクリプトの必要性を回避します。その結果、AWSインフラストラクチャ上で大規模言語モデルのファインチューニングジョブを実行するための、より信頼性が高くスケーラブルなパスが得られます。

できること

• サンプルリポジトリの確認: EKS上のRay Trainジョブの構造を理解するために、提供されているコードとマニフェストにアクセスしてください。
• クォータ可用性の確認: AWSアカウントにGファミリーインスタンスの十分なオンデマンドキャパシティ（特に2台のg6.12xlargeノード用に96 vCPU）があることを確認してください。
• 必要なツールのインストール: クラスター管理とジョブ送信のために、AWS CLI v2、eksctl、kubectl、Helm、および envsubst をセットアップしてください。
• DLCのローカルテスト: クラスターに展開する前に、ECRから公開イメージを取得し、ローカル開発環境との互換性を検証してください。
• 環境変数の設定: リージョン、インスタンスタイプ、イメージタグなどのクラスターパラメータを設定するために、提供されている scripts/env.sh を使用してください。
• ネットワークパフォーマンスの監視: マニフェスト内で NCCL_DEBUG=INFO を有効にし、トレーニング中にノード間通信にEFAが使用されていることを確認してください。