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

该示例解决的是 PyTorch 训练上集群时容器难以维护的问题：每个容器都要对齐 CUDA 运行时、对应编译的 PyTorch、NCCL、EFA 用户态库、AWS OFI NCCL 插件和 KubeRay 支持的 Ray 版本。AWS 把这些组件作为一个整体构建、测试并发布在 Amazon ECR Public，示例直接拉取 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 插件、DeepSpeed、flash-attn、Transformer Engine、PyTorch Lightning 和 Hugging Face 库，all_reduce_perf 已在路径中。AWS 每次发布都要通过多节点 EFA 与多节点 KubeRay 训练测试。镜像为公开仓库，节点无需凭证即可拉取；入口脚本会检查宿主机 NVIDIA 驱动，若旧于镜像 CUDA 则自动启用前向兼容库；同一标签既作 head 也作 worker，角色由 ray start 命令决定。 集群由一个 m7i.xlarge 系统节点和两台 g6.12xlarge 组成，全部位于同一可用区私有子网，经 NAT Gateway 访问外网拉取镜像与 Hugging Face 权重。系统节点运行 KubeRay operator 与不占 GPU 的 Ray head，作业通过 kubectl exec 在 head pod 内提交，8265 端口不对外暴露。每台 g6.12xlarge 有 4 张 L4、48 vCPU、192 GiB 内存、40 Gbps 网络和一个 EFA 接口，各跑一个占用全部 4 卡与 EFA 的 worker pod，Ray Train 在其中放置 4 个 worker，rank 0–3 与 4–7 分居两机。 训练代码只有一个 train_func，由 TorchTrainer 按 GPU 数量分发，自动组建 torch.distributed 进程组。示例采用 FSDP，将权重、梯度与优化器状态切分到 8 张卡上：Qwen2.5-1.5B 以 bfloat16 计这四类张量约需 12 GB，八路切分后每张 L4 约 1.5 GB；代价是每步都要在 rank 间聚合权重并 reduce-scatter 梯度，跨机流量由 EFA 承载。按参数规模而非特定层类做自动包装，使脚本与模型无关。 运行前需准备 AWS CLI v2、eksctl、kubectl、Helm 与 envsubst，并为两台 g6.12xlarge 申请 96 vCPU 的 G 系列按需配额，已有容量预留时可改设 CAPACITY_RESERVATION_ID。GPU 实例运行期间成本约每小时 9.50 美元，含两台实例、EKS 控制平面与 NAT Gateway。网络传输、每 pod 的 GPU 与 EFA 数量、额外 Python 包和镜像版本都通过镜像外的配置设定，需要新库时由 Ray 的 runtime_env 在提交时安装，镜像保持不变。