Running PyTorch at Scale on EKS without Building Containers

According to the AWS blog, the Ray Train Deep Learning Container (DLC) provides a pre‑built image that contains a compatible CUDA stack, NCCL, EFA libraries and Ray 2.58.0, letting you train PyTorch models on Amazon EKS without ever writing a Dockerfile. The claim matters for teams that spend days aligning driver versions, rebuilding images, and debugging silent network mismatches before they can even start a training run.

What the Ray Train DLC actually includes

The table shows that the only thing you still control is the version tag – currently public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1 – and the runtime environment variables that choose the network transport.

How KubeRay turns a public image into a multi‑node Ray cluster

KubeRay is an open‑source Kubernetes operator that adds a RayCluster custom resource. The manifest describes a head group (zero GPUs) and a worker group (four GPUs per pod). When the operator sees the manifest it creates the corresponding pods, starts Ray in the appropriate role, and wires each worker to the head’s address. Because the same Docker tag runs both head and workers, the only configuration you supply is:

• Instance type and count (e.g., two g6.12xlarge nodes)
• GPU and EFA resource requests (nvidia.com/gpu: 4, vpc.amazonaws.com/efa: 1)
• Environment variables that force NCCL to use the EFA transport (FI_PROVIDER=efa, NCCL_NET_PLUGIN=aws, etc.)

If a pod dies, KubeRay recreates it and re‑joins it to the cluster, so the Ray control store remains consistent without manual intervention.

Running Fully Sharded Data Parallel (FSDP) on eight L4 GPUs

The sample in the blog fine‑tunes the 1.5 B‑parameter Qwen2.5 model with PyTorch FSDP across eight NVIDIA L4 GPUs (four per g6.12xlarge). FSDP splits model weights, optimizer state and gradients across all GPUs, reducing the per‑GPU memory footprint to roughly 1.5 GB in bfloat16. The trade‑off is that each training step triggers a full model‑wide all‑reduce: weights are gathered before the forward pass and gradients are reduce‑scattered afterward. On a single node the traffic stays on PCIe; across nodes the traffic moves over EFA, which delivers low‑latency, high‑throughput RDMA without falling back to TCP.

Because the DLC already bundles the EFA userspace libraries and the AWS‑OFI NCCL plugin, the network stack is verified before each release. A mismatched NCCL/EFA build would silently fall back to TCP sockets, still finish the job but at a fraction of the speed – a failure mode that often goes unnoticed.

The hidden cost and trade‑off nobody spells out

• Instance dependency – The sample locks you to g6.12xlarge (four L4 GPUs, 40 Gbps networking, one EFA). Those instances are only available in selected AZs and require a quota of 96 vCPUs. If your account lacks that quota, you must request a raise or fall back to a slower instance type, which defeats the purpose of the DLC’s performance claim.
• AZ locality – EFA traffic does not cross Availability Zones. The sample places both GPU nodes in the same zone; scaling beyond a single AZ would need additional networking tricks or a different transport, eroding the simplicity the DLC promises.
• Cost visibility – The blog lists an hourly spend of about $9.50 for the two GPU nodes, the EKS control plane and a NAT gateway. That number excludes data‑egress, persistent storage, and any on‑demand price spikes in a hot region. Teams that budget per‑job need to add those line items.
• Vendor lock‑in – While the image is public, it only supports the exact CUDA, NCCL and EFA versions that AWS ships. If you need a newer driver, a different NCCL patch, or a custom CUDA kernel, you must revert to a self‑built image, losing the “no‑rebuild” advantage.

In practice the trade‑off is clear: you gain a turnkey path to multi‑node PyTorch training at the cost of committing to AWS‑specific hardware, networking and versioning. The approach is best for short‑term experiments, proof‑of‑concepts, or teams that already have EKS and KubeRay pipelines. It is less attractive for production workloads that need cross‑AZ resilience or fine‑grained version control.

What you can try today

1. Create an EKS cluster (minimum one m7i.xlarge system node). The AWS CLI command eksctl create cluster --name ray‑test --region us-east-1 --nodegroup-name system --node-type m7i.xlarge --nodes 1 does the job.
2. Install KubeRay with Helm:

   helm repo add kuberay https://ray-project.github.io/kuberay-helm/
   helm install kuberay kuberay/kuberay-operator
   

3. Pull the public DLC by adding public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1 as the image field in a simple RayCluster manifest (the blog provides a ready‑made YAML you can copy).
4. Submit the sample trainer:

   kubectl exec -it <ray-head-pod> -- bash -c "ray job submit --working-dir . --entrypoint 'python train.py --steps 10'"
   

   The job will spin up four workers on each g6.12xlarge node, run the FSDP loop, and print loss values to the Ray dashboard.
5. Watch the metrics in the Ray dashboard (port 8265) to see GPU utilization and NCCL bandwidth. If you see the NCCL_DEBUG=WARN flag fire, you know the EFA path is active.

By following those steps you can validate whether the DLC eliminates the image‑build bottleneck for your own model. If you hit a quota or AZ limitation, the failure itself tells you where the hidden cost lies.

Sources

• Distributed training with Ray Train DLCs and EFA on EKS