用 Ray 与 EFA 复现多机模型微调

把 PyTorch 训练从一台机器搬到多台，难点经常出在运行环境：CUDA、PyTorch、NCCL、EFA 和 Ray 要配套，网络没配对，任务甚至可能正常结束，只是跑得很慢。AWS 这篇文章给出的办法，是用官方 Ray Train DLC 承担镜像维护，把团队自己的工作收敛到训练函数和集群配置。

作者用两台 g6.12xlarge、八张 NVIDIA L4 跑通了 Qwen2.5-1.5B 的 FSDP 训练。这个案例值得看的是如何验证跨机通信确实走了 EFA；至于真实微调效果和吞吐收益，原文提供的证据还不完整。

不自己构建镜像，究竟省掉了什么

分布式训练镜像里，除了 CUDA 和对应的 PyTorch，还需要 NCCL、EFA 用户态库，以及连接两者的 AWS OFI NCCL plugin。Ray 的版本也要与集群配合。这些依赖分别能装上，不代表组合起来能稳定工作。

最隐蔽的问题在通信层。原文指出，依赖不匹配时，NCCL 可能退回 TCP sockets，训练仍能完成。只检查退出码，很容易把「跑通」误认为「配置正确」。

作者使用的镜像是 public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1。AWS 声称会把整套依赖一起测试，并以多节点 EFA、多节点 KubeRay 训练测试作为发布门槛。镜像从 Amazon ECR Public 拉取，无需仓库凭证；入口脚本还会检查宿主机 NVIDIA 驱动，必要时启用内置 CUDA forward-compatibility 库。

团队仍要配置硬件和网络。DLC 负责软件栈，KubeRay 根据 RayCluster 创建和维护 head、worker pod，Ray Train 则负责启动每张 GPU 上的训练进程，分配 rank、world size，并建立 torch.distributed 进程组。

三层各管一段，训练代码才不用承担整套启动逻辑。

差异配置也移到镜像外：额外 Python 包通过 runtime_env 按任务安装，通信参数通过环境变量传入，head 和 worker 使用同一个镜像。

版本固定需要分清粒度。按原文说明，vX.Y 标签会跟随安全重建移动，vX.Y.Z 标签保留初始发布构建；要求镜像字节完全一致时，应固定 digest。固定标签能约束版本选择，严格复现还需要固定镜像摘要。

八张 GPU 怎么分工，为什么必须检查 EFA

案例将两个 GPU 节点放在同一可用区，因为 EFA 流量不能跨可用区。每台 g6.12xlarge 有四张 L4、48 vCPU、192 GiB 内存、40 Gbps 网络和一个 EFA 接口，一个 worker pod 独占四张 GPU 和该接口。

另有一台 m7i.xlarge 运行系统组件、KubeRay operator 和 Ray head。head 不声明 GPU，因此不会接到训练进程。所有节点位于私有子网，通过 NAT Gateway 拉镜像和模型；任务经 kubectl exec 从 head 内提交，8265 端口无需暴露到集群外，6379 用于 Ray 控制存储。

Ray Train 在每个 worker pod 内启动四个训练进程，共八个 rank。FSDP，也就是完全分片数据并行，把参数、梯度和优化器状态拆到各张卡上。原文估算，Qwen2.5-1.5B 使用 bfloat16 时，这些张量合计约 12 GB，八路分片后每卡约 1.5 GB。这个数字不等于训练峰值显存，原文没有给出激活、临时聚合等开销的完整测量。

节省常驻显存的代价是通信：计算前聚合权重，反向传播时再做梯度归约与分片。跨机链路因此直接影响训练效率。

下面这段配置同时声明设备需求和通信路径，是整个案例最值得保留的部分：

workerGroupSpecs:
  - groupName: gpu-workers
    replicas: 2
    rayStartParams:
      num-gpus: "${GPUS_PER_NODE}"  # 让 Ray 的 GPU 数量与 pod 一致
    template:
      spec:
        nodeSelector: { role: gpu-worker }
        containers:
          - name: ray-worker
            image: "${DLC_IMAGE}"  # 使用原版 DLC，与 head 标签相同
            env:
              - { name: NCCL_DEBUG, value: "INFO" }  # 输出所选通信路径
              - { name: NCCL_SOCKET_IFNAME, value: "eth0" }  # EKS pod 网络
              - { name: FI_PROVIDER, value: "efa" }  # libfabric 使用 EFA
              - { name: NCCL_NET_PLUGIN, value: "ofi" }  # 加载 AWS OFI 插件
            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 }

关键是 vpc.amazonaws.com/efa: "1"：它让设备插件把 EFA 接口挂进容器。环境变量写对了，设备没申请到，NCCL 仍可能退回 sockets。requests 与 limits 相同则用于获得 Guaranteed QoS，但这并不保证节点永远不会故障。

head 还有一个特殊配置：NVIDIA_VISIBLE_DEVICES=void，避免没有 GPU 的 CPU 节点触发 NVIDIA 容器钩子错误。

训练函数经过 prepare_model(..., parallel_strategy="fsdp") 包装后，仍然使用普通 PyTorch 循环；自动包装策略为 size_based_auto_wrap_policy，阈值 min_num_params=int(1e7)，超过 1000 万参数的子模块会成为独立 FSDP 单元。驱动部分负责按集群 GPU 数量启动训练：

ray.init(address="auto")  # 加入 RayCluster
num_workers = int(ray.cluster_resources().get("GPU", 0))  # 此处为 8

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

要留意 ray.cluster_resources() 这一行：示例按集群全部 GPU 计算 worker 数。它适合这个专用集群，迁入共享集群时，需要重新考虑资源分配；代码中的 {...} 也表明这只是节选。

部署成功，还要证明训练走对了路

部署入口是 scripts/env.sh，集中管理 Region、Kubernetes 版本、节点、镜像、插件版本和训练参数。准备工具包括 AWS CLI v2、eksctl、kubectl、Helm，以及 GNU gettext 的 envsubst；先用 aws sts get-caller-identity 核对目标账号。

两台 GPU 实例需要 96 vCPU 的 G-family On-Demand 配额。原文估算 GPU 实例、EKS 控制平面和 NAT Gateway 合计约每小时 9.50 美元，这只是案例预算，不能当成所有 Region 的统一报价。

执行 cd training/ray-train/scripts、source ./env.sh 后，./deploy_all.sh 会先寻找容量。find_gpu_capacity.sh 检查实例供应、配额和 Spot placement score；默认传入的 --probe 还会创建并立即取消 Capacity Reservation，进一步确认容量。已有预留则可设置 CAPACITY_RESERVATION_ID。

脚本支持 --plan 只看计划、-y 跳过确认，位置写入 scripts/.placement.env。新环境搭建约需 20 到 30 分钟，训练默认单独运行，也可用 --train 串起来。

设备检查比 Helm 安装结果更重要。GPU 节点需要 role=gpu-worker 和 nvidia.com/gpu.present=true 标签，后者缺失可能让 NVIDIA device plugin 的 DaemonSet 一个节点都不覆盖。脚本会核验每个 GPU 节点实际可分配的 nvidia.com/gpu 和 vpc.amazonaws.com/efa。

./deploy_ray_train_job.sh 随后提交任务，并检查 NCCL 使用的传输路径。任务失败，或者没有走 EFA，脚本都会报错，因此可以作为 CI 门槛。./deploy_ray_train_job.sh status 则查看集群状态、pod 落点和 GPU 容量。

任务完成和通信配置正确，应当是两个独立的验收条件。

我的判断

这套方案最有价值的地方，是把镜像维护交给提供方，同时把硬件、调度和网络配置留在可审查的文件里。相比团队自己组合训练镜像，它能减少依赖维护；相比直接在 EC2 上启动，还多了 KubeRay 的 pod 管理能力，但也引入了 Kubernetes、设备插件和 operator 的运维成本。

如果团队只是偶尔训练一次，为此搭整套 EKS 未必划算。已有 Kubernetes 和 Ray 基础的团队，更容易获得收益。

实验边界也必须讲清：输入是随机 token IDs，验收标准是跨两节点执行真实 FSDP 通信并完成任务。它验证了基础设施链路，没有证明模型微调质量，也没有覆盖真实数据读取、检查点恢复和长时间运行的表现。

原文末尾称 EFA 相比 TCP 吞吐翻倍，却没给具体吞吐、测试参数、重复次数或完整对照。因此，这只能视为作者报告的结果，不能直接用于团队的容量或成本测算。换数据集之外，生产训练还需要补齐数据分片、检查点和恢复验证。

能拿来做什么

• 如果你正在负责 AI 训练平台，可以先用随机数据建立基础设施验收：设备可分配、worker 数正确、跨节点通信经过 EFA，三项都过再接真实数据。
• 如果团队已有 Ray 任务提交流程，可以用 runtime_env 管理实验新增依赖，同时记录 DLC_IMAGE 的 digest 和依赖版本，避免只保存训练脚本。
• 长任务上线前，应接入 Ray Data 或 DataLoader，向 Amazon S3 或 Amazon FSx for Lustre 保存检查点，并启用 FailureConfig 验证实例丢失后的恢复。
• 试验结束后按顺序运行 ./delete_ray_train_job.sh、./install_kuberay.sh cleanup、./install_gpu_plugins.sh cleanup、./delete_node_group.sh、./delete_cluster.sh，及时释放 GPU 和配套基础设施。