在 EKS 上用 Ray Train DLC 与 EFA 做分布式训练

在 Kubernetes 上把 PyTorch 训练扩展到多机，卡点往往不在模型代码，而在容器：CUDA 运行时、对应编译的 PyTorch、NCCL、EFA 用户态库、连接 NCCL 与 EFA 的 AWS OFI NCCL 插件，以及集群认可的 Ray 版本，任何一环变动都要重建镜像。AWS 给出的方案是 Ray Train Deep Learning Container（DLC）：由 AWS 组装并整体测试这套栈，发布在 Amazon ECR Public 并持续打补丁。

示例做了什么

示例直接使用已发布的 DLC，不做任何镜像构建，在 Amazon EKS 上跨两台 g6.12xlarge 实例、共 8 张 NVIDIA L4 GPU，用 PyTorch FSDP 微调 Qwen2.5-1.5B。集群由一个 RayCluster manifest 加一个变量文件描述，训练代码是一个函数。

架构上，一个 EKS 集群包含 AWS 托管的控制平面、一台运行 KubeRay operator 与 Ray head 的 CPU 系统节点（m7i.xlarge），以及同一可用区内两台支持 EFA 的 g6.12xlarge 节点。KubeRay 是开源 Kubernetes operator，新增 RayCluster 资源，按 manifest 中的 head 组与 worker 组创建并维护 pod；Ray Train 放置 8 个 worker，每 GPU 一个，跨节点集合通信走 EFA。

镜像里已包含什么

示例拉取的 tag 为 public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1，对应组件包括：Ray 2.58.0（含 Train、Tune、Data）、PyTorch 2.13.0 与 torchvision 0.28.0、CUDA 13.0.2 与 NCCL 2.29.7、EFA 最新版与 AWS OFI NCCL 插件、GDRCopy，以及 DeepSpeed、flash-attn、Transformer Engine、PyTorch Lightning 和 Hugging Face 库，all_reduce_perf 已在路径中。

网络这一层最容易被忽视：配置不匹配不会报错，NCCL 会退回 TCP socket，任务照常跑完但速度大幅下降。AWS 称每个版本都以多节点 EFA 和多节点 KubeRay 训练测试作为发布门槛。

镜像还省掉几项按集群重复的工作：公开镜像无需凭证拉取；启动时入口脚本检查主机 NVIDIA 驱动，若旧于镜像 CUDA 则自动启用内置的 CUDA 前向兼容库；同一 tag 既作 head 也作 worker，角色由启动方传入的 ray start 命令决定；Ray Train 与 Ray Serve、Ray LLM 的 DLC 共用 ray 仓库，位于 train-ml 前缀下。

差异靠配置而非重建

任务间的差异全部在镜像外设置：head 与 worker 角色由 KubeRay 的 rayStartParams 指定，head 设 num-gpus: "0"；GPU 与 EFA 通过 Kubernetes 资源请求声明，示例为 nvidia.com/gpu: 4、vpc.amazonaws.com/efa: 1；网络传输通过 FI_PROVIDER、NCCL_NET_PLUGIN、NCCL_SOCKET_IFNAME、NCCL_DEBUG 等环境变量控制；额外 Python 包在提交时用 Ray runtime_env 安装；镜像版本用带版本或 digest 固定的 tag；集群形态由 scripts/env.sh 中的变量决定，包括实例类型、节点数、每节点 GPU 数、Region 与可用区；训练参数作为 train.py 的参数传入。

编辑部观察

这套做法的价值在于把“镜像与驱动对齐”从每个团队的日常工作中剥离，代价是把控制权交给镜像发布方：镜像内的 Ray、PyTorch、CUDA 版本组合由 AWS 决定，团队需要确认其与自有代码和依赖兼容。示例本身是单可用区、两台节点的规模，不能据此推断更大集群或跨可用区场景的表现。

实践上可先核查三点：自有训练代码是否依赖镜像未包含的包或特定版本；EFA 是否真的被 NCCL 使用，可通过 NCCL_DEBUG 输出确认而非只看任务是否完成；镜像 tag 是否按需固定到 digest，避免上游更新带来非预期变化。