PyTorch のトレーニングをクラスタへスケールアウトする際に障害となるのは、モデルのコードではなくコンテナであることが多い。各トレーニングコンテナには、互いに整合した CUDA ランタイム、その CUDA 向けにコンパイルされた PyTorch、NCCL、EFA のユーザー空間ライブラリ、NCCL と EFA を接続する AWS OFI NCCL プラグイン、そしてクラスタ運用者と合意したバージョンの Ray が必要になる。そのうちどれか 1 つが更新されるたびにイメージを再構築して検証し直す必要があり、これはモデルをトレーニングするという本来の目的とは無関係な作業だ。AWS はこの部分を Ray Train ディープラーニングコンテナ（DLC）で引き受けている。このイメージは Amazon ECR Public で公開され、常にパッチが適用され、各コンポーネントは一体としてテストされている。
このサンプルは、そのコンテナを公開されたままの状態で使い、2 台の g6.12xlarge 上で 8 基の NVIDIA L4 GPU にまたがり、PyTorch の Fully Sharded Data Parallel（FSDP）で Qwen2.5-1.5B をファインチューニングする。プロセス全体でイメージは一切ビルドせず、トレーニングコードは 1 つの関数だけ、クラスタは 1 つの RayCluster マニフェストと変数ファイルで表現される。コンテナが EFA スタックを同梱しているため、Pod が EFA を要求すれば、2 台のインスタンス間の集合通信は EFA 上を流れる。ソリューション全体はサンプルリポジトリにあり、マルチノード Ray Serve のチュートリアルと同じ KubeRay-on-EKS の基盤を利用している。
このイメージが固定するタグは public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1 であり、本来なら層ごとに維持しなければならないビルドを置き換える。分散ランタイムは Train、Tune、Data を含む Ray 2.58.0、フレームワークは PyTorch 2.13.0 と torchvision 0.28.0、GPU ランタイムは CUDA 13.0.2、cuDNN、NCCL 2.29.7、ネットワーク層は EFA インストーラ、libfabric、その NCCL 向けにビルドされた aws-ofi-nccl、GDRCopy、トレーニングライブラリは DeepSpeed、flash-attn、Transformer Engine、PyTorch Lightning、Hugging Face の各コンポーネント、診断ツールの all_reduce_perf もパスに用意されている。
チームが最も避けたいのがネットワーク層だ。ここでバージョンが不一致になっても通常はエラーにならず、NCCL が TCP ソケットへ静かにフォールバックし、ジョブは一応完走するが、速度は本来のごく一部になる。DLC ではこれらのライブラリが一緒にビルドされ、一緒にテストされ、AWS はリリース前にマルチノード EFA とマルチノード KubeRay のトレーニングテストを実施しており、これはまさに本サンプルが使う経路をカバーしている。さらにイメージは公開されているためノードは認証情報なしで取得でき、起動時にホストの NVIDIA ドライバを確認し、イメージ内の CUDA より古い場合は同梱の CUDA 前方互換ライブラリを自動で有効にする。イメージ自体に固定の役割はなく、起動側が ray start コマンドを渡すため、同じタグを head と worker の両方に使える。Ray Train イメージは Ray Serve や Ray LLM のイメージと同じ ECR リポジトリを共有し、train-ml プレフィックスを使う。
トレーニングジョブ間の違いはすべてイメージ外の設定で実現する。Pod の役割は KubeRay の rayStartParams で指定し、サンプルでは head に num-gpus: "0" を設定し、worker と同じ DLC_IMAGE を使う。GPU と EFA は Kubernetes のリソース要求で宣言し、例えば nvidia.com/gpu: 4 と vpc.amazonaws.com/efa: 1 とする。ネットワーク転送は FI_PROVIDER、NCCL_NET_PLUGIN、NCCL_SOCKET_IFNAME、NCCL_DEBUG の環境変数で制御し、サンプルはこれら 4 つをマニフェストで明示的に設定して挙動を分かりやすくし、イメージの既定値に依存しないようにしている。実験でイメージに含まれない Python パッケージが必要な場合は、送信時に Ray の runtime_env でインストールでき、イメージは変わらない。イメージのバージョンはバージョン付きまたはダイジェスト固定のタグを使う。vX.Y はセキュリティ再ビルドで更新され、完全な vX.Y.Z は元のビルドで凍結されるため、バイト単位で同一の実行が必要ならダイジェストを固定する。
クラスタのレイアウトでは、すべてのノードが同一 VPC のプライベートサブネットにあり、インターネットへは NAT Gateway 経由でのみ到達する。ノードはこれを通じて ECR Public からイメージを取得し、worker も Hugging Face からモデルの重みをダウンロードする。1 台の m7i.xlarge がシステムノードグループを構成し、kube-system のアドオン、KubeRay オペレーター、Ray ヘッドをホストする。ヘッドはポート 8265 のジョブサーバーとポート 6379 のコントロールストアを公開し、GPU を宣言しないためトレーニング作業は割り当てられない。ジョブは EKS API 経由でヘッド Pod 内から kubectl exec で送信されるため、ポート 8265 がクラスタ外に公開されることはない。トレーニングは 2 台の g6.12xlarge 上で実行され、各インスタンスは L4 GPU 4 基、48 vCPU、192 GiB メモリ、40 Gbps ネットワーク、EFA インターフェース 1 つを備える。各インスタンスはすべての GPU と EFA インターフェースを占有する Ray ワーカー Pod を 1 つ実行し、Ray Train は各 Pod に 4 つのワーカーを配置する。rank 0〜3 が一方、rank 4〜7 がもう一方となる。EFA のトラフィックはアベイラビリティーゾーンをまたがないため、スクリプトは両インスタンスを g6.12xlarge を提供することが確認された同一ゾーンに配置する。
コンテナ内部では、Ray Train が本来ランチャーコードとして書く必要のある部分を担う。マルチノード PyTorch では通常、GPU ごとにプロセスを起動し、それぞれに rank、world size、ランデブーアドレスを渡す。Ray Train ではトレーニングループを関数として書き、ScalingConfig とともに TorchTrainer に渡すだけで、Ray がクラスタ内の GPU のある場所に GPU ごとのワーカーを起動し、torch.distributed のプロセスグループを形成して各ワーカーで関数を実行する。サンプルは FSDP を要求し、重み、勾配、オプティマイザの状態を各 GPU に複製するのではなく 8 基の GPU に分割する。Qwen2.5-1.5B の bfloat16 版はこれらのテンソルに約 12 GB を必要とし、8 分割すると各 L4 GPU は約 1.5 GB を保持するだけになる。代償は通信だ。FSDP は各層の重みを使用前に gather し、逆伝播後に勾配を reduce-scatter するため、毎ステップでモデル全体が rank 間を流れ、2 台のインスタンス間のこのトラフィックを DLC の EFA スタックが担う。
開始前の準備として、AWS CLI v2 にログインし、eksctl、kubectl、Helm、GNU gettext の envsubst を用意する。G ファミリーのオンデマンド容量を 96 vCPU 分確保し、少なくとも 1 つのゾーンが g6.12xlarge を提供するリージョンが必要だ。オンデマンドキャパシティ予約が既にある場合は CAPACITY_RESERVATION_ID を設定すると、スクリプトはそれをそのまま使い容量検索を省略する。予算としては GPU インスタンス稼働中に 1 時間あたり約 9.50 ドルで、2 台のインスタンス、EKS コントロールプレーン、NAT Gateway を含む。aws sts get-caller-identity で CLI が想定のアカウントを向いているか素早く確認できる。
トレーニングコードは code/train.py にあり、コンテナ、クラスタ、ネットワークに固有の内容は含まれない。train_func は Ray によって GPU ごとに 1 回実行され、rank を取得し、モデルを読み込み、prepare_model に渡して FSDP でラップし、通常の PyTorch ループを回す。各ステップでランダムな入力、勾配のクリア、順伝播、逆伝播、そして各 rank が自身のシャードのみを更新するローカルな最適化を行い、ray.train.report でステップ数と損失を報告する。サンプルはパラメータ数に基づく自動ラップポリシーを使い、1000 万パラメータを超えるサブモジュールをそれぞれ FSDP ユニットにするため、スクリプトは特定モデルの層クラスに依存しない。ドライバーは実行中のクラスタに参加し、GPU 数を数えてトレーナーを起動する。ワーカー数はクラスタから決まるため、同じスクリプトが単一ノードでも複数ノードでも変更なしに動作する。