Amazon SageMaker AI 推出 training session APIs：自己写训练循环、无需管理 GPU

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

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

伴随该服务，我们还发布了 SageMaker 训练会话手册，其中包含一组可直接运行的 SFT 和 RL 方案，你可以把它们作为自己训练循环的蓝图。

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

什么是 SageMaker AI 训练会话 API？

借助 SageMaker AI 训练会话 API，你可以在本地机器上通过自己的 Python 进程驱动一个活跃的、由 GPU 支撑的会话。你调用这些操作，SageMaker AI 在托管加速器上运行它们，其上已加载你选择的基础模型以及你的低秩适配（LoRA）适配器。没有预先描述整个任务的配置文件，也没有创建集群的步骤。你通过代码以交互方式控制每一步发生什么，因此可以快速迭代和调试。

职责划分如下：

你专注于	你编写	SageMaker AI 负责
你的数据、RL 环境、智能体和算法	一个调用 `forward_backward()`、`optim_step()` 和 `sample()` 等操作的 Python 训练循环，用于需要 GPU 的步骤	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 客户端代码可以直接从中采样。

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

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

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

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

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

在 SageMaker AI 的模型自定义选项中，训练会话 API 处于一个独特的位置：比完全托管的 Serverless Model Customization 对训练算法的控制力强得多，而比在 Amazon SageMaker Training Jobs 或 Amazon SageMaker HyperPod 上训练所需承担的运维责任少得多。

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

图 3：训练会话 API 在 SageMaker AI 模型自定义选项中的定位

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

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

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

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

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

打开会话

创建一个会话，然后获取分词器（tokenizer）和其中的一个 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 的部署笔记本。

在我们运行该示例时，对 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 工作进程在训练器执行一个步骤的同时持续生成。从较旧策略采样的一组响应，只要不超过 k 步之旧就会被使用，否则重新排队，cispo 损失会对策略差异进行校正。
流式梯度计算：各组完成后以子批次提交 forward_backward，而不是每步提交一次。
留出验证、检查点恢复和每步指标，可选择发布到 SageMaker 托管 MLflow。

图 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 进行强化学习，以及基于二元反馈的 KTO 偏好对齐。每个方案的 README 都列出了其命令、超参数和预期指标。

结论

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

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

开始使用：

从 SageMaker 训练会话 cookbook 中选择一个方案，例如 Code RL。
运行 Amazon SageMaker 示例仓库中的 GSM8K 示例，并替换为你自己的奖励函数。
阅读训练会话 API 文档，了解前提条件、配额、区域、支持的模型和 API 参考。

致谢

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