在 Amazon SageMaker HyperPod 上为 Amazon Nova 部署多轮 RL 基础设施
Deploying Multi-Turn RL Infrastructure for Amazon Nova on Amazon SageMaker HyperPod
Amazon 推出基于 Amazon SageMaker HyperPod 和 Amazon Nova Forge 的多轮强化学习(RL)基础设施,用于训练执行多步骤工作流的 agent。该方案通过事件驱动管道自动配置资源:数据上传至 S3 后触发 Step Functions 编排,SageMaker HyperPod 集群运行 GRPO 权重更新,ECS on Fargate 执行奖励环境(如 Wordle),Nova Forge SDK 负责模型与环境间的消息路由和对话状态跟踪。基础设施采用两阶段部署,长期资源(VPC、EKS、ECS、S3、IAM)通过 AWS CDK 一次性部署,每次训练运行启动临时计算资源。训练参数如 `max_steps`(默认10)、`generation_replicas`(默认4)、`global_batch_size`(默认64)可在 `cdk.json` 中配置或通过 `cdk deploy -c` 覆盖。成本约为每小时 $786-$1,180(10-12 个 ml.p5.48xlarge 实例)。
当你构建执行多步骤工作流的企业级 agent 时,会面临一个根本性的训练挑战。这些 agent 需要查询数据库、调用 API、交叉引用结果,并从流程中途的失败中恢复。任何单一动作的质量都取决于后续几步的结果。标准的基于人类反馈的强化学习(RLHF)孤立地优化单一响应。这种方法对于多步骤工作流来说是不够的,因为 agent 在继续之前验证数据可以防止下游错误的级联。多轮强化学习(RL)通过优化整个交互序列来解决这一差距。你的 agent 通过试错学习工具编排、错误恢复和多步推理。监督微调(SFT)、检索增强生成(RAG)和继续预训练是补充技术,但它们通常本身并不能教授这些序列决策能力。Amazon SageMaker AI 也提供多轮 RL 作为完全托管、无服务器的能力,将此技术引入 SageMaker 训练作业,无需管理基础设施。当你需要对训练堆栈进行完全控制时:你自己的 agent 环境、自定义编排或特定实例配置。对于这些情况,Amazon Nova 在 Amazon SageMaker HyperPod 上的多轮 RL 基础设施为你提供了计算、编排和奖励路由层,以在这些复杂工作流上训练 agent。Amazon Nova 提供前沿智能和行业领先的性价比,而 Amazon Nova Forge 通过多轮 RL 训练能力扩展了这一点。在这篇文章中,你将使用 Amazon Nova Forge 在 Amazon SageMaker HyperPod 上部署一个两阶段的多轮 RL 基础设施。最终,你将拥有一个事件驱动的管道,当你将数据上传到 Amazon Simple Storage Service (Amazon S3) 时,它会自动开始训练。训练作业教会模型玩 Wordle,这是你自己 RL 任务的占位符。
解决方案概览
该解决方案是一个事件驱动的管道:你将数据集上传到 Amazon S3,基础设施会自动配置计算资源、路由奖励并运行多轮 RL 训练。三个层负责执行工作。一个 SageMaker HyperPod 集群生成响应并应用 GRPO(Group Relative Policy Optimization)权重更新。AWS Fargate 上的 ECS 运行你的奖励环境。Nova Forge SDK 在模型和该环境之间路由消息,同时跨轮次跟踪对话状态。AWS Step Functions 编排运行,由 Amazon EventBridge 在数据到达 S3 时触发。该架构分为两个阶段:一次性的 AWS Cloud Development Kit (AWS CDK) 部署提供长期存在的基础(VPC、EKS/HyperPod、ECS、S3、IAM 和管道),而每次训练运行则启动自己的临时资源。这可以防止 GPU 计算资源在运行之间闲置,并让你无需重新部署即可迭代。下图显示了这些组件在训练运行期间如何交互。
Amazon SageMaker HyperPod 基础设施架构上的多轮 RL
该管道在三个计算平面上编排一个多轮对话循环:
- Amazon SageMaker HyperPod (EKS):训练主节点、工作节点和 vLLM 生成副本在 P5 实例上运行。模型生成响应。训练节点使用奖励信号执行 GRPO 权重更新。
- ECS on Fargate:奖励工作节点运行你的环境(例如,Wordle 或自定义的 Bring Your Own Orchestrator (BYOO) 环境)。它们通过 SQS 接收模型响应,根据你的评分标准对其进行评分,并返回奖励信号。
- Amazon Nova Forge:SDK 的代理层在模型和奖励环境之间路由消息,跨轮次跟踪对话状态。
前提条件
在部署之前,请确保你具备以下条件:
- Amazon Nova Forge 订阅:你需要此订阅才能访问 Nova Forge SDK 和模型训练 API。
- SageMaker HyperPod 实例配额:对于
generation_replicas: 4,至少需要 10 ×ml.p5.48xlarge。对于生产工作负载,请求 12-14 个。也支持ml.p5en.48xlarge。 - 已引导的 CDK 环境:在首次部署之前运行
cdk bootstrap。 - Python 3.12+:CDK 应用和 AWS Lambda 运行时需要此版本。
- AWS CDK v2:使用
npm install -g aws-cdk安装。用于合成和部署 AWS CloudFormation 堆栈。 - AWS Command Line Interface (AWS CLI) v2:配置了具有创建 VPC、EKS 集群、SageMaker HyperPod 集群、IAM 角色和 Step Functions 权限的凭证。
- Docker:你使用 Docker 来构建打包 Nova Forge SDK 的 Lambda 容器镜像。
重要提示: 此基础设施在运行时每小时大约花费 $786-$1,180(10-12 个 ml.p5.48xlarge 实例)。请查看成本分解部分,并计划在不进行主动训练时销毁堆栈。
部署基础设施
克隆并安装
首先克隆示例仓库并为 AWS CDK 应用安装 Python 依赖项。在你的本地机器或开发环境中运行以下命令:
git clone https://github.com/aws-samples/sample-nova-multi-turn-rl-infra.git
cd nova-multi-turn-rl-infra
pip install -r requirements.txt
配置
在 cdk.json 中的 context 键下设置所有参数。首次部署前需要两个参数:
| 参数 | 描述 |
|---|---|
project_tag |
所有资源名称的唯一前缀(例如,my-nova-rl)。将成为你的 EKS 集群名称、HyperPod 集群名称和 S3 存储桶标签的一部分。 |
sdk_resource_prefix |
SDK 创建的资源的前缀(例如,nrl-myproject)。Nova Forge SDK 在命名其 CloudFormation 堆栈、Lambda 函数和 SQS 队列时使用此前缀。 |
关键基础设施参数:
| 参数 | 默认值 | 描述 |
|---|---|---|
instance_type |
ml.p5.48xlarge |
HyperPod 实例类型 |
instance_count |
10 | 受限实例组中的实例数量 |
nova_model |
NOVA_LITE_2 |
要训练的模型:NOVA_MICRO、NOVA_LITE、NOVA_LITE_2 或 NOVA_PRO |
vf_env_id |
wordle |
用于验证的内置奖励环境 |
use_custom_env / custom_env_id |
— | 将 use_custom_env 设置为 "true" 并指定 custom-environments/ 下的目录名称用于 BYOO |
reward_cpu / reward_memory |
2048 CPU, 4096 MiB | Fargate 任务大小调整 |
eks_kubernetes_version |
1.32 | EKS 集群版本 |
训练参数。这些参数通过管道事件流向训练作业:
| 参数 | 默认值 | 描述 |
|---|---|---|
training_method |
— | RFT_MULTITURN_FULL 或 RFT_MULTITURN_LORA |
max_steps |
10 | 训练步数 |
generation_replicas |
4 | vLLM 生成副本数 |
global_batch_size |
64 | 每步训练的样本数 |
你可以在部署时覆盖任何参数:
cdk deploy -c instance_count=1 -c max_steps=20
部署
cdk deploy --require-approval never
两阶段部署模型
该基础设施遵循两阶段部署模型,将长期存在的基础资源与临时的每次运行资源分开。这种分离通过仅在需要时创建昂贵的计算资源来降低成本,通过避免为每次训练运行完全重新部署来加快迭代速度,并通过为每一层提供独立的生命周期来简化管理。
阶段 1:AWS CDK 部署(一次性): 当你运行 cdk deploy 时,你将配置基础基础设施:
- 一个带有私有子网和 VPC 端点的 Amazon Virtual Private Cloud (Amazon VPC)。
- 一个带有 HyperPod 受限实例组 (RIG) 的 Amazon Elastic Kubernetes Service (Amazon EKS) 集群。
- 一个用于奖励工作节点的 AWS Fargate 上的 Amazon Elastic Container Service (Amazon ECS) 集群。
- 一个用于训练数据和检查点的 Amazon Simple Storage Service (Amazon S3) 存储桶。
- AWS Identity and Access Management (IAM) 角色。
- 一个 AWS Step Functions 管道。
- Amazon EventBridge 规则。
部署大约需要 30-40 分钟。大部分时间花在 EKS 集群创建(约 15 分钟)上,其次是 SageMaker HyperPod Helm chart 安装(通过 AWS CodeBuild,约 5 分钟)、SageMaker HyperPod 集群配置(约 15-25 分钟,取决于 P5 容量)和 Lambda 容器镜像构建(约 5 分钟)。
阶段 2:运行时(每次训练运行): 当你将一个 .jsonl 文件上传到 S3 存储桶的 training-data/ 前缀时,EventBridge 会触发一个 Step Functions 管道,该管道为该训练运行创建运行时资源。Nova Forge SDK 部署其自己的 AWS CloudFormation 堆栈,其中包含 AWS Lambda 函数(对话代理)、Amazon Simple Queue Service (Amazon SQS) FIFO 队列(模型和环境之间的消息路由)、一个 Amazon DynamoDB 表(对话状态跟踪)和 ECS Fargate 任务(奖励工作节点)。SDK 管理这些资源的生命周期。
触发训练运行
部署基础设施后,你只需一次 S3 上传即可开始训练运行。EventBridge 监视存储桶的 training-data/ 前缀,并自动启动 Step Functions 管道。一次成功的训练运行会在 Amazon CloudWatch 指标中显示奖励分数随连续步骤的增加而增加。对于 Wordle 环境,预计模型会在 50-100 步内收敛,平均奖励从接近零提高到 0.6-0.8,因为模型学会了根据反馈缩小猜测范围。
准备训练数据
.jsonl 文件使用基于元数据的格式:每行包含一个提示和奖励环境用于评分的真实答案:
{"id": "wordle_train_001", "metadata": {"prompt": "Guess the 5-letter word", "answer": "crane"}}
{"id": "wordle_train_002", "metadata": {"prompt": "Guess the 5-letter word", "answer": "slate"}}
{"id": "wordle_train_003", "metadata": {"prompt": "Guess the 5-letter word", "answer": "plumb"}}
id 字段在所有记录中必须是唯一的。模型在每个对话开始时看到 metadata.prompt,奖励环境使用 metadata.answer 来跨轮次对模型的响应进行评分。每个训练示例都成为一个多轮对话:模型生成一个猜测,从奖励环境接收反馈,并迭代直到解决谜题或用完所有轮次。
上传并自动触发
# 从 CDK 输出获取存储桶名称
BUCKET=$(aws cloudformation describe-stacks --stack-name NovaMultiTurnRlStack \
--query "Stacks[0].Outputs[?OutputKey=='TrainingBucketName'].OutputValue" --output text)
# 上传:管道自动启动
aws s3 cp training-data.jsonl s3://$BUCKET/training-data/training-data.jsonl
当一个文件到达 training-data/ 前缀时,EventBridge 检测到 PutObject 事件并调用 S3TriggerFn Lambda 函数。此函数提取文件路径,将其与 cdk.json 中的训练参数合并,并将生成的事件数据传递给 Step Functions,后者通过五个阶段执行管道。你可以在 Step Functions 控制台中查看进度,其中每个步骤都显示其输入、输出、持续时间和重试次数。
手动触发以进行临时运行
要迭代参数而无需重新上传数据,请使用设置脚本绕过 EventBridge 并直接调用管道:
# 覆盖训练参数以进行快速测试运行
./scripts/setup.sh \
--data-path s3://$BUCKET/training-data/training-data.jsonl \
--max-steps 5 \
--global-batch-size 32 \
--training-method RFT_MULTITURN_LORA
这在开发期间很有帮助:指向现有数据集并试验不同的步数、批量大小或训练方法(全微调与 LoRA),而无需修改 cdk.json 或重新部署。
监控和调试
你可以跨 SageMaker HyperPod、ECS、Lambda 和 SQS 监控和排查此多阶段管道的每一层。基础设施在每一层都提供可观测性,以便你快速隔离问题。
管道监控
Step Functions 控制台是主要仪表板。每个执行都显示逐步进度,包括时间、输入和输出数据以及重试历史。每个步骤都写入自己的 Amazon CloudWatch 日志组。在正常操作期间,每个 Step Functions 步骤都在其预期时间窗口内完成:基础设施设置 2-3 分钟,奖励工作节点部署 1-2 分钟,数据验证不到 1 分钟,训练提交 3-5 分钟。SQS 队列应显示消息稳定流动,没有持续积压。
失败告警
一个使用 AWS Key Management Service (AWS KMS) 加密的 Amazon Simple Notification Service (Amazon SNS) 主题,在 Step Functions 执行失败时发布通知。订阅以通过电子邮件、Slack 或 PagerDuty 接收告警:
# 从 CDK 输出获取主题 ARN
TOPIC_ARN=$(aws cloudformation describe-stacks --stack-name NovaMultiTurnRlStack \
--query "Stacks[0].Outputs[?OutputKey=='AlertTopicArn'].OutputValue" --output text)
# 订阅并通过电子邮件确认
aws sns subscribe \
--topic-arn $TOPIC_ARN \
--protocol email \
--notification-endpoint your-team@example.com
触发失败
如果 S3 上传未启动管道,请检查死信队列。失败的 EventBridge 到 Lambda 调用会存储在此处,并包含原始 S3 事件数据和错误详细信息:
DLQ_URL=$(aws cloudformation describe-stacks --stack-name NovaMultiTurnRlStack \
--query "Stacks[0].Outputs[?OutputKey=='TriggerDlqUrl'].OutputValue" --output text)
aws sqs receive-message --queue-url $DLQ_URL --max-number-of-messages 5
常见原因包括文件名格式错误(触发器期望 .jsonl 扩展名)、IAM 权限漂移和 Lambda 并发限制。
训练停滞调试
如果管道启动但训练停滞,根本原因通常在于模型和奖励环境之间的消息路由层。使用 SDK 内置诊断检查 SQS 队列健康状况:
from rft_infra import check_all_queues
# 返回所有 FIFO 队列的消息计数(正在处理、可用、延迟)
check_all_queues()
请求队列中不断增长的积压加上空的响应队列表明奖励工作节点未在处理。响应队列中不断增长的积压表明模型未在消耗奖励。任一模式都告诉你应该调查哪个组件。
HyperPod 健康检查
验证所有训练和生成 pod 是否在 SageMaker HyperPod 集群上运行:
kubectl get pods -n kubeflow
所有 pod 应显示 Running 或 Completed 状态。处于 Pending 状态的 pod 表示受限实例组中的实例容量不足。处于 CrashLoopBackOff(一种 Kubernetes 状态,表示容器反复崩溃并重启)状态的 pod 表示容器级错误。使用 kubectl logs -n kubeflow 检查 pod 日志以获取详细信息。
清理和管理成本
清理脚本按以下顺序停止活动资源:Step Functions 执行、ECS 任务、SDK CloudFormation 堆栈和 CDK 堆栈。为避免产生费用,请运行以下命令销毁整个堆栈:
./cleanup.sh
要销毁基础设施同时保留 S3 中的训练数据(存储费用仍将适用):
./cleanup.sh --retain-data
重要提示: 在不进行主动训练时销毁堆栈。SageMaker HyperPod 实例和 NAT Gateway 的空闲成本会迅速累积。
成本分解
| 资源 | 每小时成本 | 备注 |
|---|---|---|
| SageMaker HyperPod 8 × ml.p5.48xlarge | ~$786/hr | 最低配置(10 个实例,8 个用于计算) |
| SageMaker HyperPod 12 × ml.p5.48xlarge | ~$1,180/hr | 生产配置(12 个实例,有预留空间) |
| EKS 控制平面 | $0.10/hr | |
| NAT Gateway | ~$0.045/hr | |
| ECS Fargate 奖励工作节点 | ~$0.15-0.30/hr + 数据传输 | |
| S3、Lambda、CodeBuild、SQS、DynamoDB | 可忽略不计 |
有关当前 Amazon Nova Forge 定价,请参阅 Amazon Nova 定价页面。
优化成本
你可以通过两种策略降低空闲成本:
- 缩放到零:配置一个具有自动扩缩策略的单实例集群,该策略在非活动期间缩放到零,从而消除空闲计算成本。
- 实例类型调度:在设置和低峰阶段使用较低成本的实例类型,然后在训练工作负载开始时过渡到更高性能的实例。
结论
你现在拥有一个生产就绪、事件驱动的基础设施,用于在 Amazon SageMaker HyperPod 上使用 Amazon Nova Forge 进行多轮强化学习。将数据集上传到 Amazon S3,管道会自动处理配置、训练编排和奖励路由。要将此方案适应你的用例,请将 Wordle 环境替换为你自己的 API 调用 agent 或企业工作流。
后续步骤:
- 克隆示例仓库:
aws-samples/nova-multi-turn-rl-infra。 - 阅读 Amazon SageMaker HyperPod 文档。
- 阅读关于 Amazon Nova 功能和定价的信息。
- 使用 AWS Step Functions 开发者指南自定义你的管道。
关于作者
Maria Masood Maria 专注于 agentic AI、强化微调和多轮 agent 训练。她在机器学习方面拥有专业知识,涵盖大型语言模型定制、奖励建模以及为 AI agent 构建端到端训练管道。她内心是一名可持续发展爱好者,喜欢园艺和制作拿铁咖啡。
Nick Biso Nick 是 AWS Professional Services 的高级机器学习工程师。他使用数据科学和工程解决复杂的组织和技术挑战。此外,他还在 AWS 云上构建和部署 AI/ML 模型。他的热情还延伸到对旅行和多元文化体验的爱好。
Christian Kamwangala Christian 是 AWS 的 AI/ML 和生成式 AI 专家解决方案架构师,他与企业客户合作,架构、优化和部署生产级 AI 解决方案。他的专长在于推理优化——为大规模部署平衡性能、成本和延迟。工作之余,他喜欢探索大自然并与家人和朋友共度时光。
Manoj Gupta Manoj 是 AWS 的高级解决方案架构师,常驻旧金山。他在 AWS 拥有超过 4 年的经验,与客户密切合作,构建优化的 AI/ML 驱动解决方案和云基础设施。他的主要关注领域是数据、AI/ML 和安全,帮助组织现代化其技术栈。工作之余,他喜欢户外活动和与家人一起旅行。