Distributed training with Ray Train DLCs and EFA on EKS

AWS published a walkthrough showing that its Ray Train Deep Learning Container (DLC), tag public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1, can fine-tune Qwen2.5-1.5B with PyTorch Fully Sharded Data Parallel (FSDP) across eight NVIDIA L4 GPUs on two g6.12xlarge instances on Amazon EKS without building any custom image. The container bundles Ray 2.58.0 with Train, Tune and Data, PyTorch 2.13.0, torchvision 0.28.0, CUDA, NCCL and the AWS OFI NCCL plugin, so collectives between the two instances run over the Elastic Fabric Adapter (EFA) once the pods request it. The sample, deployed with the open-source KubeRay operator via a single RayCluster manifest plus a variables file, is published in an AWS sample repository and addresses the container-alignment work that typically blocks PyTorch training from scaling to a cluster.

Artificial Intelligence https://aws.amazon.com/blogs/machine-learning/ Distributed training with Ray Train DLCs and EFA on EKS What usually stops PyTorch https://pytorch.org/ training from scaling out to a cluster is the container, not the model code. Every container needs the same CUDA runtime https://docs.nvidia.com/cuda/cuda-programming-guide/01-introduction/cuda-platform.html , a PyTorch build compiled against it, NCCL https://developer.nvidia.com/nccl NVIDIA Collective Communications Library , the Elastic Fabric Adapter EFA https://aws.amazon.com/hpc/efa/ userspace libraries, the AWS OFI NCCL https://github.com/aws/aws-ofi-nccl plugin that connects NCCL to EFA, and a Ray https://docs.ray.io/en/latest/ 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 https://aws.github.io/deep-learning-containers/ray-train/ 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. 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 https://aws.github.io/deep-learning-containers/ray-train/deployment/eks/ 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 https://aws.amazon.com/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 https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa.html as soon as the pods request it. The entire solution is in this sample repository https://github.com/aws-samples/sample-aws-deep-learning-containers/tree/main/training/ray-train , and it uses the same KubeRay-on-EKS foundation as our multi-node Ray Serve walkthrough https://aws.amazon.com/blogs/machine-learning/run-multi-node-inference-workloads-using-ray-llm-deep-learning-containers/ . 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 https://aws.github.io/deep-learning-containers/ray/ and Ray LLM DLCs https://aws.github.io/deep-learning-containers/ray-llm/ , 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 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 https://aws.amazon.com/ec2/instance-types/g6/ 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. What Ray Train does inside the container The DLC provides the stack. Ray Train https://docs.ray.io/en/latest/train/train.html provides the part that would otherwise be launcher code. Multi-node PyTorch normally means starting a process per GPU on every host and giving each one its rank, the world size and a rendezvous address. With Ray Train you write the training loop as a function and pass it to TorchTrainer with a ScalingConfig . Ray starts one worker per GPU wherever the cluster has GPUs, forms the torch.distributed process group, and runs your function on each worker. This sample asks Ray Train for FSDP https://docs.pytorch.org/docs/stable/fsdp.html , 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. Prerequisites You must have the following in place before you start: - Tools: AWS Command Line Interface AWS CLI v2 signed in to the target account, eksctl, kubectl, Helm, and envsubst from GNU gettext. - Quota: 96 vCPUs of G-family On-Demand capacity for the two g6.12xlarge instances, in an AWS Region where at least one zone offers them. If you already hold an On-Demand Capacity Reservation for the instances, set CAPACITY RESERVATION ID instead. The scripts launch the node group into it and skip the capacity search. - Budget: about $9.50 an hour while the GPU instances are up, covering both instances, the EKS control plane and the NAT Gateway. A quick way to confirm that the CLI is pointed at the account that you expect: The training code: One function and a trainer code/train.py is the whole job, and it contains nothing specific to the container, the cluster or the network. Ray runs train func once per GPU. It asks Ray for its rank, loads the model, hands it to prepare model for FSDP, and runs an ordinary PyTorch loop: code/train.py — the per-GPU function abridged The driver joins the running cluster, counts its GPUs and starts the trainer. Because the worker count comes from the cluster, the same script runs unchanged on one node or many: code/train.py — the driver abridged The batches are random token IDs so that the run depends on nothing but the cluster. The sample is judged on whether real FSDP collectives crossed both nodes and the job completed. For a real fine-tune, replace the random tensor with a Ray Data or DataLoader pipeline. Nothing else in the function changes. The manifest: The only place the DLC is configured manifest/raycluster.yaml is where the image meets the cluster, and it is short because the image needs so little. The head’s only special setting is NVIDIA VISIBLE DEVICES=void , which stops the container runtime’s NVIDIA hook from failing on a node that has no GPU. The worker group holds everything else: manifest/raycluster.yaml — worker group abridged Read it as three decisions. Which image: ${DLC IMAGE} , unmodified, for both head and workers. Which hardware: four GPUs and the node’s single EFA interface per pod. Requesting vpc.amazonaws.com/efa is what makes the EFA device plugin mount the interface into the container. Without it, NCCL can’t see EFA and falls back to sockets. Which transport: four environment variables, so nothing is left to auto-detection. Two smaller details keep a multi-node job stable. First, requests should equal limits, which places the workers in Kubernetes’ Guaranteed Quality of Service QoS class. This matters because a single throttled or evicted rank stalls the other seven. Second, num-gpus<