AWS Details Distributed Training on Amazon EKS with Ray Train
AWS highlights how the Ray Train Deep Learning Container helps streamline distributed PyTorch training across multiple nodes on Amazon EKS using Elastic Fabric Adapter.
Radwan AbdulhadiEditor and technology writer
Overcoming Distributed Training Container Challenges
Scaling PyTorch training across a cluster has traditionally been hindered by container configuration complexity rather than model code itself. Every container requires matching versions of the CUDA runtime, a compiled PyTorch build, NCCL, Elastic Fabric Adapter (EFA) userspace libraries, the AWS OFI NCCL plugin, and a compatible Ray version. Maintaining and rebuilding these images whenever a component changes introduces significant overhead.
To address this operational burden, AWS provides the Ray Train Deep Learning Container. AWS assembles, tests, and publishes the complete stack as a single unit in Amazon Elastic Container Registry (Amazon ECR) Public, allowing engineering teams to run distributed workloads using the container exactly as published without building custom images.
Cluster Architecture with KubeRay and Amazon EKS
Deploying Ray clusters on Amazon Elastic Kubernetes Service (Amazon EKS) relies on KubeRay, an open-source Kubernetes operator. KubeRay introduces a RayCluster resource that enables operators to define both head and worker groups within a single manifest. The operator automatically builds, maintains, and replaces failed pods to keep the cluster functioning reliably.
In a typical configuration, a CPU system node runs the KubeRay operator and the Ray head, while dedicated GPU instances handle training execution. The Ray head coordinates the global control store, dashboard, and job server without handling training tasks directly.
Fine-Tuning Qwen2.5-1.5B with PyTorch FSDP
A reference sample demonstrates how to fine-tune the Qwen2.5-1.5B model using PyTorch Fully Sharded Data Parallel (FSDP) across eight NVIDIA GPUs on two EFA-enabled g6.12xlarge instances. Because the container already includes the necessary networking stack, multi-node collectives route efficiently over EFA without requiring manual network configuration inside the image.
The training code is encapsulated in a single function paired with a RayCluster manifest and a variable file. Because the environment variables and runtime settings are defined externally, teams can customize package dependencies or network transports per job without modifying the underlying base image.
Operational Benefits of Pre-Built DLCs
Using published Deep Learning Containers eliminates several common friction points in machine learning infrastructure. Nodes pull the public image without requiring registry credentials, and the container automatically checks host NVIDIA drivers against bundled CUDA forward-compatibility libraries upon startup.
Additionally, the containers avoid separate role definitions by allowing the runtime environment to supply the start command. AWS tests these releases extensively against multi-node EFA and KubeRay configurations, helping ensure consistent performance and preventing silent fallbacks to slower TCP sockets.
Sources
- AWS Machine Learning BlogDistributed training with Ray Train DLCs and EFA on EKS
Continue chronologically
Related entity coverage
Related stories
TorchNeuron Brings Native PyTorch Support to AWS Trainium
TorchNeuron arrives as a public beta, introducing a native PyTorch backend for AWS Trainium with eager execution and torch.compile support.
Amazon Launches TorchNeuron to Challenge Nvidia in AI
Amazon Web Services has released TorchNeuron, a new development tool designed to help AI startups run applications on custom Trainium chips.
Train Qwen3 on AWS Trainium with torchtitan
Learn how to train Qwen3 on AWS Trainium using torchtitan across three distinct configuration approaches on a single trn2.48xlarge instance.
Hugging Face Launches Space to Check AWS Trainium Support
Discover which Hugging Face models run on AWS Trainium using the new torch-neuronx model support space, checking compatibility, scripts, and more.
