用 Python 编排 SageMaker 强化学习训练

微调大模型时，现成配方往往不够灵活，自己写训练循环又得接手 GPU 运维。AWS 这次给出的方案，是把采样、梯度计算和参数更新变成 API：训练流程留在你的 Python 代码里，GPU 交给 SageMaker AI。

它值得关注的地方，是自定义奖励、多轮工具调用和逐步调试终于有了更轻的基础设施入口。但「不用管 GPU」不等于「不用懂训练」，文中的两组效果数据，也需要放回各自的实验条件里看。

自己写训练循环，具体还能控制什么

Amazon SageMaker AI training session APIs 把托管边界下移到了训练操作。过去，全托管服务通常要求你提供模型、数据和超参数，再按既定配方执行；现在，你可以从本地 Python 进程控制一个持续运行、有状态的远程会话，模型及其 LoRA（低秩适配）适配器会在多个步骤之间留在 GPU 上。

你决定下一批用什么数据、生成多少回答、怎样评分、何时更新参数。服务负责运行环境、分布式执行、梯度同步、故障恢复、租户隔离，以及 GPU 运维。

被托管的是计算基础设施，训练算法仍然由你负责。

SDK 用三个依次创建的客户端组织这些能力：

• ServiceClient 管执行角色、model package group、Amazon S3 输出位置，以及可选的 AWS KMS 密钥和 VPC 配置，再启动会话；部分配置也能按会话覆盖。
• TrainingClient 提供 forward、forward_backward、optim_step、save_weights_for_sampler 等训练和保存操作。
• SamplingClient 提供 sample 和 compute_logprobs，负责生成回答及计算对数概率。

活动会话还提供兼容 OpenAI 的 chat completions endpoint，已有 OpenAI 客户端代码可以通过它采样。这项兼容说的是生成接口，不能据此推断训练操作也能直接套用 OpenAI API。

强化学习的一步可以拆成四段：远程生成回答，本地计算奖励，远程计算梯度并更新参数，再刷新采样器。监督微调（SFT）则可以去掉生成和奖励环节，直接对标注样本调用带 cross_entropy 损失的 forward_backward。

它在 SageMaker 产品线中的位置也很清楚：比 Serverless Model Customization 给你更多算法控制权，比 Training Jobs 和 HyperPod 少承担基础设施责任。支持哪些基础模型、哪些 AWS Regions 可用，原文要求查文档，没有列出完整清单。

四个操作能串起训练，但数据不能随便拼

作者用 GSM8K 演示最小强化学习循环。这个基准约有 8,500 道小学数学应用题，每题有一个数值答案。实验针对 openai/gpt-oss-20b 训练 LoRA，设置 rank=32、alpha=64，使用 GRPO（组相对策略优化）。

做法是让模型对同一道题生成多个回答，按正确与否给出 1.0 或 0.0，再减去组内奖励均值，并除以奖励标准差加 1e-8。高于组内平均水平的回答获得正向强化。这里需要自己定义的，是答案提取、奖励和训练数据构造。

下面这段代码展示了生成、评分、更新和刷新采样器如何接起来：

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()

要盯住 save_weights_and_get_sampling_client() 这一行：更新训练权重后，需要刷新采样器，下一步才能从更新后的策略生成回答。循环使用 loss_fn="ppo"，并不妨碍它按文中方式实现 GRPO 的组内优势计算。

异步行为也很关键。GPU 操作立即返回 APIFuture，调用 .result() 才等待结果。代码先提交一批采样请求，再逐个等待，使这些请求能够并发生成。训练操作则按调用顺序执行，所以可以在 forward_backward 完成前提交 optim_step。

容易被短代码遮住的是 build_training_datum。它必须打包提示词和回答的 token、采样时逐 token 的 logprobs、优势值，以及排除提示词损失的权重掩码。API 简化了执行，却没有替你处理这些训练语义。

保存也有成本之外的管理含义：每次保存都会在 model package group 注册版本，包括逐步刷新采样器。训练后，save_weights_for_sampler(merge_weights=True).result() 保存可部署的合并模型，save_state().result() 保存权重和优化器状态；恢复使用 create_training_client_from_state_with_optimizer。会话在 stop() 或到期前持续占用并发配额，异常路径也需要清理，或使用 with 管理客户端。

从数学题到代码训练，多出来的是评测和调度

作者报告，GSM8K 示例在 40 道题上训练 10 步后，留出集准确率从 56% 升到 75%。这证明示例可以产生可观察的变化，但原文没有交代这组留出集大小、重复实验结果和波动范围。

更接近复杂任务的是 cookbook 中的 code_rl 配方。训练题来自 DeepCoder 数据集中的 PrimeIntellect、TACO 和 LiveCodeBench，测试使用 Codeforces 和 LiveCodeBench。模型生成 Python 解法，在本机的 SandboxFusion 容器中执行单元测试。

全部测试通过得到奖励 1，否则为 0；回答没有代码块时另有 0.1 惩罚。两轮交互允许模型调用 check_solution，读取测试结果后修改答案。工具调用和结果使用模型原生聊天格式，训练与推理使用相同的 token 格式。

为了提高吞吐，采样与训练会重叠运行。旧策略生成的回答，只要落后不超过 k 步就可以使用，否则重新排队；cispo 损失用于校正策略差异。完成的回答组还会按子批次提交 forward_backward，不必等整步数据全部到齐。

作者给出的实验配置是：GPT-OSS 20B、LoRA rank 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"

其中 --max_steps_off_policy 3 和 --stale_policy requeue 决定旧策略样本如何处理；--num_rollout_workers 80 是示例设置，不能直接当成所有团队的推荐值。

cookbook 还封装了批处理、聊天模板、优势估计、验证、断点恢复和逐步指标，可选接入 SageMaker managed MLflow；其他配方覆盖 Tulu 3、UltraChat 的 SFT、完整规模 GSM8K RL，以及基于二元反馈的 KTO 偏好对齐。

我的判断

我认可它把奖励逻辑留在用户代码里的取舍。测试执行、外部工具调用、多轮交互，很难都塞进一份固定配置。对于已经需要改训练算法的团队，这种边界比继续扩展配置文件更自然。

按处理 token 计费，也确实对本地评分耗时长的任务有吸引力。原文明确说，本地奖励计算耗时 10 秒，这段等待不计 GPU 时间；数据加载、日志和调试同理。

但「更省」仍需要账单验证。原文没有给出 token 单价、两组实验总成本，也没有与按小时租 GPU 做同条件对照。复杂奖励模型如果需要额外计算资源，其费用仍要另算。多轮生成和每题多个回答也会增加处理量。

这两组结果支持配方具备可行性，还不足以证明它比其他训练方案更划算。

此外，算法控制权有边界：文章展示的是受支持基础模型上的 LoRA，以及服务暴露的训练操作。它没有说明是否支持任意模型结构、任意底层算子或所有训练方式。普通业务若标准 SFT 已经够用，全托管配方仍可能更省工程精力。

对程序员而言，更有价值的启发是评测设计。代码能通过一组测试，是一个可计算的奖励，但测试覆盖不到的行为，也不会自动进入优化目标。把测试变成训练反馈前，要先确认测试是否代表你真正需要的能力。

能拿来做什么

• 如果你已经在做代码模型后训练，可以先参考「生成 Python 解法、执行单元测试、返回奖励」这条链路，优先检查评测器是否可靠，再增加多轮工具交互。
• 如果现有微调平台卡在自定义损失、数据课程或多阶段流程，可以对照 sample、forward_backward、optim_step 等操作，判断自己的算法能否表达，再核实模型、Region 和配额。
• 如果准备复现 GSM8K 示例，应先跑通原始奖励与数据构造，再替换自己的奖励函数，同时检查 logprobs、优势值和提示词掩码是否正确。
• 如果要接入团队实验流程，应同时记录配置、留出集指标、处理 token 和实际费用，并把会话异常清理、检查点恢复纳入运行脚本。