使用 Ray Train DLC 和 EFA 在 EKS 上进行分布式训练
- 作者
- 姓名
- Nino
- 职业
- Senior Tech Editor
为大规模分布式机器学习构建和维护自定义深度学习容器往往会成为研发流程中的效率瓶颈。在扩展诸如 Qwen2.5-1.5B 或 DeepSeek-V3 等大语言模型的微调任务时，手动管理 CUDA 驱动、NCCL 版本以及 AWS 弹性Fabric适配器（EFA）库可能会严重拖慢工程进度。幸运的是，预构建的 Ray Train 深度学习容器原生集成了完整的从 CUDA 到 EFA 的技术栈，开发者无需编写任何自定义的 Dockerfile 即可快速启动企业级集群。
在本篇技术评测中，我们将探讨一个完全可复现的流水线，它在运行于 Amazon EKS 上的两个 g6.12xlarge 实例的八个 GPU 之间，利用 Ray Train、PyTorch 全分片数据并行（FSDP）和 EFA 实现了高效训练。此外，我们还将探讨如何将诸如 n1n.ai 等高性能 LLM API 整合到开发与评估循环中，从而进一步简化工作流。
架构概览：在 EKS 上结合使用 Ray Train、PyTorch FSDP 与 EFA
为了在多节点 Kubernetes 集群中实现线性扩展，我们的基础设施紧密结合了 AWS 网络硬件与 Ray 编排原语。
+-------------------------------------------------------+
| Amazon EKS Cluster |
| |
| +-----------------------+ +---------------------+ |
| | g6.12xlarge (Node 1)| | g6.12xlarge (Node 2)| |
| | - 4x NVIDIA L4 GPUs | | - 4x NVIDIA L4 GPUs | |
| | - EFA Network Devices| | - EFA Network Devices| |
| +-----------^-----------+ +----------^----------+ |
| | | |
| +------------+-------------+ |
| | |
| Ray Train Cluster |
+-------------------------------------------------------+
Amazon EKS 上的每个 g6.12xlarge 实例都提供了四个 NVIDIA L4 GPU 以及超低延迟的 EFA 网络接口。通过使用官方的 Ray Train DLC，底层的 MPI、NCCL 和 libfabric 层已经被预先配置好，能够直接开箱即用 EFA 的内核旁路（kernel bypass）功能。
前提条件与基础设施配置
在部署分布式训练任务之前，请确保您的环境满足以下基本要求：
- 一个处于活动状态的 Amazon EKS 集群（1.28 或更高版本），并配置了 Karpenter 或托管节点组。
- 在
kube-system命名空间中安装了 AWS EFA Kubernetes 设备插件。 - 通过 Helm 部署了用于集群生命周期管理的 Ray Operator。
- 能够访问强大且稳定的 LLM API 聚合平台，例如 n1n.ai，以便在准备阶段持续进行模型评估和合成数据生成。
验证 EFA 接口
在提交 Ray 任务之前，请验证您的 Kubernetes Pod 是否可以检测到 EFA 设备。您可以使用以下命令在测试 Pod 内部运行快速诊断检查：
kubectl run efa-test --rm -it --image=amazon/aws-efa-installer:latest --command -- fi_info -p efa
如果配置正确，输出结果应列出处于活动状态的 efa 提供程序，且不会抛出任何初始化错误。
准备 Ray Train 集群配置
我们使用标准的 Kubernetes 自定义资源（CR）定义 Ray 集群。该配置部署了一个头节点和一个工作节点，并将它们与启用 EFA 的实例配置文件紧密绑定。
apiVersion: ray.io/v1
kind: RayCluster
metadata:
name: ray-train-cluster
spec:
rayVersion: '2.35.0'
headGroupSpec:
rayStartParams: {}
template:
spec:
containers:
- name: ray-head
image: public.ecr.aws/ray/ray:2.35.0-py310-gpu
resources:
limits:
cpu: '16'
memory: 64Gi
nvidia.com/gpu: '0'
workerGroupSpecs:
- groupName: worker-group
replicas: 1
minReplicas: 1
maxReplicas: 2
rayStartParams: {}
template:
spec:
containers:
- name: ray-worker
image: public.ecr.aws/ray/ray:2.35.0-py310-gpu
securityContext:
capabilities:
add: ["IPC_LOCK"]
resources:
limits:
cpu: '48'
memory: 192Gi
nvidia.com/gpu: '4'
aws.amazon.com/efa: '4'
使用 PyTorch FSDP 实现 Qwen2.5-1.5B 微调脚本
接下来，我们使用 Ray Train 的 TorchTrainer 编写训练脚本。PyTorch FSDP 将模型的参数、梯度和优化器状态分片到我们的八个 GPU 上，从而大幅降低了每个设备的显存开销。
import os
import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer
from ray.train.torch import TorchTrainer
from ray.train import ScalingConfig, CheckpointConfig
def train_func(config):
model_id = config.get("model_id", "Qwen/Qwen2.5-1.5B")
# 初始化分词器和模型
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.bfloat16,
device_map="cuda"
)
# 设置 FSDP 包装
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
model = FSDP(model)
# 此处省略训练循环实现...
print("已成功跨分布式节点初始化 FSDP 模型包装器。")
def main():
scaling_config = ScalingConfig(
num_workers=2,
use_gpu=True,
resources_per_worker={"CPU": 8, "GPU": 4}
)
trainer = TorchTrainer(
train_loop_per_worker=train_func,
train_loop_config={"model_id": "Qwen/Qwen2.5-1.5B"},
scaling_config=scaling_config,
)
result = trainer.fit()
print(f"训练完成。检查点路径: {result.checkpoint.path}")
if __name__ == "__main__":
main()
使用统一 API 优化大语言模型工作流
在由 Ray 和 EFA 管理的基础设施上训练自定义权重的同时，针对生产级模型评估中间检查点至关重要。许多工程团队不再维护多个供应商的 SDK，而是选择利用 n1n.ai，通过单一的统一端点轻松地将评估流量路由到 Claude 3.5 Sonnet 和 OpenAI o3 等领先模型。
总结与后续步骤
通过将 Ray Train 深度学习容器、PyTorch FSDP 以及 AWS EFA 结合并在 Amazon EKS 上运行，工程团队可以绕过繁琐的容器镜像编译开销，直接投入到高性能的分布式训练中。无论您是在扩展 Qwen2.5-1.5B 还是在尝试海量的专家混合（MoE）架构，这套技术栈都能提供可预测的硬件利用率和极低的网络延迟。
Get a free API key at n1n.ai