跳到正文
AWS Machine Learning Blog· Nicolas Jourdan·· 2 小时前AI 评分58

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

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

AI 导读

Amazon SageMaker AI 推出训练会话 API，把 forward_backward()、optim_step()、sample() 等操作变成 API 调用，模型在 GPU 上保持加载，用户只需在本地 Python 进程里编写训练循环，GPU 分配、分布式执行、故障恢复和检查点存储由服务负责。

正文 · 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，用于 rollout 和评估期间的生成与评分）。每个 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 约束的操作时才计费。第二，成本随数据集和生成量扩展，而不是随墙钟时间或集群规模扩展，这使得支出更容易预测，实验也更容易比较。对于本地计算繁重的工作负载（例如复杂奖励模型、多工具智能体 rollout 或人在回路评估），相比为整个会话预留 GPU，节省可能非常显著。

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

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

支持的基础模型以及各模型可用的 AWS 区域列在训练会话 API 文档中。

训练会话 API 在 Amazon SageMaker AI 中的定位

在 SageMaker AI 的模型定制选项中，训练会话 API 处于一个独特的位置：对训练算法的控制程度远高于全托管的无服务器模型定制，而运维负担则远低于在 Amazon SageMaker 训练作业或 Amazon SageMaker HyperPod 上进行训练。

选项 对训练的控制 基础设施
SageMaker 无服务器模型定制 – 托管 API 低 – 选择模型、配置可用超参数并提供数据 无服务器，无需选择实例
SageMaker 无服务器模型定制 – SageMaker AI 训练会话 API 高 – 你编写循环、调用原语 无服务器，无需选择实例
SageMaker 训练作业 完全 托管、临时作业，由你选择实例类型
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：训练会话 API 在 SageMaker AI 模型定制选项中的定位

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

本节拆解一个最小化的强化学习循环，以展示各项操作如何组合在一起。任务是 GSM8K，一个包含约 8,500 道小学水平数学应用题的基准，每道题只有一个数值答案，如下例所示：

Natalia 在 4 月向她的 48 位朋友出售了发夹，然后在 5 月出售的发夹数量是 4 月的一半。Natalia 在 4 月和 5 月总共出售了多少个发夹？ → 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 logprob、优势值，以及一个将提示词排除在损失之外的 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%。

真实工作负载：在竞赛编程上进行 RL

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 worker 在训练器运行一步时持续生成。从较旧策略采样的一组响应，如果最多旧 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 和 ML 解决方案。他参与过不同领域的项目，包括 MLOps、计算机视觉和 NLP，涉及广泛的 AWS 服务。业余时间，Giuseppe 喜欢踢足球。

Pooja Karadgi

Pooja Karadgi

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

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

开源生态部署工程数据与训练产品更新#产品更新#部署/工程#数据/训练#开源/仓库