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. 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 . What you get without building anything 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: Layer What you would otherwise build and keep aligned In the Ray Train DLC Distributed runtime Ray with the dashboard and job server, at a version KubeRay supports Ray 2.58.0 with Train, Tune and Data Framework PyTorch compiled for your CUDA version PyTorch 2.13.0, torchvision 0.28.0 GPU runtime CUDA, cuDNN and NCCL releases that agree with each other CUDA 13.0.2, cuDNN, NCCL 2.29.7 Networking EFA installer, libfabric and aws-ofi-nccl , built against that NCCL EFA latest, AWS OFI NCCL plugin, GDRCopy Training libraries DeepSpeed, flash-attn and Transformer Engine compiled for your CUDA DeepSpeed, flash-attn, Transformer Engine, PyTorch Lightning, Hugging Face libraries Diagnostics NCCL tests built from source all_reduce_perf already on the path The networking row is the one most teams would least like to own. A mismatch there doesn’t produce an error. NCCL drops to TCP sockets, and the job runs to completion at a fraction of the speed. In the DLC, those libraries are built and tested together, and AWS gates each release on multi-node EFA and multi-node KubeRay training tests, which are exactly the paths this sample uses. A few more properties remove work that you would otherwise do per cluster: No registry setup – The image is public, so nodes pull it without credentials and nothing is built or pushed on the way. No driver matching – On startup the entrypoint checks the host’s NVIDIA driver and, if it is older than the image’s CUDA, turns on the bundled CUDA forward-compatibility libraries automatically. No head-versus-worker images – The image has no fixed role. Whatever starts the container supplies the ray start command, so one tag serves every pod. No separate inference image family – Ray Train shares the ray ECR repository with the Ray Serve DLC and Ray LLM DLCs , under the train-ml tag prefix. Configure it, don’t rebuild it Everything that differs between one training job and the next is set around the image. The sample uses each of these mechanisms: What you control How you set it In this sample Head or worker role KubeRay rayStartParams per pod group Same DLC_IMAGE for both. Head runs with num-gpus: "0" GPUs and EFA per pod Kubernetes resource requests nvidia.com/gpu: 4 , vpc.amazonaws.com/efa: 1 Network transport Environment variables on the pod or the job FI_PROVIDER , NCCL_NET_PLUGIN , NCCL_SOCKET_IFNAME , NCCL_DEBUG Extra Python packages Ray runtime_env at submit time Any package a job needs, installed at submit time Image version Versioned or digest-pinned tag in DLC_IMAGE train-ml-cuda-v1.1 Cluster shape Variables in scripts/ env.sh Instance type, node count, GPUs per node, Region, AZ Training job Arguments to train.py Model, steps, sequence length, batch size Packages per job. The image ships the HuggingFace training stack, so this sample’s job needs nothing added. When an experiment does need a library the image does not carry, Ray installs it for that job at submit time through runtime_env and the image is unchanged. No new image, no new tag. Transport per job. The image defaults to FI_PROVIDER=efa , and its NCCL configuration excludes the loopback and Docker bridge interfaces. On EKS you point NCCL_SOCKET_IFNAME at the pod’s eth0 . The sample sets all four networking variables in the manifest anyway, so the behavior is visible in one file and doesn’t depend on image defaults. Because these are plain environment variables, a single job can override them through its runtime_env . Versions you choose. The sample pins train-ml-cuda-v1.1 , the current minor release, so a cluster is reproducible from one run to the next. The vX.Y tags move with security rebuilds. The fully qualified vX.Y.Z tags stay frozen at the original release build, so pin a digest when you need a byte-identical run. The same image is also available from private ECR URIs. How the cluster is laid out 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. 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. Training runs on two g6.12xlarge instances, each with four NVI