AWS 发布教程：在 EKS 上用 Ray Train DLC 和 EFA 分布式微调 Qwen2.5-1.5B

通常阻碍 PyTorch 训练扩展到集群的是容器，而不是模型代码。每个容器都需要相同的 CUDA 运行时、针对它编译的 PyTorch 构建版本、NCCL（NVIDIA 集合通信库）、Elastic Fabric Adapter（EFA）用户态库、将 NCCL 连接到 EFA 的 AWS OFI NCCL 插件，以及一个集群运维方认可的 Ray 版本。构建这个镜像，并在其中任何一个组件更新时重新构建，本身就是一项工程，而且与你最初要训练的模型毫无关系。Ray Train Deep Learning Container（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 部署八个工作进程（每个 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，通过 PyTorch Fully Sharded Data Parallel（FSDP）在两个 g6.12xlarge 实例上的八块 NVIDIA L4 GPU 上微调 Qwen2.5-1.5B。训练代码只是一个函数，集群只需一个 RayCluster 清单加一个变量文件，全程无需构建任何镜像。因为 DLC 已经自带 EFA 技术栈，一旦 pod 提出请求，两个实例之间的集合通信就会通过 EFA 进行。完整的解决方案在这个 示例代码库中，它使用的 KubeRay-on-EKS 基础与我们的 多节点 Ray Serve 指南相同。

无需构建任何东西即可获得什么

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

Networking 这一行是大多数团队最不愿意自己维护的部分。这里出现不匹配不会报错——NCCL 会退回到 TCP socket，任务照常跑完，只是速度只有原来的一小部分。而在 DLC 中，这些库是一起构建、一起测试的，并且 AWS 对每次发布都用多节点 EFA 和多节点 KubeRay 训练测试做门禁，而这正是本示例所使用的路径。

还有几项特性省去了你原本要在每个集群上做的事：

• 无需配置 registry – 镜像是公开的，节点无需凭据即可拉取，中途不需要构建或推送任何东西。
• 无需匹配驱动 – 启动时，entrypoint 会检查主机的 NVIDIA 驱动，如果驱动比镜像的 CUDA 旧，就自动启用内置的 CUDA 前向兼容库。
• 无需区分 head 与 worker 镜像 – 镜像没有固定角色。启动容器的一方提供 ray start 命令，因此同一个标签可用于所有 pod。
• 无需单独的推理镜像系列 – Ray Train 与 Ray Serve DLC 和 Ray LLM DLCs 共享 ray ECR 仓库，使用相同的 train-ml 标签前缀。

配置它，而不是重建它

训练任务之间所有差异的内容都在镜像外部设置。本示例用到了以下每种机制：

按任务的包。镜像自带 HuggingFace 训练栈，因此本示例的任务无需额外安装任何东西。当实验需要镜像中没有的库时，Ray 会在提交时通过 runtime_env 为该任务安装，镜像保持不变。无需新镜像，无需新标签。

按任务的传输配置。镜像默认使用 FI_PROVIDER=efa，其 NCCL 配置排除了 loopback 和 Docker bridge 接口。在 EKS 上，你将 NCCL_SOCKET_IFNAME 指向 pod 的 eth0。本示例反正会在 manifest 中设置全部四个网络变量，这样行为在单个文件中清晰可见，也不依赖镜像的默认值。由于这些都是普通环境变量，单个任务可以通过其 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 网关访问互联网——每个节点正是通过它从 Amazon ECR Public 拉取 DLC，worker 也通过它从 Hugging Face 下载模型权重。一个 m7i.xlarge 构成系统节点组，托管 kube-system 附加组件、KubeRay operator 和 Ray head。head 在 8265 端口运行作业服务器，在 6379 端口运行 Ray 的控制存储，并且它不声明任何 GPU，因此不会有训练任务落在它上面。作业通过 kubectl exec 从 head pod 内部经 EKS API 提交，因此 8265 端口永远不会暴露到集群之外。

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

Ray Train 在容器内做了什么

DLC 提供了技术栈。Ray Train 提供了原本需要自己编写启动器代码的部分。多节点 PyTorch 通常意味着在每台主机上为每个 GPU 启动一个进程，并为每个进程指定 rank、world size 和一个会合地址。使用 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。
• 配额：在至少有一个可用区提供 G 系列 On-Demand 容量的 AWS 区域中，为两台 g6.12xlarge 实例预留 96 个 vCPU 的 G 系列 On-Demand 容量。如果你已持有这些实例的 On-Demand Capacity Reservation，请改为设置 CAPACITY_RESERVATION_ID。脚本会将节点组启动到该预留中，并跳过容量搜索。
• 预算：GPU 实例运行期间每小时约 $9.50，涵盖两台实例、EKS 控制平面和 NAT 网关。

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

aws sts get-caller-identity

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

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()})

驱动程序加入正在运行的集群，统计 GPU 数量并启动训练器。由于 worker 数量来自集群，同一脚本在单节点或多节点上运行都无需修改：

code/train.py — 驱动程序（有删节）

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 组中：

manifest/raycluster.yaml — worker 组（有删节）

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，只能回退到 socket。

用哪种传输方式：四个环境变量，这样就不会留下任何东西交给自动检测。

两个较小的细节保证了多节点作业的稳定。第一，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 的容量每天都在区域和可用区之间变动，所以这是起点，而不是在 env.sh 中预先猜测某个区域。它首先运行 find_gpu_capacity.sh，后者针对每个候选区域检查：每个可用区是否提供 g6.12xlarge、账户的按需 G/VT 系列vCPU 配额是否够用所请求的节点数量，以及作为容量信号的 EC2 Spot 竞价放置分数。配合 --probe（deploy_all.sh 默认会传入），它会通过在每个候选可用区中创建并立即取消一个真实的容量预留来获得确定的答案。然后它展示计划并询问一次：

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

按顺序在推荐的 Region 中运行 y，即 deploy_cluster.sh、deploy_node_group.sh、install_gpu_plugins.sh 和 install_kuberay.sh，每一步都会跳过自身的重复确认，因为这一步已经询问过（-y 也会跳过此提示，而 --plan 只显示计划而不创建任何东西）。如果 GPU 节点组无法在最佳可用区中启动，脚本会将其删除并尝试探测确认过的下一个可用区。GPU_AZ 还会与集群实际拥有的私有子网进行交叉核对，而不仅仅看 g6.12xlarge 在哪里有售，因此 GPU 节点组不会最终被固定在集群没有子网的可用区上。选定的 Region 和可用区会保存到 scripts/.placement.env，其他所有脚本都会读取它，因此后续命令会指向同一位置，无需重新导出任何内容。在全新的 Region 上，这四个步骤总共需要 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 device plugin 图表的 DaemonSet 会运行在零个节点上，而且没有任何错误指出原因。install_gpu_plugins.sh 确保两个设备插件都已存在。近期的 eksctl 会将它们捆绑进启用 EFA 的节点组，因此在全新集群上它会发现它们并跳过安装；否则它会以固定图表版本安装它们。无论哪种方式，它都会验证 nvidia.com/gpu 和 vpc.amazonaws.com/efa 在每个 GPU 节点上可分配，而不是仅仅信任 Helm release 安装成功。install_kuberay.sh 将 KubeRay 操作器安装到系统节点上，使其远离 GPU 实例。

deploy_all.sh 默认不会运行训练作业。这一步保持独立，你可以随想随重跑，因为它是真正的作业运行而非基础设施，这样作业失败也绝不会让你重新走一遍集群搭建流程。传入 --train 可将其串联起来：

./deploy_ray_train_job.sh ——即 RayCluster、作业以及各项检查。

使用 envsubst 渲染清单并应用。随后，在等待任何 Pod 之前，它会重新读取每个 GPU 节点的可分配资源：如果设备插件没有启动，你会在这里得到一个明确的错误，而不是一个一直处于 Pending 直到超时的 Pod。一旦 head 节点 Ready 且两个 worker 都 Ready，它会将 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 网关一起，每小时约 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）。镜像和代码保持不变。

关于作者