AWS Ray Train DLCs and EFA on EKS simplify distributed PyTorch training

Simplifying Distributed Training Container Management

Scaling PyTorch training across a cluster typically involves complex container setup, requiring specific versions of CUDA, PyTorch, NCCL, Elastic Fabric Adapter (EFA) libraries, and the AWS OFI NCCL plugin. AWS addresses this by providing Ray Train Deep Learning Containers (DLCs). These DLCs are pre-built container images that include all necessary components, tested as a single unit, and published in Amazon ECR Public.

The DLCs eliminate the need for users to build and rebuild images for each component update, allowing them to focus on model development rather than infrastructure. AWS maintains and patches these images, ensuring compatibility and stability for distributed training workloads.

Ray Cluster Architecture on Kubernetes

Ray operates a job across machines as a single cluster with distinct roles: a Ray head node for coordination and control, and Ray worker nodes that supply GPUs and run training processes. On Kubernetes, each of these roles is represented by a pod. The open-source KubeRay operator manages the creation and maintenance of these Ray clusters on Kubernetes, including on Amazon Elastic Kubernetes Service (EKS).

KubeRay allows users to define head and worker groups in a single manifest, which it then uses to build and maintain the cluster, replacing failed pods as needed. This abstraction simplifies the deployment and management of Ray clusters within a Kubernetes environment.

Flexible Configuration with Pre-built Images

The same tagged DLC image can function as both the Ray head and Ray workers, whether deployed on Amazon EKS via KubeRay or directly on Amazon EC2. Users configure the environment by specifying instance types, GPU allocation per pod, NCCL network transport, and additional Python packages required for a job.

This approach separates configuration from the image itself, providing flexibility without requiring image rebuilds. The inclusion of the EFA stack within the DLC enables collectives between instances to utilize EFA for high-performance communication as soon as pods request it.

Example Use Case

A sample walkthrough demonstrates fine-tuning Qwen2.5-1.5B with PyTorch Fully Sharded Data Parallel (FSDP) across eight NVIDIA L4 GPUs on two g6.12xlarge instances using the published DLC. The training code is encapsulated in a single function, and the cluster is defined by one RayCluster manifest and a variable file. No custom image building is required for this process.

Impact on Distributed Training

This offering from AWS simplifies the operational overhead associated with distributed deep learning training. By providing pre-configured and maintained container images, AWS reduces the complexity of managing dependencies and ensuring compatibility across various components. This allows developers and researchers to deploy and scale PyTorch training jobs more efficiently on EKS, leveraging high-performance networking like EFA without extensive setup.