Distributed training with Ray Train DLCs and EFA on EKS

Building and maintaining custom Deep Learning Containers (DLCs) for large-scale distributed machine learning is often an operational bottleneck. When scaling LLM fine-tuning tasks such as training models like Qwen2.5-1.5B or DeepSeek-V3, managing CUDA drivers, NCCL versions, and AWS Elastic Fabric Adapter (EFA) libraries manually can quickly derail engineering velocity. Fortunately, pre-built Ray Train DLCs ship with the entire CUDA-through-EFA stack natively integrated, allowing developers to spin up enterprise-grade clusters without writing a single custom Dockerfile.

In this technical review, we will walk through a fully reproducible pipeline that leverages Ray Train, PyTorch Fully Sharded Data Parallel (FSDP), and EFA across eight GPUs on two g6.12xlarge instances running on Amazon EKS. Furthermore, we will examine how integrating high-performance LLM APIs like those from n1n.ai can streamline your development and evaluation loops.

Architecture Overview: Ray Train, PyTorch FSDP, and EFA on EKS

To achieve linear scaling across multi-node Kubernetes clusters, our infrastructure relies on a tight integration of AWS networking hardware and Ray orchestration primitives.

+-------------------------------------------------------+
|                   Amazon EKS Cluster                  |
|                                                       |
|  +-----------------------+   +---------------------+  |
|  |   g6.12xlarge (Node 1)|   | g6.12xlarge (Node 2)|  | 
|  |  - 4x NVIDIA L4 GPUs  |   | - 4x NVIDIA L4 GPUs |  | 
|  |  - EFA Network Devices|   | - EFA Network Devices| | 
|  +-----------^-----------+   +----------^----------+  | 
|              |                          |             | 
|              +------------+-------------+             | 
|                           |                           | 
|                   Ray Train Cluster                   | 
+-------------------------------------------------------+

Each g6.12xlarge instance on Amazon EKS provides four NVIDIA L4 GPUs and ultra-low-latency EFA network interfaces. By utilizing the official Ray Train DLC, the underlying MPI, NCCL, and libfabric layers are pre-configured to utilize EFA kernel bypass capabilities directly out of the box.

Prerequisites and Infrastructure Provisioning

Before deploying our distributed training job, ensure your environment meets the following baseline requirements:

1. An active Amazon EKS cluster (version 1.28 or later) with Karpenter or Managed Node Groups.
2. The AWS EFA Kubernetes device plugin installed in the kube-system namespace.
3. The Ray Operator deployed via Helm for cluster lifecycle management.
4. Access to a robust LLM API aggregator like n1n.ai for continuous model evaluation and synthetic data generation during preparation phases.

Validating EFA Interfaces

Before submitting your Ray job, verify that your Kubernetes pods can detect the EFA devices. You can run a quick diagnostic check inside a test pod using the following command:

kubectl run efa-test --rm -it --image=amazon/aws-efa-installer:latest --command -- fi_info -p efa

If configured correctly, the output should list the active efa provider without throwing initialization errors.

Preparing the Ray Train Cluster Configuration

We define our Ray cluster using a standard Kubernetes Custom Resource. This configuration provisions a head node and a worker node, binding them tightly to our EFA-enabled instance profiles.

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ray-train-cluster
spec:
  rayVersion: '2.35.0'
  headGroupSpec:
    rayStartParams: {}
    template:
      spec:
        containers:
          - name: ray-head
            image: public.ecr.aws/ray/ray:2.35.0-py310-gpu
            resources:
              limits:
                cpu: '16'
                memory: 64Gi
                nvidia.com/gpu: '0'
  workerGroupSpecs:
    - groupName: worker-group
      replicas: 1
      minReplicas: 1
      maxReplicas: 2
      rayStartParams: {}
      template:
        spec:
          containers:
            - name: ray-worker
              image: public.ecr.aws/ray/ray:2.35.0-py310-gpu
              securityContext:
                capabilities:
                  add: ["IPC_LOCK"]
              resources:
                limits:
                  cpu: '48'
                  memory: 192Gi
                  nvidia.com/gpu: '4'
                  aws.amazon.com/efa: '4'

Implementing the Qwen2.5-1.5B Fine-Tuning Script with PyTorch FSDP

Next, we write our training script using Ray Train's TorchTrainer. PyTorch FSDP partitions model parameters, gradients, and optimizer states across our eight GPUs, drastically reducing memory overhead per device.

import os
import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer
from ray.train.torch import TorchTrainer
from ray.train import ScalingConfig, CheckpointConfig

def train_func(config):
    model_id = config.get("model_id