AWS’s Ray Train example runs a Qwen2.5-1.5B PyTorch Fully Sharded Data Parallel (FSDP) job on Amazon EKS using the published Ray Train Deep Learning Container (DLC), with no custom image build. It places eight NVIDIA L4 GPUs across two g6.12xlarge worker instances. To use the sample’s Elastic Fabric Adapter (EFA) path, keep those workers in the same VPC and Availability Zone.

What the sample deploys

A single m7i.xlarge system node hosts the KubeRay operator and a Ray head with no GPUs. Two g6.12xlarge instances run the GPU workers, four GPUs per instance. Ray Train assigns eight workers across the two machines; the sample uses bfloat16 for Qwen2.5-1.5B.

ComponentSample configurationRole
System nodeOne m7i.xlargeHosts the KubeRay operator and GPU-free Ray head.
GPU workersTwo g6.12xlarge instances, each with four NVIDIA L4 GPUs and one EFA interfaceRun the eight distributed training workers.

Prerequisites and EFA placement

Plan for 96 G-family On-Demand vCPUs for the two GPU instances, or an existing On-Demand Capacity Reservation. AWS lists g6.12xlarge as EFA-supported, but the supported instance types available vary by Region, so check regional support before provisioning.

EFA traffic cannot cross Availability Zones or VPCs. Place both GPU workers in the same VPC and Availability Zone. Ordinary IP traffic through the ENA part of an EFA interface is separate and remains routable.

Keep EFA workers in one VPC and Availability Zone

A VPC can span Availability Zones, but the sample’s EFA connection between GPU workers must stay within one of them.

Configure Ray Train and EFA

Conceptual network diagram showing two GPU workers linked by EFA within one VPC and Availability Zone, with crossed-out paths across both boundaries.

The sample uses the published image public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1 unchanged for the Ray head and workers. Its listed software stack is Ray 2.58.0, PyTorch 2.13.0, torchvision 0.28.0, CUDA 13.0.2 and NCCL 2.29.7. AWS notes that vX.Y image tags can move with security rebuilds; a fully qualified vX.Y.Z tag stays on its original release build, while a digest is the option for byte-identical runs.

Each worker pod requests four GPUs, one EFA device, 40 CPUs and 160Gi of memory. The sample sets the same values as limits, and derives Ray’s num-gpus from the Kubernetes GPU request. That alignment helps prevent Kubernetes and Ray from scheduling different amounts of GPU capacity.

The EFA configuration sets FI_PROVIDER=efa, NCCL_NET_PLUGIN=ofi, NCCL_SOCKET_IFNAME=eth0 and NCCL_DEBUG=INFO. Each worker pod also requests vpc.amazonaws.com/efa: 1, which makes the EFA device available to the pod. Without that resource request, NCCL can use sockets instead.

Deploy the cluster and submit the job

Follow the sample’s dependencies in order:

  1. Set the configuration in scripts/env.sh. Choose the cluster, Region, node, image and training variables so the deployment scripts use the intended settings.
  2. Run deploy_all.sh to check capacity and create the cluster and node group. The script probes capacity before provisioning. Use --plan to preview infrastructure actions; --train chains training after infrastructure setup.
  3. Install or verify the GPU and EFA device plugins. The GPU plugin exposes GPU resources to Kubernetes, and the EFA plugin makes the requested EFA device available to worker pods.
  4. Install KubeRay. The operator manages the RayCluster resource inside Kubernetes.
  5. Run deploy_ray_train_job.sh to render and apply the worker manifest. The deployment checks allocatable GPU and EFA resources before starting the RayCluster.
  6. Submit the job from inside the Ray head pod. The sample uses kubectl exec; it does not expose Ray’s port 8265 outside the cluster.

The sample places nodes in private subnets in one VPC and uses a NAT Gateway for downloads, including the published container and model weights from Hugging Face.

Check that NCCL is using EFA

Before waiting for worker pods, confirm that each GPU node reports allocatable GPU and EFA resources. The sample also checks the GPU node label nvidia.com/gpu.present=true and keeps the Kubernetes GPU request aligned with Ray’s num-gpus setting.

Then inspect NCCL’s selected transport. The job script exits with an error if EFA is not on the path. If the EFA device is unavailable to a pod or its resource request is missing, NCCL can fall back to sockets and the job may still run on TCP.

What the training run demonstrates

The code generates random token IDs rather than reading a real training corpus. AWS describes replacing that input with a Ray Data or DataLoader pipeline for a data-driven fine-tune. The example therefore exercises the distributed cluster and training setup, not fine-tuning on a real dataset.

AWS reports that the EFA path doubled throughput over TCP in this sample. That is AWS’s result for the example, not a general performance guarantee.