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

通常阻碍 PyTorch 训练扩展到集群的是容器，而不是模型代码。每个容器都需要相同的 CUDA 运行时、针对它编译的 PyTorch 构建、NCCL（NVIDIA 集合通信库）、Elastic Fabric Adapter (EFA) 用户态库、将 NCCL 连接到 EFA 的 AWS OFI NCCL 插件，以及集群运维方认可的 Ray 版本。构建该镜像，并在其中任何一个组件变动时重新构建，本身就是一项工程，而且与你原本要训练的模型毫无关系。Ray Train 深度学习容器（DLC）就是那个镜像，由 AWS 构建。AWS 组装整个技术栈，将其作为一个整体进行测试，发布到 Amazon Elastic Container Registry (Amazon ECR) Public 并持续为其打补丁。

图 1：一个 EKS 集群，其控制平面由 AWS 管理，一个 CPU 系统节点（m7i.xlarge）运行 KubeRay operator 和 Ray head，以及同一可用区内的两个启用 EFA 的 g6.12xlarge 节点。KubeRay 根据一个 RayCluster 清单构建它，Ray Train 放置八个 worker（每个 GPU 一个），节点间集合通信通过 EFA 进行。

Ray 以单一集群、两种角色的方式跨机器运行作业。Ray head 负责协调并运行跟踪集群状态的全局控制存储、仪表板以及你提交作业的作业服务器。Ray worker 节点提供 GPU 并运行训练进程。每个 worker 在启动时通过连接 head 的地址加入集群。在 Kubernetes 上，它们每一个都是一个 pod。Kubernetes 本身没有 Ray 集群的概念，因此必须有东西来创建 head 和 worker pod、在每个 pod 中以正确的角色启动 Ray、将 worker 指向 head，并替换失败的 pod。KubeRay 就是做这个的。它是一个开源 Kubernetes operator，添加了 RayCluster 资源：你在一个清单中描述 head 组和 worker 组，KubeRay 据此构建并维护集群。

同一个带标签的镜像既作为 Ray head 运行，也作为 Ray worker 运行，无论 KubeRay 是在 Amazon Elastic Kubernetes Service (Amazon EKS) 上启动它，还是你自己在 Amazon Elastic Compute Cloud (Amazon EC2) 上启动它。你需要添加的是配置：使用哪些实例、每个 pod 获得多少 GPU、NCCL 使用哪种网络传输，以及作业需要哪些额外的 Python 包。所有这些都在镜像外部设置。

本文介绍一个示例，它完全按发布状态使用 DLC，在两个 g6.12xlarge 实例上跨八个 NVIDIA L4 GPU，用 PyTorch 全分片数据并行（FSDP）微调 Qwen2.5-1.5B。训练代码是一个函数，集群是一个 RayCluster 清单加一个变量文件，整个过程不构建任何镜像。由于 DLC 已包含 EFA 技术栈，只要 pod 请求，两个实例之间的集合通信就会通过 EFA 进行。整个解决方案位于此示例仓库中，它使用了与我们的多节点 Ray Serve 演练相同的 KubeRay-on-EKS 基础。

无需构建任何东西即可获得的内容

Ray Train DLC 是在 PyTorch 训练镜像之上叠加了 Ray 的训练库。此示例拉取的标签 public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1 替代了原本需要像这样构建的镜像：

网络这一行是大多数团队最不愿意自己负责的部分。这里出现不匹配不会产生错误。NCCL 会降级为 TCP socket，作业会以极低的速度跑完。在 DLC 中，这些库是一起构建和测试的，并且 AWS 每次发布都要通过多节点 EFA 和多节点 KubeRay 训练测试，而这正是本示例使用的路径。

还有几个特性可以省去你原本需要按集群逐个完成的工作：

• 无需设置镜像仓库 – 该镜像是公开的，因此节点无需凭据即可拉取，过程中也不需要构建或推送任何内容。
• 无需匹配驱动 – 启动时，入口点会检查主机的 NVIDIA 驱动，如果它比镜像的 CUDA 旧，就会自动启用内置的 CUDA 前向兼容库。
• 无需区分 head 与 worker 镜像 – 该镜像没有固定角色。无论什么启动容器，都会提供 ray start 命令，因此一个标签即可服务所有 pod。
• 无需单独的推理镜像系列 – Ray Train 与 Ray Serve DLC 和 Ray LLM DLC 共享 ray ECR 仓库，位于 train-ml 标签前缀下。

配置它，而不是重建它

每个训练作业之间所有不同的内容都围绕镜像进行设置。本示例使用了以下每种机制：

按作业安装包。 镜像自带 HuggingFace 训练栈，因此本示例的作业无需添加任何内容。当实验确实需要镜像未包含的库时，Ray 会在提交时通过 runtime_env 为该作业安装，镜像保持不变。无需新镜像，无需新标签。

按作业设置传输。 镜像默认使用 FI_PROVIDER=efa，其 NCCL 配置排除了回环和 Docker 网桥接口。在 EKS 上，你将 NCCL_SOCKET_IFNAME 指向 pod 的 eth0。本示例仍在清单中设置了全部四个网络变量，因此行为在一个文件中可见，不依赖镜像默认值。由于这些是普通环境变量，单个作业可以通过其 runtime_env 覆盖它们。

由你选择版本。示例固定使用 train-ml-cuda-v1.1，即当前次要版本，因此集群在每次运行之间都可复现。vX.Y 标签会随安全重建而更新。完全限定的 vX.Y.Z 标签则冻结在原始发布构建上，所以当你需要字节级一致的运行时，请固定摘要。同一镜像也可通过私有 ECR URI 获取。

集群的布局方式

图 2：示例的集群。每个 Pod 都运行来自 Amazon ECR Public 的同一个 Ray Train DLC。CPU 系统节点运行 KubeRay operator 和一个不带 GPU 的 Ray head。每个启用了 EFA 的 g6.12xlarge（两者位于同一个可用区）各运行一个四 GPU 的 Ray worker Pod。

所有节点都位于同一个虚拟私有云（VPC）的私有子网中，并且仅通过 NAT Gateway 访问互联网，每个节点正是通过这种方式从 Amazon ECR Public 拉取 DLC，worker 也正是通过这种方式从 Hugging Face 下载模型权重。一台 m7i.xlarge 构成系统节点组，承载 kube-system 附加组件、KubeRay operator 和 Ray head。head 在端口 8265 上运行 job server，在 6379 上运行 Ray 的控制存储，并且它不对外声明任何 GPU，因此不会有训练工作落到它上面。作业通过 EKS API 在 head Pod 内部使用 kubectl exec 提交，因此端口 8265 从不暴露在集群之外。

训练运行在两台 g6.12xlarge 实例上，每台配备四个 NVIDIA L4 GPU、48 个 vCPU、192 GiB 内存、40 Gbps 网络和一个 EFA 接口。每台运行一个 Ray worker Pod，该 Pod 占用全部四个 GPU 和 EFA 接口，Ray Train 在每个 Pod 中放置四个 worker：一台实例上为 rank 0–3，另一台上为 4–7。由于 EFA 流量不跨可用区，脚本将两台实例放在同一个已确认提供 g6.12xlarge 的可用区中。

Ray Train 在容器内做什么

DLC 提供整套技术栈。Ray Train 提供原本需要由启动器代码完成的部分。多节点 PyTorch 通常意味着要在每台主机上为每个 GPU 启动一个进程，并为每个进程指定其 rank、world size 和一个 rendezvous 地址。使用 Ray Train 时，你把训练循环写成一个函数，并用 ScalingConfig 将其传给 TorchTrainer。Ray 会在集群中有 GPU 的地方为每个 GPU 启动一个 worker，组建 torch.distributed 进程组，并在每个 worker 上运行你的函数。

本示例要求 Ray Train 使用 FSDP，它将模型的权重、梯度和优化器状态拆分到全部八个 GPU 上，而不是复制到每个 GPU。bfloat16 格式的 Qwen2.5-1.5B 需要约 12 GB 来存放这四个张量。按八路分片后，每个 L4 GPU 大约持有 1.5 GB。代价是通信：FSDP 在使用每一层之前收集该层的权重，之后对梯度进行 reduce-scatter，因此整个模型在每一步都会在 rank 之间传输。在两台实例之间，这部分流量正是由 DLC 的 EFA 栈承载的。

前提条件

开始之前，你必须准备好以下内容：

• 工具：已登录目标账户的 AWS Command Line Interface（AWS CLI）v2、eksctl、kubectl、Helm，以及来自 GNU gettext 的 envsubst。
• 配额：为两台 g6.12xlarge 实例准备 96 个 vCPU 的 G 系列按需容量，且位于至少有一个可用区提供这些实例的 AWS 区域。如果你已持有这些实例的按需容量预留，请改为设置 CAPACITY_RESERVATION_ID。脚本会将节点组启动到该预留中并跳过容量搜索。
• 预算：GPU 实例运行期间每小时约 $9.50，涵盖两台实例、EKS 控制平面和 NAT Gateway。

快速确认 CLI 指向的是你预期的账户的方法：

aws sts get-caller-identity

训练代码：一个函数和一个 trainer

code/train.py 就是全部工作，其中没有任何针对容器、集群或网络的内容。Ray 会为每个 GPU 运行一次 train_func。它向 Ray 询问自己的 rank，加载模型，将其交给 prepare_model 进行 FSDP，并运行一个普通的 PyTorch 循环：

code/train.py — 每个 GPU 的函数（节选）

def train_func(config):
    ctx = ray.train.get_context()
    world_size, rank = ctx.get_world_size(), ctx.get_world_rank()
    device = ray.train.torch.get_device()

    model = AutoModelForCausalLM.from_pretrained(config["model_id"], torch_dtype=torch.bfloat16)

    # Any submodule above 10M parameters becomes its own FSDP unit. Wrapping by
    # size, not by a model-specific layer class, keeps the script model-agnostic.
    auto_wrap_policy = partial(size_based_auto_wrap_policy, min_num_params=int(1e7))
    model = ray.train.torch.prepare_model(
        model,
        parallel_strategy="fsdp",
        parallel_strategy_kwargs={"auto_wrap_policy": auto_wrap_policy},
    )
    optimizer = torch.optim.AdamW(model.parameters(), lr=config["learning_rate"])

    for step in range(config["steps"]):
        input_ids = torch.randint(0, model.config.vocab_size,
            (config["batch_size"], config["seq_len"]), device=device)
        optimizer.zero_grad()
        loss = model(input_ids=input_ids, labels=input_ids).loss
        loss.backward()  # parameters gathered again, gradients reduce-scattered
        optimizer.step()  # local: each rank updates only its own slice
        ray.train.report({"step": step, "loss": loss.item()})

driver 加入正在运行的集群，统计其 GPU 数量并启动 trainer。由于 worker 数量来自集群，同一个脚本可以在单节点或多节点上不加修改地运行：

code/train.py — driver（节选）

ray.init(address="auto")  # join the RayCluster
num_workers = int(ray.cluster_resources().get("GPU", 0))  # 8 here

trainer = TorchTrainer(
    train_func,
    train_loop_config={...},
    scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=True),
    run_config=RunConfig(name="ray-train-fsdp-llm"),
)
result = trainer.fit()

这些批次是随机 token ID，因此运行只依赖于集群本身。样本的评判标准是真实的 FSDP 集合通信是否跨越了两个节点，以及作业是否完成。对于真实的微调，请将随机张量替换为 Ray Data 或 DataLoader 管道。函数中其他内容不变。

清单：唯一配置 DLC 的地方

manifest/raycluster.yaml 是镜像与集群交汇的地方，它很短，因为镜像需要的配置很少。head 唯一的特殊设置是 NVIDIA_VISIBLE_DEVICES=void，它防止容器运行时的 NVIDIA hook 在没有 GPU 的节点上失败。worker group 包含其余所有内容：

manifest/raycluster.yaml — worker group（节选）

workerGroupSpecs:
  - groupName: gpu-workers
    replicas: 2
    rayStartParams:
      num-gpus: "${GPUS_PER_NODE}"  # keeps Ray's GPU count equal to the pod's
    template:
      spec:
        nodeSelector: { role: gpu-worker }
        containers:
          - name: ray-worker
            image: "${DLC_IMAGE}"  # the DLC as published; same tag as the head
            env:
              - { name: NCCL_DEBUG, value: "INFO" }  # print the chosen transport
              - { name: NCCL_SOCKET_IFNAME, value: "eth0" }  # pod network on EKS
              - { name: FI_PROVIDER, value: "efa" }  # libfabric: EFA, nothing else
              - { name: NCCL_NET_PLUGIN, value: "ofi" }  # load the AWS OFI plugin
            resources:
              limits: { nvidia.com/gpu: "${GPUS_PER_NODE}", vpc.amazonaws.com/efa: "1",
                cpu: "40", memory: 160Gi }
              requests: { nvidia.com/gpu: "${GPUS_PER_NODE}", vpc.amazonaws.com/efa: "1",
                cpu: "40", memory: 160Gi }

可以将其理解为三个决策。

使用哪个镜像：${DLC_IMAGE}，不做修改，head 和 worker 都使用它。

使用什么硬件：每个 pod 四个 GPU 和节点的单个 EFA 接口。请求 vpc.amazonaws.com/efa 才能使 EFA 设备插件将接口挂载到容器中。没有它，NCCL 无法看到 EFA，会回退到套接字。

使用哪种传输：四个环境变量，因此不留给自动检测。

两个较小的细节可保持多节点作业稳定。首先，requests 应等于 limits，这会将 worker 置于 Kubernetes 的 Guaranteed 服务质量（QoS）类别中。这很重要，因为单个被限流或被驱逐的 rank 会拖住其他七个。其次，num-gpus 与 GPU 请求派生自同一个变量，因此 Ray 对容量的视图始终与 Kubernetes 一致。

启动：一条命令启动基础设施，然后运行作业

配置集中在一个地方，scripts/env.sh：集群名称、Region、Kubernetes 版本、节点类型和数量、镜像、KubeRay 和两个设备插件的固定版本，以及训练参数。导出其中任何一个即可为一次运行更改它。所有五个脚本都是幂等的，可以安全地重新运行。deploy_cluster.sh、deploy_node_group.sh 和 deploy_ray_train_job.sh 在创建或触碰集群资源之前都会提示。install_gpu_plugins.sh 和 install_kuberay.sh 是协调器，直接运行。

cd training/ray-train/scripts
source ./env.sh

./deploy_all.sh – 找到容量实际所在的位置，然后在那里构建基础设施。

g6.12xlarge 的容量每天在 Region 和可用区之间变化，因此这是起点，而不是在 env.sh 中预先猜测的 Region。它首先运行 find_gpu_capacity.sh，按候选 Region 检查：每个可用区是否提供 g6.12xlarge，账户的按需 G/VT 系列 vCPU 配额是否有空间容纳请求的节点数，以及 EC2 Spot 放置分数作为容量信号。使用 --probe（deploy_all.sh 默认传递），它通过在每个候选区创建并立即取消真实的容量预留来获得确定答案。然后它显示计划并询问一次：

Proceed with the whole chain without further prompts? (y/N):

回答 y 会按顺序在推荐区域中运行 deploy_cluster.sh、deploy_node_group.sh、install_gpu_plugins.sh 和 install_kuberay.sh，每个脚本都会跳过自己的重复确认，因为这一个已经询问过了（-y 也会跳过此提示，而 --plan 只显示计划而不创建任何内容）。如果 GPU 节点组无法在最佳可用区启动，脚本会将其移除并尝试探测确认的下一个可用区。GPU_AZ 还会与集群实际的私有子网进行交叉核对，而不仅仅是 g6.12xlarge 恰好被提供的区域，因此 GPU 节点组不会最终被固定到集群没有子网的可用区。所选区域和可用区会保存到 scripts/.placement.env，其他所有脚本都会读取它，因此后续命令无需重新导出任何内容即可定位到同一位置。在新区域中，这四个步骤总共需要 20 到 30 分钟。现有的集群或节点组会被检测并复用。

deploy_cluster.sh 创建 EKS 控制平面、VPC、OpenID Connect (OIDC) 提供程序、核心附加组件和 m7i.xlarge 系统节点组，并在实际提供 g6.12xlarge 的可用区中选择子网。deploy_node_group.sh 添加两个启用 EFA 的 g6.12xlarge 实例，标记为 role=gpu-worker 和 nvidia.com/gpu.present=true。第二个标签是团队经常忽略的一个。没有它，NVIDIA 设备插件 chart 的 DaemonSet 会在零个节点上运行，且没有任何错误指向原因。install_gpu_plugins.sh 确保两个设备插件都存在。最近的 eksctl 将它们捆绑到一个启用 EFA 的节点组中，因此在新集群上它会找到它们并跳过。否则，它会以固定的 chart 版本安装它们。无论哪种方式，它都会验证 nvidia.com/gpu 和 vpc.amazonaws.com/efa 在每个 GPU 节点上可分配，而不是信任 Helm 发布成功。install_kuberay.sh 将 KubeRay operator 安装到系统节点上，使其远离 GPU 实例。

deploy_all.sh 默认不会做的是运行训练作业。这仍然是一个单独的步骤，你可以随意重新运行，因为它是一次真实的作业运行而非基础设施，因此失败的作业永远不会让你重新经历集群设置。传入 --train 将其串联起来：

./deploy_ray_train_job.sh —— RayCluster、作业和检查。

使用 envsubst 渲染清单并应用它。然后它在等待任何 pod 之前重新读取每个 GPU 节点的可分配资源：如果设备插件没有启动，你会在这里得到明确的错误，而不是一个一直处于 Pending 状态直到超时的 pod。一旦 head 就绪且两个 worker 都就绪，它就会将 code/ 复制到 head 上，并通过 kubectl exec 从内部提交：

ray job submit --address http://localhost:8265 --working-dir /tmp/ray-train-code \
python3 train.py --model_id Qwen/Qwen2.5-1.5B

提交调用会流式传输所有八个 rank 的日志，并在作业完成时返回。然后脚本会检查 NCCL 使用了哪种传输方式，如果作业失败或 EFA 不在路径中，则以错误退出，这也使其可用作持续集成 (CI) 门禁。./deploy_ray_train_job.sh status 打印 RayCluster 状态、每个 pod 所在的主机以及 GPU 容量。正确配置的集群如下所示：

$ kubectl get nodes -l role=gpu-worker \
-o custom-columns='NAME:.metadata.name,GPU:.status.allocatable.nvidia\.com/gpu,EFA:.status.allocatable.vpc\.amazonaws\.com/efa'
NAME               GPU   EFA
ip-192-168-70-51   4     1
ip-192-168-92-26   4     1

图 3：此集群上的一个 FSDP 步骤。权重在前向和反向传播中收集，梯度进行 reduce-scatter。实例之间的流量走 DLC 提供的路径，NCCL → AWS OFI NCCL 插件 → libfabric → EFA，当该路径缺失时回退到 TCP 套接字。

将其适配到你的工作负载

由于镜像按发布原样使用，且作业从集群获知其规模，大多数更改只涉及一个文件：

除非另有说明，表中的变量位于 env.sh 中。切换实例系列时，GPUS_PER_NODE 必须与硬件匹配，并且该系列需要支持 EFA，还要出现在 EFA 设备插件允许的实例类型列表中。此处固定的插件版本已知包含 g6.12xlarge。迁移到 p5 或其他支持 GPUDirect 的系列无需更改镜像或训练代码。

拆除

GPU 实例是开销最大的部分，连同控制平面和 NAT Gateway 每小时约 9.50 美元。请按与创建时相反的顺序移除资源：

./delete_ray_train_job.sh        # RayCluster
./install_kuberay.sh cleanup     # KubeRay operator
./install_gpu_plugins.sh cleanup # NVIDIA and EFA device plugins
./delete_node_group.sh           # GPU instances
./delete_cluster.sh              # EKS cluster, VPC and system node

下一步

两个实例和八个 GPU 在 EKS 上运行了一次分片微调，而容器无需自行构建。Ray Train DLC 提供了从 CUDA 到 EFA 的经过测试的软件栈。KubeRay 使用同一个标签将其启动为 head 和 worker。Ray Train 放置了 worker 并组建了进程组。从硬件到传输再到 Python 包，所有可变的部分都从镜像外部设置。镜像自带的 EFA 路径通过四个环境变量生效，在代码不变的情况下将吞吐量提升至 TCP 的两倍。

自然的下一步：

• 将随机批次替换为真实数据集，并将检查点保存到 Amazon Simple Storage Service（Amazon S3）或 Amazon FSx for Lustre。
• 将 DLC_IMAGE 固定到带版本号的标签，以实现可复现的运行。
• 开启 FailureConfig，使长时间运行在实例丢失后仍能继续。
• 迁移到更大的模型，使用支持 GPUDirect RDMA 的系列，例如 p5。镜像和代码保持不变。

关于作者

Sunita Koppar

Sunita 是 AWS 生成式 AI 和机器学习团队的高级专家解决方案架构师，她与不同行业的客户合作，设计解决方案、构建概念验证并推动可衡量的业务成果。在工作之外，她对学习和教授梵语充满热情，积极与学生社区互动，帮助他们提升技能和成长。

Kevin Wang

Kevin 是 AWS 深度学习容器/AMI 团队的软件开发工程师，负责构建和维护客户在 Amazon Elastic Compute Cloud（Amazon EC2）、Amazon SageMaker AI 和 Amazon Elastic Kubernetes Service（Amazon EKS）上运行的深度学习镜像。他的团队在 GitHub 上以开源方式管理这些镜像及其构建系统。他拥有北京大学学士学位和南加州大学硕士学位，均为计算机科学专业。

Sukhmeet Singh

Sukhmeet 是 AWS 生成式 AI 和机器学习团队的高级专家解决方案架构师。他与独立软件供应商（ISV）客户合作，设计 AI 解决方案、构建概念验证，并将其投入生产，取得可衡量的业务成果。他的兴趣涵盖生成式 AI、计算机视觉、IoT、ML 基础设施和自然语言处理。他尤其喜欢将这些领域联系起来，以大规模解决现实世界的问题。