ReadingResearchRadarInvestment frameworkAI Company Primary Sources
Sign in / Sign up
← Back to reading
NewsMaterial4 min readDisconfirmAI Editorial

AWS Ray Train DLC: Distributed PyTorch Training Usually Stalls at the Container, Not the Model Code

AWS's pre-built Ray Train containers bundle CUDA, PyTorch, NCCL, EFA userspace libraries, and the AWS OFI NCCL plugin into a single tested unit, enabling multi-node FSDP training to run with only environment variable and resource request configuration, without building any images. The article provides a complete, reproducible method for fine-tuning Qwen2.5-1.5B using eight NVIDIA L4 GPUs across two g6.12xlarge instances.

2026-10-29 02:19 GMT+8

People doing distributed PyTorch training often focus their efforts on model code, but what actually prevents training from scaling from a single machine to a multi-node cluster is usually the container. Each container requires the same set of components: a matching CUDA runtime, a PyTorch build compiled against that CUDA version, NCCL (NVIDIA Collective Communications Library), Elastic Fabric Adapter (EFA) userspace libraries, the AWS OFI NCCL plugin that connects NCCL to EFA, and a Ray version approved by the cluster operator. Building this image and rebuilding it whenever any one of these components changes is an engineering project in itself, entirely unrelated to the model you intend to train.

AWS has packaged this stack into the Ray Train Deep Learning Container (DLC), publishing it on Amazon ECR Public and keeping it patched. The article runs an example using this image as-is: fine-tuning Qwen2.5-1.5B with PyTorch FSDP across eight NVIDIA L4 GPUs on two g6.12xlarge instances. The training code consists of a single function, the cluster is defined by a RayCluster manifest plus a variables file, and no image is built at any point.

The networking stack is the part most teams least want to maintain themselves. The key risk here is not an error message, but silent degradation: if the EFA userspace library and NCCL versions do not match, NCCL falls back to TCP sockets. The job completes successfully, but at a fraction of normal speed. The DLC builds and tests these libraries together; AWS gates each release with multi-node EFA and multi-node KubeRay training tests, which are exactly the paths used in this example.

Sign in / Sign up
ReadingResearchRadarInvestment frameworkAI Company Primary Sources

All differences outside the image are resolved through configuration rather than rebuilding. The cluster layout is described by a KubeRay manifest: a CPU system node runs the KubeRay operator and the Ray head (requesting no GPU), while two g6.12xlarge nodes each run a worker pod requesting four GPUs and one EFA interface. Network transport is determined by four environment variables: FI_PROVIDER is set to efa so libfabric uses only EFA; NCCL_NET_PLUGIN is set to ofi to load the AWS OFI plugin; NCCL_SOCKET_IFNAME points to the pod's eth0; and NCCL_DEBUG is set to INFO to verify the actual transport chosen. These variables are written into the manifest, making behavior visible in a single file and independent of image defaults.

Inside the container, Ray Train handles the launcher work. Multi-node PyTorch typically means starting a process per GPU on every host and assigning each process a rank, world size, and rendezvous address. With Ray Train, you write the training loop as a function and pass it to TorchTrainer; Ray automatically starts one worker per GPU wherever the cluster has GPUs, forms torch.distributed process groups, and runs the function on each worker. The example requests an FSDP strategy, sharding model weights, gradients, and optimizer state across all eight GPUs. For Qwen2.5-1.5B in bfloat16, these four tensors occupy about 12 GB; sharded eight ways, each L4 GPU holds approximately 1.5 GB.

The cost is communication: FSDP gathers layer weights before use and reduce-scatters gradients afterward, meaning the entire model crosses between ranks on every step. This inter-instance traffic is carried by the EFA stack included in the DLC. The article concludes with an attribution result: after enabling the EFA path via four environment variables, throughput doubled compared to TCP with unchanged code. This is an AWS demonstration conclusion; the example uses random token batches, and the evaluation criteria are whether FSDP collectives actually crossed both nodes and whether the task completed, not training convergence quality.

The value of this entire solution lies in removing container building from the training workflow. The image is public, allowing nodes to pull it without credentials; the entrypoint script checks the host NVIDIA driver at startup and automatically enables bundled CUDA forward-compatibility libraries if the host driver is older than the image's CUDA. The same tag runs both head and worker pods with no role distinction. When additional Python packages are needed, they are installed via Ray's runtime_env at submission time, leaving the image unchanged. To change models, adjust steps, modify sequence length or batch size, only the parameters in the variables file need changing. Switching GPU instance families requires only confirming that the new family supports EFA and is on the device plugin allowlist; neither code nor image needs modification.

Sources

  • aws.amazon.comAmazon Web Services↗
Continue researching →