AWS releases Ray Train DLC to simplify distributed AI training on EKS

Amazon Web Services (AWS) has released the Ray Train Deep Learning Container (DLC), a pre-configured image designed to streamline distributed machine learning training on Amazon Elastic Kubernetes Service (EKS). Published in October 2026, this container bundles Ray, PyTorch, and the necessary networking libraries for high-performance computing, allowing engineers to deploy multi-node training jobs without managing complex dependency chains.

What happened

The primary bottleneck in scaling PyTorch training across clusters is often the container infrastructure rather than the model code itself. Engineers typically spend significant time aligning CUDA runtimes, PyTorch builds, the NVIDIA Collective Communications Library (NCCL), and Elastic Fabric Adapter (EFA) userspace libraries. Mismatches in these components can cause silent performance degradation, where NCCL falls back to slower TCP sockets instead of using high-speed EFA connections. To address this, AWS has assembled, tested, and published a unified image in Amazon Elastic Container Registry (ECR) Public.

This new DLC is built and maintained by AWS, which patches the stack and validates it through multi-node EFA and KubeRay training tests. The release includes Ray 2.58.0, PyTorch 2.13.0, CUDA 13.0.2, and NCCL 2.29.7, along with the AWS OFI NCCL plugin that connects NCCL to EFA. By providing this as a public image, AWS removes the need for teams to build and maintain their own container images for every update to these underlying libraries.

The accompanying technical guide demonstrates how to use this DLC to fine-tune the Qwen2.5-1.5B model using PyTorch Fully Sharded Data Parallel (FSDP) across eight NVIDIA L4 GPUs. The setup uses two g6.12xlarge instances in a single Availability Zone, managed by the KubeRay operator. This configuration allows inter-node collectives to run over EFA immediately, as the DLC already carries the required networking stack.

How it works

The Ray Train DLC functions as a single, role-agnostic image that serves both the Ray head and worker pods. On Kubernetes, the KubeRay operator uses a RayCluster manifest to create the necessary pods, pointing workers to the head node for coordination. Because the image contains all necessary drivers and libraries, users do not need to build custom images. Instead, they configure the cluster through environment variables and resource requests in the Kubernetes manifest.

Networking is handled automatically through pre-configured environment variables. The DLC defaults to using EFA for communication via the FI_PROVIDER=efa setting and excludes loopback interfaces from NCCL configuration. If a job requires additional Python packages not included in the base image, Ray’s runtime_env feature installs them at submit time, keeping the base image unchanged. This separation ensures that the container remains stable while allowing flexibility for specific experimental needs.

Inside the container, Ray Train manages the distributed process group formation. Engineers write their training loop as a standard Python function, which Ray executes across all available GPUs. For FSDP workloads, Ray splits model weights, gradients, and optimizer states across GPUs, reducing memory requirements per device. The DLC ensures that the heavy communication traffic generated by FSDP gather and reduce-scatter operations travels over the high-bandwidth EFA interface rather than standard network sockets.

Key details

• Image Tag: The sample uses public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1, which includes Ray 2.58.0 and PyTorch 2.13.0.
• Hardware Support: The reference architecture uses g6.12xlarge instances with NVIDIA L4 GPUs and EFA interfaces, deployed in a single Availability Zone.
• Networking Stack: Includes CUDA 13.0.2, NCCL 2.29.7, and the AWS OFI NCCL plugin, tested to ensure EFA usage instead of TCP fallback.
• Configuration Method: Cluster shape and transport settings are defined via environment variables and Kubernetes resource requests, not by rebuilding the image.
• Cost Estimate: The sample setup costs approximately $9.50 per hour, covering two GPU instances, the EKS control plane, and a NAT Gateway.
• Driver Compatibility: The container entrypoint automatically enables CUDA forward-compatibility libraries if the host NVIDIA driver is older than the image’s CUDA version.

Why it matters

For software engineers and ML teams, this release significantly reduces the operational overhead of distributed training. Managing compatible versions of CUDA, NCCL, and EFA drivers is a fragile process that often leads to subtle performance issues. By providing a tested, public image, AWS allows teams to focus on model development rather than infrastructure debugging. The ability to pin specific versions or use moving tags for security updates provides flexibility for both reproducible research and production stability.

The integration with KubeRay and EKS also simplifies cluster management. Engineers can define their entire training infrastructure in a single manifest, specifying GPU counts and EFA interfaces as standard Kubernetes resources. This approach leverages existing Kubernetes tooling and avoids the need for custom scripts to manage Ray cluster lifecycle events. The result is a more reliable and scalable path to running large language model fine-tuning jobs on AWS infrastructure.

What you can do

• Review the sample repository: Access the provided code and manifests to understand the structure of a Ray Train job on EKS.
• Check quota availability: Ensure your AWS account has sufficient On-Demand capacity for G-family instances, specifically 96 vCPUs for the two g6.12xlarge nodes.
• Install required tools: Set up AWS CLI v2, eksctl, kubectl, Helm, and envsubst to manage the cluster and submit jobs.
• Test the DLC locally: Pull the public image from ECR to verify compatibility with your local development environment before deploying to the cluster.
• Configure environment variables: Use the provided scripts/env.sh to set cluster parameters such as region, instance type, and image tag.
• Monitor network performance: Enable NCCL_DEBUG=INFO in your manifest to confirm that EFA is being used for inter-node communication during training.