Distributed training with Ray Train DLCs and EFA on EKS

What usually stops PyTorch training from scaling out to a cluster is the container, not the model code. Every container needs the same CUDA runtime , a PyTorch build compiled against it, NCCL (NVIDIA Collective Communications Library), the Elastic Fabric Adapter (EFA) userspace libraries, the AWS OFI NCCL plugin that connects NCCL to EFA, and a Ray version the cluster operator agrees with. Building that image and rebuilding it each time one of those pieces moves is a project of its own, and it has nothing to do with the model you set out to train. The Ray Train Deep Learning Container (DLC) is that image, built by AWS. AWS assembles the stack, tests it as one unit, publishes it in Amazon Elastic Container Registry (Amazon ECR) Public and keeps it patched.

Figure 1: One EKS cluster with a control plane managed by AWS, a CPU system node (m7i.xlarge) running the KubeRay operator and Ray head, and two EFA-enabled g6.12xlarge nodes in one Availability Zone. KubeRay builds it from one RayCluster manifest, Ray Train places eight workers, one per GPU, and inter-node collectives cross over EFA.

Ray runs a job across machines as a single cluster with two roles. The Ray head coordinates and runs the global control store that tracks cluster state, the dashboard and job server you submit jobs to. Ray worker nodes supply the GPUs and run the training processes. Each worker joins the cluster by connecting to the head’s address when it starts. On Kubernetes, each of these is a pod. Kubernetes has no built-in idea of a Ray cluster, so something must create the head and worker pods, start Ray in the right role in each, point the workers at the head, and replace pods that fail. KubeRay does this. It’s an open source Kubernetes operator that adds a RayCluster resource: you describe a head group and a worker group in one manifest, and KubeRay builds and maintains the cluster from it.

All nodes sit in private subnets of one virtual private cloud (VPC) and reach the internet only through a NAT Gateway, which is how every node pulls the DLC from Amazon ECR Public and how workers download model weights from Hugging Face. One m7i.xlarge forms the system node group, hosting kube-system add-ons, the KubeRay operator and the Ray head. The head runs the job server on port 8265 and Ray’s control store on 6379, and it advertises no GPUs, so no training work lands on it. Jobs are submitted from inside the head pod with kubectl exec , through the EKS API, so port 8265 is never exposed outside the cluster.

This sample asks Ray Train for FSDP , which splits the model’s weights, gradients and optimizer state across all eight GPUs rather than copying them to each one. Qwen2.5-1.5B in bfloat16 needs about 12 GB for those four tensors. Sharded eight ways, each L4 GPU holds roughly 1.5 GB. The price is communication: FSDP gathers each layer’s weights before using it and reduce-scatters gradients afterwards, so the whole model crosses between ranks on every step. Between the two instances, that traffic is what the DLC’s EFA stack carries.

Configuration lives in one place, scripts/env.sh : cluster name, Region, Kubernetes version, node types and counts, the image, pinned versions of KubeRay and both device plugins, and the training parameters. Export any of them to change it for a run. All five scripts are idempotent and safe to re-run. deploy_cluster.sh , deploy_node_group.sh , and deploy_ray_train_job.sh each prompt before creating or touching cluster resources. install_gpu_plugins.sh and install_kuberay.sh are reconcilers and run straight through.

Figure 2: The sample’s cluster. Every pod runs the same Ray Train DLC from Amazon ECR Public. The CPU system node runs the KubeRay operator and a Ray head with no GPUs. Each EFA-enabled g6.12xlarge, both in one Availability Zone, runs a single four-GPU Ray worker pod.

Training runs on two g6.12xlarge instances, each with four NVIDIA L4 GPUs, 48 vCPUs, 192 GiB of memory, 40 Gbps of networking and one EFA interface. Each runs one Ray worker pod that claims all four GPUs and the EFA interface, and Ray Train places four workers in each pod: ranks 0–3 on one instance, 4–7 on the other. Because EFA traffic doesn’t cross Availability Zones, the scripts place both instances in a single zone that has been checked to offer g6.12xlarge.

g6.12xlarge capacity shifts between Regions and Availability Zones day to day, so this is the starting point rather than a Region guessed up front in env.sh . It first runs find_gpu_capacity.sh , which checks, per candidate Region: whether g6.12xlarge is offered in each Availability Zone, whether the account’s On-Demand G/VT-family vCPU quota has room for the requested node count, and the EC2 Spot placement score as a capacity signal. With --probe , which deploy_all.sh passes by default, it gets a definitive answer by creating and immediately canceling a real Capacity Reservation in each candidate zone. It then shows the plan and asks once:

A single tagged image runs as both the Ray head and the Ray workers, whether KubeRay starts it on Amazon Elastic Kubernetes Service (Amazon EKS) or you start it yourself on Amazon Elastic Compute Cloud (Amazon EC2). What you add is configuration: which instances, how many GPUs each pod gets, which network transport NCCL uses, and which extra Python packages a job needs. All of it is set from outside the image.

This post walks through a sample that uses the DLC exactly as published to fine-tune Qwen2.5-1.5B with PyTorch Fully Sharded Data Parallel (FSDP) across eight NVIDIA L4 GPUs on two g6.12xlarge instances. The training code is one function, the cluster is one RayCluster manifest plus a file of variables, and no image is built at any point. Because the DLC already carries the EFA stack, collectives between the two instances run over the EFA as soon as the pods request it. The entire solution is in this sample repository , and it uses the same KubeRay-on-EKS foundation as our multi-node Ray Serve walkthrough .

The Ray Train DLC is the PyTorch training image with Ray’s training libraries on top. The tag this sample pulls, public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1 , replaces a build that would otherwise look like this: