Distributed training with Ray Train DLCs and EFA on EKS

wpnews.pro

Artificial Intelligence #

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.

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:

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 theray start command, so one tag serves every pod.
No separate inference image family – Ray Train shares theray ECR repository with theRay Serve DLC andRay LLM DLCs , under thetrain-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:

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 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 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, 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, andenvsubst 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, setCAPACITY_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 Data 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<

source & further reading

aws.amazon.com — original article Building an Agentic AI Platform on Amazon Bedrock AgentCore Announcing SOCI-indexed images for faster SageMaker AI jobs Fix evaluation bias to unlock better models on Amazon Bedrock