Ray train EKS fine-tuning is now possible without building a custom training image. AWS has published a walkthrough that fine-tunes the Qwen2.5-1.5B model with PyTorch Fully Sharded Data Parallel across eight NVIDIA L4 GPUs on two Amazon EKS instances, using a prebuilt Ray Train container image, the KubeRay operator and the Elastic Fabric Adapter for inter-node networking. 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. This article explains how those pieces fit together, what setup the approach requires, and where its limits lie.
What stops PyTorch training from scaling out
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, the Elastic Fabric Adapter 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 networking part of that stack is the one most teams would least like to own, because a mismatch there does not produce an error. NCCL drops to TCP sockets, and the job runs to completion at a fraction of the speed. The Ray Train Deep Learning Container, built by AWS, is that image: AWS assembles the stack, tests it as one unit, publishes it in Amazon Elastic Container Registry Public and keeps it patched, gating each release on multi-node EFA and multi-node KubeRay training tests.
The container the sample pulls is public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1. Inside it are Ray 2.58.0 with Train, Tune and Data, PyTorch 2.13.0 with torchvision 0.28.0, CUDA 13.0.2 with cuDNN and NCCL 2.29.7, the latest EFA libraries with the AWS OFI NCCL plugin and GDRCopy, and training libraries including DeepSpeed, flash-attn, Transformer Engine, PyTorch Lightning and Hugging Face libraries.
Three further properties remove work per cluster. The image is public, so nodes pull it without credentials. On startup, its entrypoint turns on bundled CUDA forward-compatibility libraries automatically if the host driver is older than the image’s CUDA version. And the image has no fixed role: whatever starts the container supplies the ray start command, so one tag serves every pod.
Ray Train EKS fine-tuning: the cluster layout
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, together with 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 roles is a pod, but Kubernetes has no built-in idea of a Ray cluster. 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 is an open source Kubernetes operator that adds a RayCluster resource, and the operator builds and maintains the cluster from one manifest describing a head group and a worker group.
The sample cluster has one EKS cluster, a CPU system node running the KubeRay operator and the Ray head, and two EFA-enabled g6.12xlarge nodes in one Availability Zone. One m7i.xlarge forms the system node group, hosting system add-ons, the KubeRay operator and the Ray head, which runs the job server on port 8265 and the control store on port 6379 and 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 instance 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 to 3 on one instance and ranks 4 to 7 on the other. Because EFA traffic does not cross Availability Zones, both instances sit in a single zone checked to offer g6.12xlarge. All nodes sit in private subnets of one virtual private cloud and reach the internet only through a NAT Gateway, which is how every node pulls the container image and how workers download model weights from Hugging Face.
What Ray Train does inside the container
The container 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.
The sample asks Ray Train for Fully Sharded Data Parallel, or 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 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 container’s EFA stack carries.
Ray Train’s documented pattern is a training function passed to a trainer with a scaling configuration: the function runs on each worker, and the trainer launches the distributed job. The following is a conceptual illustration of that documented mechanism, not the sample’s code:
trainer = TorchTrainer(
train_func,
scaling_config=ScalingConfig(num_workers=8, use_gpu=True),
)
result = trainer.fit()In the sample cluster, that pattern runs with one worker per GPU across the eight GPUs.
Configure it, don’t rebuild it
Everything that differs between one training job and the next is set around the image. The head or worker role comes from KubeRay rayStartParams per pod group, with the same container image for both and the head running with zero GPUs. GPUs and EFA per pod come from Kubernetes resource requests, four GPUs and one EFA interface per worker pod. Network transport comes from environment variables on the pod or the job. Extra Python packages come from Ray runtime_env at submit time. The image version comes from a versioned or digest-pinned tag. Cluster shape comes from variables in a scripts file, and training job parameters come as arguments to the training script.
The image ships the Hugging Face training stack, so the sample’s job needs nothing added. When an experiment needs a library the image does not carry, Ray installs it for that job at submit time through runtime_env and the image is unchanged. The image defaults to FI_PROVIDER=efa, and its NCCL configuration excludes the loopback and Docker bridge interfaces. On EKS, NCCL_SOCKET_IFNAME points at the pod’s eth0. The resource and transport settings below are the values the research establishes; the configuration surface also includes the networking variables NCCL_DEBUG and NCCL_NET_PLUGIN, whose sample values the research does not establish:
resources:
limits: { nvidia.com/gpu: "4", vpc.amazonaws.com/efa: "1" }
requests: { nvidia.com/gpu: "4", vpc.amazonaws.com/efa: "1" }
env:
- { name: FI_PROVIDER, value: "efa" }
- { name: NCCL_SOCKET_IFNAME, value: "eth0" }Read that manifest as three decisions: the unmodified image for head and workers; four GPUs and the node’s single EFA interface per pod, where requesting the EFA resource is what makes the device plugin mount the interface, without which NCCL falls back to sockets; and environment variables that make the stated transport configuration explicit. Because these are plain environment variables, a single job can override them through its own runtime_env.
One allocation detail is worth stating plainly: each worker pod is allocated four GPUs and one EFA interface.
Versioning deserves the same care. The sample pins train-ml-cuda-v1.1, the minor release current at publication. Minor tags move with security rebuilds, so they do not by themselves guarantee an identical run; fully qualified release tags stay frozen at the original release build, and a digest should be pinned when you need a byte-identical run.
Prerequisites before you start
Before starting, you need the tools in place: AWS Command Line Interface version 2 signed in to the target account, eksctl, kubectl, Helm, and envsubst from GNU gettext. You also need 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, that is an alternative to relying on On-Demand quota.
The stated running cost is about $9.50 an hour while the GPU instances are up, covering both instances, the EKS control plane and the NAT Gateway. That is a reason to remove infrastructure you are no longer using, rather than leaving GPU instances running.
Configuration lives in one scripts environment file. The four infrastructure steps take 20 to 30 minutes on a fresh Region.
Running the training job stays a separate step, and its script checks which transport NCCL used, exiting with an error if the job failed or EFA was not in the path, which also makes it usable as a continuous integration gate.
What the sample run actually shows
The sample fine-tunes Qwen2.5-1.5B with FSDP across eight GPUs on two EFA-enabled instances, with no image built at any point. The EFA path the image ships engaged with four environment variables and doubled throughput over TCP on unchanged code. That is the headline result, and it is worth stating its boundary plainly: the source gives the doubling relative to TCP, not an absolute step time or throughput figure, so this article does not invent one.
The sample’s batches are random token IDs, so the run depends on nothing but the cluster. It 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. That is a deliberate design choice for a stack demonstration, and it is also the sample’s main limit as evidence: it proves the container, the cluster layout and the networking path work together, not that a particular model trains well on particular data.
What the sample does prove is practical: the configuration surfaces are separate from the image, so model, job parameters and cluster shape are set outside it. The documented next step for scale is a larger model on a GPUDirect RDMA family such as p5, with the image and code staying the same.
Limits and sensible next steps
The natural next moves follow directly from the sample’s limits. Replace the random batches with a real dataset, and checkpoint to Amazon S3 or Amazon FSx for Lustre, so a run produces something durable. Pin the container image to a frozen release tag, or to a digest for byte-identical runs. Turn on FailureConfig so long runs survive a lost instance. And move to a larger model on a GPUDirect RDMA family such as p5, with the image and code staying the same.
The cost figure is a stated hourly estimate for the whole setup, not a per-run bill, and suitable GPU capacity is a prerequisite that varies by Region and zone, with a capacity reservation as an alternative to On-Demand quota. The pinned tag was the current minor release at publication and minor tags move with security rebuilds, so check the current release before reusing it.
For teams that already run PyTorch training on Kubernetes, the value here is not a new training algorithm. It is the removal of a whole class of work: keeping a custom image aligned across CUDA, PyTorch, NCCL, EFA and Ray, and debugging the silent TCP fallback when that alignment slips. A prebuilt, tested image, a single RayCluster manifest and four environment variables replace that work, and the sample shows the stack carrying real FSDP collectives across two instances.
