跳到正文
原文
AWS Machine Learning Blog· Nicolas Jourdan·· 3 小时前精选AI 评分61

AWS 推出 SageMaker AI 训练会话 API，可自写训练循环而无需管理 GPU

Write your own training loop — without managing GPUs — with Amazon SageMaker training session APIs

AI 导读

AWS 在 Amazon SageMaker AI 的无服务器模型定制中新增训练会话 API，把 forward_backward()、optim_step()、sample() 等操作变成 API 调用，模型在多次调用间常驻 GPU，由服务负责运行时、分布式执行、梯度同步、故障恢复和检查点存储。

推荐理由

官方同时放出可运行 cookbook 和两个实测案例，读者能据此判断自写训练循环的代价与收益。

正文 · AI 翻译

如今微调大型语言模型（LLM）通常要在训练循环控制与运维简便性之间做出取舍。使用全托管服务时，你只需选择基础模型、提供数据和超参数，其余全部由服务负责，当标准化的预置方案适合你的问题时，这种方式效果很好。托管程度较低的服务给你更多控制权，但你也需要管理更多底层基础设施和训练框架。有时标准方案并不够用：你需要为监督微调（SFT）自定义损失函数或数据课程，需要自己的偏好优化、强化学习（RL）奖励函数，或者上周刚发布的对齐算法。一旦你需要自行塑造训练过程，你就得承担这些基础设施工作，以换取算法上的自由。

如今，Amazon SageMaker AI 中的无服务器模型定制能力通过训练会话 API 得到扩展，它将抽象边界下移了一层，从而消除了这一取舍。它们不再抽象掉算法，而是抽象掉计算。forward_backward()、optim_step() 和 sample() 等操作成为 SageMaker API 调用，你可以将它们组合成你的技术所需的任意流程：SFT、偏好优化、RL 或其他新方法。这些调用针对一个活跃的有状态会话，因此模型在各步骤之间保持加载在 GPU 上，你可以对它发起任意多次调用。该服务负责运行时设置、分布式执行、梯度同步、故障恢复、租户隔离以及大规模运行 GPU 的运维工作。你编写训练循环，SageMaker AI 运行基础设施。

除了该服务，我们还发布了 SageMaker 训练会话 cookbook，其中包含一组可直接运行的 SFT 和 RL 方案，你可以将其用作自己训练循环的蓝本。

在本文中，我们介绍编程模型和三个 SDK 客户端，并说明训练会话 API 在 SageMaker AI 的模型训练与定制选项中处于什么位置。然后我们剖析一个针对小学数学问题（GSM8K）的最小 RL 循环，展示这些操作如何组合在一起，最后以 cookbook 中的一个真实工作负载收尾：针对竞赛编程问题的强化学习，其中生成的代码通过运行单元测试来评分。

什么是 SageMaker AI 训练会话 API？

借助 SageMaker AI 训练会话 API，你可以在本地机器的 Python 进程中驱动一个由 GPU 支持的活跃会话。你调用这些操作，SageMaker AI 便在加载了你所选基础模型和低秩适配（LoRA）适配器的托管加速器上运行它们。无需预先描述整个作业的配置文件，也没有集群创建步骤。你可以在代码中交互式地控制每一步发生的事情，从而快速迭代和调试。

职责划分如下：

你专注于 你编写 SageMaker AI 负责
你的数据、RL 环境、智能体和算法 一个 Python 训练循环，为需要 GPU 的步骤调用 forward_backward()、optim_step() 和 sample() 等操作 GPU 分配、模型和 LoRA 加载、分布式执行、故障恢复、检查点存储、安全性

API 围绕 SageMaker Python SDK 中的三个客户端组织，按顺序创建。一个 ServiceClient 持有平台级设置（执行角色、模型包组、Amazon Simple Storage Service (Amazon S3) 输出位置，以及可选的 AWS Key Management Service (AWS KMS) 密钥和虚拟私有云 (VPC) 配置，这些设置你也可以按会话覆盖）并启动一个会话。ServiceClient 实例化一个 TrainingClient（训练和检查点操作，例如 forward、forward_backward、optim_step 和 save_weights_for_sampler），后者又实例化一个 SamplingClient（sample 和 compute_logprobs，用于 rollouts 和评估期间的生成与评分）。每个 API 调用对应你循环中的一个操作。活动会话还提供一个兼容 OpenAI 的聊天补全端点，因此现有的 OpenAI 客户端代码可以从中采样。

SageMaker Python SDK client hierarchy: ServiceClient creates a TrainingClient, which in turn creates a SamplingClient

图 1：用于训练会话 API 的 SageMaker Python SDK 客户端层级

对于强化学习，循环具有一种典型形态。下图展示了一个训练步骤以及各部分运行的位置：绿色的步骤在服务管理的 GPU 上调用 API，紫色的步骤（你的循环控制和奖励逻辑）在你的机器上运行。

One reinforcement learning training step, showing which operations run on managed GPUs and which run on the local client

图 2：一个训练步骤，采样、梯度和检查点操作在受管 GPU 上运行，循环控制和奖励逻辑在你的机器上运行

每一步加载一批提示，并 (1) 在 GPU 上从当前模型采样。然后你的 Python 代码 (2) 对这些输出进行评分：验证器、奖励模型调用、测试框架，任何你能用代码表达的东西。(3) forward_backward 和 optim_step 将评分转化为 GPU 上的梯度更新，(4) save_weights_and_get_sampling_client() 刷新采样器，以便下一步从更新后的模型采样。

你拥有数据、循环控制和奖励；SageMaker AI 运行 GPU 密集型的采样、梯度和检查点操作。相同的原语也涵盖监督微调 (SFT) 和其他 LLM 训练技术。对于 SFT，去掉采样和奖励步骤，并在带标签的样本上使用 cross_entropy 损失调用 forward_backward。

这是为谁准备的，在 Amazon SageMaker AI 中处于什么位置？

当你需要快速迭代周期和算法自由度（例如自定义损失、新的 RL 奖励信号或多阶段流水线），但又不想自行预置或运维 GPU 基础设施时，训练会话 API 是合适的起点。

使用按处理的 token 数计费，因此你只需为循环消耗的计算付费。如果你的奖励函数在你的机器上对一个批次评分需要 10 秒，那么在这运行期间不会计费 GPU 时间。与按小时 GPU 计费相比，这有两个实际优势。第一，没有空闲 GPU 成本：奖励计算、数据加载、日志记录和调试都在不计费的时间进行，只有当你调用 GPU 密集型操作时才计费。第二，成本随你的数据集和生成量扩展，而不是随墙钟时间或集群规模扩展，这使得支出更容易预测，实验更容易比较。对于本地计算繁重的工作负载（例如复杂的奖励模型、多工具智能体 rollouts 或人在回路评估），相比为整个会话预留 GPU，节省可能非常显著。

Training session APIs 专为那些工作由训练算法而非仅由数据集定义的团队而构建。由于循环由你自己编写，同一套原语既支持在标注数据上使用 cross_entropy 损失的监督微调（SFT），也支持使用 ppo、importance_sampling 或 cispo（GRPO、PPO 及其他组相对方法）的强化学习，以及由相同组件构建的自定义目标，例如偏好优化或蒸馏。它们适合：

  • 从事后训练与对齐的研究人员，实现 RL 方法（GRPO、PPO、重要性采样变体）、自定义优势估计器或新目标。
  • 进行监督微调或偏好优化的团队，需要非标准损失、数据课程或多阶段流水线，而这些是配置驱动的配方无法表达的。
  • 需要自定义 rollout 逻辑、奖励塑造和课程的智能体与多轮 RL 构建者。
  • 已不再满足于配置驱动微调、但又不想承担 GPU 集群运维负担的机器学习（ML）工程师。
  • 移植现有训练循环代码的研究人员，因为基于操作的编程模型可直接映射到针对底层 forward_backward、optim_step 和 sample API 编写的循环。

支持的基础模型以及每个模型可用的 AWS 区域列在 training session APIs 文档中。

Training session APIs 在 Amazon SageMaker AI 中的定位

在 SageMaker AI 的模型定制选项中，training session APIs 处于一个独特的位置：对训练算法的控制程度远高于完全托管的 Serverless Model Customization，而运维负担则远低于在 Amazon SageMaker Training Jobs 或 Amazon SageMaker HyperPod 上训练。

选项 对训练的控制 基础设施
SageMaker Serverless Model Customization – 托管 API 低 – 选择模型、配置可用超参数并提供数据 无服务器，无需选择实例
SageMaker Serverless Model Customization – SageMaker AI training session APIs 高 – 你编写循环并调用原语 无服务器，无需选择实例
SageMaker Training Jobs 完全 托管的临时作业，由你选择实例类型
SageMaker HyperPod 完全 托管持久集群，具备高级故障恢复能力

Chart plotting control over the training algorithm against infrastructure you manage, showing training session APIs at high control and low infrastructure relative to other SageMaker AI options

图 3：Training session APIs 在 SageMaker AI 模型定制选项中的定位

训练循环剖析：在数学问题上运行 GRPO

本节拆解一个最小化的强化学习循环，以展示各项操作如何协同工作。任务是 GSM8K，这是一个约 8,500 道小学数学生成题的基准，每道题只有一个数值答案，如下例所示：

Natalia 在四月向她的 48 位朋友出售了发夹，然后在五月她卖出的发夹数量是四月的一半。Natalia 在四月和五月总共卖出了多少发夹？ → 72

我们使用组相对策略优化（GRPO）为 OpenAI 的 GPT-OSS 20B 训练一个 LoRA 适配器：对于每个问题，模型生成多个回答，每个回答都会进行正确性评分，优于组平均水平的回答会被强化。完整可运行的示例位于 Amazon SageMaker 示例仓库中。

在打开会话之前，请完成文档中的先决条件。

打开会话

创建一个会话，然后从中获取分词器和 SamplingClient：

from sagemaker.train.training_session import ServiceClient, SamplingParams, AdamParams

svc = ServiceClient(role_arn=role, region=REGION)

training_client = svc.create_lora_training_client(
    "openai/gpt-oss-20b",                        # a name from svc.list_base_models()
    training_session_name=SESSION_NAME,
    rank=32, alpha=64,                           # LoRA rank and scaling factor
    model_package_group_arn=MODEL_PACKAGE_GROUP_ARN,
    s3_output_path=OUTPUT_S3_PATH,
)
tokenizer = training_client.get_tokenizer()
sampling_client = training_client.create_sampling_client()

定义奖励和优势

任务特定逻辑是运行在你自己机器上的两个短函数：

def compute_reward(response_text, ground_truth):
    predicted = extract_answer(response_text)      # parse the number after "#### "
    return 1.0 if predicted == ground_truth else 0.0

def compute_advantages(rewards):
    rewards = np.array(rewards)
    advantages = rewards - rewards.mean()          # center within the group (the GRPO baseline)
    return advantages / (rewards.std() + 1e-8)     # normalize

循环

每一步都遵循图 2 所示的四个阶段。draw、build_prompt 和 build_training_datum 是示例中的简短辅助函数：

for step in range(NUM_STEPS):
    datums, pending = [], []
    # (1) sample K responses per question in one request each (GPU); returns a future right away
    for question, ground_truth in draw(PROMPTS_PER_STEP):          # data loading is local
        prompt = tokenizer.encode(build_prompt(question), add_special_tokens=False)
        pending.append((prompt, ground_truth, sampling_client.sample(
            prompt, num_samples=K,
            sampling_params=SamplingParams(max_tokens=512, temperature=0.7, top_p=0.9))))

    for prompt, ground_truth, future in pending:
        result = future.result()                                    # wait once the batch is in flight
        # (2) score each response, then center the rewards within the group (local)
        rewards = [compute_reward(tokenizer.decode(s.tokens), ground_truth) for s in result.sequences]
        advantages = compute_advantages(rewards)
        datums += [build_training_datum(prompt, s, a) for s, a in zip(result.sequences, advantages)]

    # (3) one gradient step over the batch (GPU); the optimizer step is sequenced behind it
    fwd_bwd_op = training_client.forward_backward(datums, loss_fn="ppo")
    optim_op = training_client.optim_step(AdamParams(learning_rate=1e-5))
    fwd_bwd_op.result(), optim_op.result()

    # (4) refresh the sampler so the next step's rollouts come from the updated policy (GPU)
    sampling_client = training_client.save_weights_and_get_sampling_client()

training_client.stop()

有三个细节让这个循环得以运转：

操作返回 future。每个绑定 GPU 的调用都会立即返回一个 APIFuture：请求会被排队并由后台提交器发送，而 .result() 会等待该值。循环会先提交某一步的所有采样请求，然后再等待其中任何一个，因此整批请求会并发生成。训练操作按你调用的顺序执行，所以 optim_step 可以在 forward_backward 完成之前提交。

数据项携带损失所需的一切。build_training_datum 打包提示词和响应 token、来自 sample() 的每 token logprobs、优势值，以及一个将提示词排除在损失之外的 weights 掩码。

检查点是模型包。每次保存都会在你的模型包组中注册一个版本，包括上面循环中每步的采样器刷新。训练结束后，另外两次保存会添加一个用于部署的合并模型，以及一个包含权重和优化器的可恢复状态。

training_client.save_weights_for_sampler(merge_weights=True).result()   # merged, deployable model
training_client.save_state().result()                                   # resumable state (weights + optimizer)

要恢复，请使用 create_training_client_from_state_with_optimizer 从最新检查点启动一个新会话。

正在运行的会话会占用你的并发会话配额，直到你调用 stop() 或它过期，所以完成后要停止它，包括在错误路径中，或者把客户端放在 with 块中打开。合并后的检查点可部署到实时端点。若要将检查点作为 LoRA 适配器提供服务，请参阅 cookbook 的部署 notebook。

在我们运行该示例时，对 40 个问题进行 10 个训练步骤，将留出准确率从 56% 提升到了 75%。

真实工作负载：竞赛编程上的强化学习

GSM8K 让机制保持可见，但生产级 RL 工作负载会增加更多内容：更长的生成、多轮交互、外部评分器，以及足以让 GPU 保持忙碌的吞吐量。SageMaker 训练会话 cookbook 将这套工程封装成基于相同操作构建的可复用训练循环。对于新的 RL 任务，你需要编写一个环境：模型看到什么、其答案如何评分，以及它是否能获得下一轮。该 cookbook 处理批处理、聊天模板、优势估计、重叠采样与训练、验证、带恢复的检查点以及每步指标，并记录每次运行的配置以便复现。code_rl 配方用它来在竞赛编程问题上训练模型。

任务。问题来自 DeepCoder 数据集（训练使用 PrimeIntellect、TACO 和 LiveCodeBench；测试使用 Codeforces 和 LiveCodeBench）。模型编写一个 Python 解法，该解法会针对问题的单元测试在你机器上的 SandboxFusion 容器中运行。当所有测试通过时奖励为 1，否则为 0，当响应中不包含代码块时扣 0.1。

多轮工具使用。经过两轮后，模型可以调用一个 check_solution 工具，读取测试结果，并在最终评分前修改答案。工具调用和结果以模型原生的聊天格式呈现，因此模型训练时使用的 token 格式与推理时相同。

该配方在最小循环之上增加了什么：

  • 重叠采样与训练： Rollout 工作线程在训练器运行一步时持续生成。从较旧策略采样的一组响应，如果最多落后 k 步就会被使用，否则会重新排队，并且 cispo 损失会校正策略差异。
  • 流式梯度计算：forward_backward 在各组完成时以子批次形式提交，而不是每个步骤提交一次。
  • 留出集验证、检查点恢复和每一步的指标，可选择发布到 SageMaker 托管 MLflow。

Code RL setup: a local training loop and SandboxFusion container on the left, and managed GPUs running GPT-OSS 20B with a LoRA adapter on the right

图 4：Code RL 设置，左侧为本地训练循环与沙箱，右侧为托管 GPU

我们为 GPT-OSS 20B 训练了一个 LoRA 适配器（秩为 32），训练 50 步，每一步包含 128 个问题和每个问题 8 个回答，使用两轮交互，陈旧度上限为 k=3。训练奖励从 0.46 上升到 0.60，在 200 个留出测试问题上的准确率从 40% 上升到 50.5%。

要运行它，请启动沙箱并启动该配方：

bash tools/sandbox_up.sh
python -m training_sessions_cookbook.recipes.code_rl.train \
    --config training_sessions_cookbook/recipes/code_rl/configs/coderl_oss20b_service.json \
    --max_steps_off_policy 3 --stale_policy requeue --num_rollout_workers 80 --max_steps 50 \
    --base_model_arn "$BASE_MODEL_ARN" \
    --model_package_group_arn "$MODEL_PACKAGE_GROUP_ARN" \
    --role_arn "$ROLE_ARN" \
    --output_s3_path "$OUTPUT_S3_PATH"

该 cookbook 还有更多采用相同结构的配方，包括在聊天数据集（Tulu 3、UltraChat）上进行监督微调、在全规模 GSM8K 上进行 RL，以及基于二元反馈的 KTO 偏好对齐。每个配方的 README 都列出了其命令、超参数和预期指标。

结论

SageMaker AI 训练会话 API 将抽象边界下移到训练操作：你在自己的机器上用 Python 编写训练循环，服务在托管 GPU 上运行它，处理模型和 LoRA 加载、分布式执行、故障恢复和检查点保存。一个完整的 GRPO 循环归结为四个操作，sample、forward_backward、optim_step 和 save_weights_and_get_sampling_client，奖励和优势逻辑在你自己的代码中。如 Code RL 配方所示，这些操作同样可以扩展到带有沙箱评分的多轮、使用工具的 RL。

自定义奖励、智能体式 rollout、逐步调试以及新发布的技术都使用相同的编程模型：算法留在你的循环中，由 SageMaker AI 运行 GPU。

要开始使用：

致谢

感谢 Arun Nagarajan、Xichao Wang 和 Shruti Sharma 对本文的贡献和指导。

关于作者

Nicolas Jourdan

Nicolas Jourdan

Nicolas 是 AWS 驻德国法兰克福的资深 AI/ML 专业解决方案架构师。他支持客户在 AWS 上构建 AI 和机器学习工作负载，从早期设计决策一直到生产。他的重点领域包括模型定制、分布式训练和推理。在加入 AWS 之前，他曾担任研究科学家和工程经理，从事自动驾驶、无人机和制造业应用，并拥有目标检测和异常检测方面的专利。

Bruno Pistone

Bruno Pistone

Bruno 是 AWS 驻意大利米兰的全球资深生成式 AI/ML 专业解决方案架构师。他与 AWS 产品团队和大型客户合作，帮助他们充分理解技术需求，并设计充分利用 AWS 云和 Amazon ML 技术栈的 AI 和机器学习解决方案。他的专长包括分布式训练和推理工作负载、模型定制、生成式 AI 和端到端 ML。他喜欢与朋友共度时光、探索新地方和前往新目的地旅行。

Giuseppe Angelo Porcelli

Giuseppe Angelo Porcelli

Giuseppe 是 Amazon Web Services 的首席机器学习专家解决方案架构师。凭借多年的软件工程和机器学习背景，他与各种规模的客户合作，了解他们的业务和技术需求，并设计能够充分利用 AWS Cloud 和 Amazon Machine Learning 技术栈的 AI 和机器学习解决方案。他参与过不同领域的项目，包括 MLOps、计算机视觉和 NLP，涉及广泛的 AWS 服务。业余时间，Giuseppe 喜欢踢足球。

Pooja Karadgi

Pooja Karadgi

Pooja 是 AWS 的高级产品经理，负责 Amazon SageMaker AI 的模型定制。她专注于让基础模型易于访问并适用于生产环境，通过强化学习等技术帮助客户将基础模型适配到其特定领域。她还参与过 SageMaker Studio 中的开发人员和管理员体验工作，涵盖从实验到生产的机器学习生命周期。

来源：AWS Machine Learning Blog · aws.amazon.com

#产品更新#数据/训练#部署/工程#Agent