Distributed training with Ray Train DLCs and EFA on EKS

Distributed Training on EKS Made Simple with Ray Train DLCs and EFA

Running large‑scale PyTorch jobs on Kubernetes often stalls at the container level: you must keep CUDA, NCCL, the Elastic Fabric Adapter (EFA) stack, and a compatible Ray version perfectly aligned. The Ray Train Deep Learning Container (DLC) eliminates that friction by providing a pre‑built, fully tested image that includes everything needed for multi‑node training on Amazon EKS.

---

1. Why Use the Ray Train DLC?

| Feature | What It Means for You | |---|---| | All‑in‑one stack | The image bundles a matching CUDA runtime, cuDNN, NCCL, the EFA userspace libraries, the AWS OFI NCCL plugin, and Ray 2.58 with Train, Tune, and Data. No manual version juggling. | | Public ECR image | Pulls without credentials, so you never need to maintain a private registry or push custom builds. | | Driver‑forward compatibility | On startup the entrypoint checks the host NVIDIA driver; if it’s older than the image’s CUDA, bundled forward‑compatibility libraries are enabled automatically. | | Single image for head and workers | The same tag runs both the Ray head and the GPU workers, simplifying manifests and avoiding duplicate images. | | Runtime‑environment package installs | Any extra Python dependency can be installed at job submit time via `runtime_env`, keeping the base image immutable. |

By offloading the networking and driver concerns to the DLC, you focus on the model code and experiment configuration.

---

2. Setting Up the Cluster with KubeRay

1. Create an EKS cluster with at least one CPU‑only node for the KubeRay operator and one or more EFA‑enabled GPU nodes (e.g., `g6.12xlarge`).
2. Install the KubeRay operator using the official Helm chart. It watches a custom `RayCluster` resource and creates the head and worker pods automatically.
3. Define a `RayCluster` manifest that references the DLC image (`public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1`). Example snippets:

```yaml apiVersion: ray.io/v1 kind: RayCluster metadata: name: qwen-train spec: headGroupSpec: serviceType: ClusterIP rayStartParams: num-gpus: "0" template: spec: containers: - name: ray-head image: public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1 env: - name: FIPROVIDER value: efa - name: NCCLSOCKET_IFNAME value: eth0 workerGroupSpecs: - replicas: 2 minReplicas: 2 maxReplicas: 2 groupName: gpu-workers rayStartParams: num-gpus: "4" template: spec: containers: - name: ray-worker image: public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1 resources: limits: nvidia.com/gpu: 4 vpc.amazonaws.com/efa: 1 ```

1. Submit a training job using `ray job submit` or the Ray Dashboard. Provide a `runtime_env` block if you need extra packages:

```json { "runtime_env": { "pip": ["transformers==4.40.0", "accelerate"] } } ```

The operator ensures that if any pod crashes, it is recreated and re‑joined to the Ray head, preserving cluster state.

---

3. Running a Real‑World Example – Fine‑Tuning Qwen2.5‑1.5B

The sample repository (linked in the original blog) demonstrates fine‑tuning the Qwen2.5‑1.5B model with PyTorch Fully Sharded Data Parallel (FSDP) across eight L4 GPUs spread over two `g6.12xlarge` nodes. Key takeaways:

• No custom Dockerfile – the DLC already contains PyTorch 2.13, torchvision, DeepSpeed, flash‑attn, and the Transformer Engine compiled for CUDA 13.
• EFA‑accelerated collectives – NCCL automatically uses the EFA fabric because the image ships the AWS OFI NCCL plugin and the manifest sets `FI_PROVIDER=efa`.
• Reproducibility – Pin the image tag (`train-ml-cuda-v1.1`) or a digest to guarantee byte‑identical runs across experiments.
• Scalable resource spec – Adjust `replicas`, `nvidia.com/gpu`, and `vpc.amazonaws.com/efa` in the manifest to scale up or down without touching the container.

Running the job is as simple as: ```bash ray job submit \ --address http://<head-service>:8265 \ --runtime-env-json '{"pip": ["datasets", "evaluate"]}' \ -- python train.py --model qwen2.5-1.5b --batch-size 8 --steps 500 ``` All communication between the two nodes travels over the low‑latency EFA network, delivering near‑bare‑metal throughput.

---

4. Best Practices and Gotchas

• Validate EFA availability: Ensure the subnet has the `EFA` service enabled and that the instance type supports it.
• Monitor NCCL diagnostics: The DLC includes `allreduceperf`. Running it inside a pod can quickly confirm that NCCL is using EFA rather than falling back to TCP.
• Pin image digests for production: Tags like `v1.1` receive security patches; for strict reproducibility, use the immutable digest (`sha256:…`).
• Leverage Ray Dashboard: It provides real‑time visibility into worker health, GPU utilization, and network traffic, helping you spot bottlenecks early.
• Avoid mixing driver versions: If the host driver is newer than the image’s CUDA, the forward‑compatibility shim works, but the reverse (older driver, newer CUDA) will cause failures.

---

Closing Thoughts

The combination of Ray Train DLCs, EFA, and KubeRay turns what used to be a multi‑day container‑building exercise into a few YAML edits and a single `ray job submit` command. By relying on AWS‑maintained images that are continuously tested on multi‑node EFA setups, data science teams can iterate faster, scale more reliably, and keep their focus on model innovation rather than low‑level stack compatibility.