使用 Amazon SageMaker 训练课程 API 写入自己的训练循环——不管理 GPU

人工智能

使用 Amazon SageMaker 训练课程 API 写入自己的训练循环——不管理 GPU

微调一种大型语言模型(LLM)今天通常强制在训练循环控制与操作简便之间进行权衡。 有了完全管理好的服务后，您会选择一个基础模型，提供数据和超参数，服务拥有其他一切，当一个标准，预先包装的配方适合您的问题时，这个服务效果很好。 管理较少的服务会给你更多的控制，但你也管理更多的基础基础设施和培训框架。 有时标准食谱还不够：您需要自定义的丢失或数据课程，一旦你需要自己塑造训练程序，你便承担基础设施的工作，以获得算法自由。

如今，Amazon SageMaker AI 扩展中的无服务器模型定制与培训会话API，通过将抽象边界绘制为更低一层来消除这种权衡。 与其抽象地去掉算法，不如抽象地去掉计算。 前向后向( ), opatim step( )，和样本( ) 等操作成为 SageMaker API 的调用，由您编写到您技术需要的任何程序：SFT，偏好优化，RL，或者新的东西。 通话目标是一个活的，状态良好的会话，所以模型在GPU上停留在跨台阶的加载上，你按需要驱动尽可能多的通话。 该服务处理运行时间设置，分布执行，梯度同步，断层恢复，租户隔离，以及运行GPU的操作规模。 你写训练环，和SageMakerAI运行基础设施。

SageMaker训练课程的烹饪手册，一套随时可运行的 SFT 和 RL 食谱

在此文章中，我们引入了编程模型和三个SDK客户端，并显示培训课API在SageMaker AI的模型培训和定制选项中的位置。 然后我们拆开一个关于年级数学问题的最小RL循环(GSM8K)，以显示操作是如何组合起来的，并用烹饪本上的现实世界工作量完成：增强对竞争性编程问题的学习，通过运行单位测试来分级生成的代码。

SageMaker AI培训课的API是什么？

通过 SageMaker AI 训练课程 APIs，您可以在本地机器上从自己的Python 进程驱动一个实时的，GPU支持的会话。 SageMaker AI 运行在装有你所选择的基础模型的管理加速器上加上你的低亮适应(LoRA)适配器没有描述整个工作的配置文件，你从代码中交互控制每个步骤发生的事情，这样你就可以快速地进行脚踏和调试。

职责分工是这样的：

API在SageMaker Python SDK中围绕三个客户组织，按顺序创建。 A ServiceClient 持有平台级设置(执行角色，模型包组，Amazon简单存储服务(Amazon S3)输出位置，以及可选的AWS密钥管理服务(AWS KMS)密钥和虚拟私有云(VPC)配置，您也可以在每会话中推翻)并开始一个会话。 ServiceClient 即时操作 a Trainclient(培训和检查操作，如前向，前向后向，优化步骤，以及保存 qu量 for sampler)，这反过来会即时操作 样板操作(样本和计算 logprobs，用于在推出和评价时生成和评分). API 调用地图到您循环中的一个操作。 一个活动会话还服务于OpenAI兼容的聊天完成端点，因此现有的OpenAI客户端代码可以从中取样。

对于强化学习，循环具有特征形状。 下图显示 1 个训练步骤以及每个部分运行的地方：在服务管理的GPU上使用绿色调用API的步数，以及紫色的步数，你的循环控制和奖励逻辑，运行在你的机器上。

每个步骤都从GPU上的当前模型中装入一批 prompt和(1)样本。 您的 Python 代码然后(2) 对这些输出进行评分：验证器、 奖励模式调用、 测试控制带、 任何您可以用代码表达的内容。 (3) 向后退和 optimm step 将分数转换为 GPU 的梯度更新，(4) 保存重量和给样本客户端() 刷新采样器，以便下一步从更新的模型中提取。

您拥有数据、 循环控制和奖励；SageMaker AI 运行 GPU 的采样、 梯度和检查点操作。 同样的原始人覆盖了监督的微调(SFT)和其他LLM训练技术。 对于 SFT，将采样和奖励步骤放下，并用一个横跨的 entropy 丢失来调用标签上的例子。

Amazon SageMaker AI中

当您需要快速的迭代周期和算法自由时，训练课程API 是正确的起点，例如自定义丢失、 新的 RL 奖励信号、 或多阶段管道，但不想自己提供或操作 GPU 基础设施。

使用费每处理一个符号就收费，如果你的奖励功能需要10秒才能在你的机器上分出一个批次，则在运行时没有GPU时间计费。 与小时GPU计费相比，这有两个实际好处。 首先，没有闲置的GPU成本：奖励计算、数据加载、记录和调试都发生在时钟之外，你只有在调用GPU捆绑操作时才计费。 第二，成本尺度与你的数据集和生成量，而不是与时钟时间或集群大小，这使得支出更容易预测和实验比较。 为整个会话保留GPU所产生的节余对于当地计算繁重的工作量可能很大，例如复杂的奖励模式、多工具代理推出，或人员即时评价。

培训课程 API 是为工作由培训算法，而不仅仅是数据集定义的团队而构建的。 因为你自己编程环路，同样的原始人支持监督微调(SFT)，在标注数据上带有交叉 entropy损失，加固学习用po，重要性采样，或cispo(GRPO,PPO，以及其他类群关系方法)，以及自定义的目标，如偏好优化或蒸馏，由同一块构建。 它们适合：

- 实施RL方法(GRPO、PPO、重要抽样变体)、定制优势估计器或新目标的培训和调整后研究人员。

需要非标准损失、数据课程或多阶段管道

- 代理和多回合RL建设者，他们需要定制的推出逻辑、奖励塑造和课程。

- 机器学习(ML)工程师

- 研究人员移植现有的训练-浏览码，因为基于操作的编程模型地图直接写入了针对低级前向、优化步骤和API样本的循环。

培训课的API文件列出了所支持的基础模型和AWS区域。

Amazon SageMaker AI范围内的培训课程

APIs在SageMaker AI的模型定制选项中处于一个独特的位置：对培训算法的控制大大高于完全管理的无服务器模式定制，业务所有权远低于Amazon SageMaker培训工作或Amazon SageMaker HyperPod的培训。

解剖一个训练循环：关于数学问题的GRPO

本节拆分一个最小的强化学习循环，以显示操作如何组合。 任务为GSM8K，一个约8500个年级数学字问题的基准测试，每个字都有单数答案，如下例所示：

娜塔莉娅在4月向48位朋友出售剪辑，之后她在5月卖出了一半的剪辑。 4月和5月Natalia总共卖了多少片段？

我们为OpenAI的GPT-OSS 20B训练LoRA适配器，采用集团相对政策优化(GRPO)：对于每个问题，模型产生多个响应，每个响应的得分正确，击败集团平均水平的响应得到强化。 完整，可运行的例子在 Amazon SageMaker 示例库中。

在开始开会之前，请在文件中填写先决条件。

开幕会话

创建会话，然后从中获取标定器和采样机：

确定奖励和好处

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

循环

每个步骤都遵循图2所示的四个阶段，绘图，构建即时，构建训练 datum是样本中的短助手：

三个细节使这个循环工作：

业务还原期货。 每个 GPU 绑定的调用立即返回一个 APIFuture：请求由背景提交者排队发送，并且。 results () 等待值。 循环在等待任何步骤之前提交所有步骤的取样请求，所以整个批次同时生成。 训练操作按您调用的顺序运行，因此在前向后向完成前可以提交 optim step。

基准携带损失所需的一切。 建立 training datum 将快速和响应符，样本中的每根日志()，优势，以及一个能使快速排除丢失的权重掩码打包。

检查点是模型包。 每个保存都会在您的模型包组中注册一个版本，包括上面循环中刷新的每步采样器。 经过训练，再省两笔资金，增加了一个部署的合并模型和一个带有重量和优化器的可恢复状态。

要恢复，请从最近的检查点开始新的会话，并使用创建 training client from state with optiger。

运行的会话会违反您的并行会话配额，直到您调用停止( ) 或到期为止，所以当您完成时停止它，包括在错误路径中，或在块中打开客户端。 合并后的检查站部署在实时端点。 将检查点作为LoRA适配器

在我们的抽样中，关于40个问题的10个培训步骤被搁置的精确度从56%到75%。

现实世界的工作量：竞争性方案拟订的RL

GSM8K保持了力学的可见性，但生产RL工作量新增更多：更长的代数，多转动的相互作用，一个外部的分级器，以及足够的吞吐量使GPU保持繁忙。 SageMaker培训课的烹饪手册套件，在同样的操作基础上进行再利用培训循环。 对于新的RL任务，您会写一个环境：模型所看到的，其答案如何得分，以及它是否得到另一个转折。 烹饪手册处理批量，聊天模板，优势估计，重复采样和培训，验证，检查与恢复，以及每步的度量标准，并记录每个运行的配置，以便重复。 code rl配方使用它来训练一种关于竞争性编程问题的模型。

任务问题来自DeepCoder数据集(用于培训的PrimeIntellect,TACO，和LiveCode Bench；用于测试的Codeforce和LiveCode Bench). 该模型会写一个 Python 解决方案，它与您机器上的 Sandbox Fusion 容器中的问题单元测试相悖。 当所有测试通过时，奖励为1分，否则为0分，答复中未包含代码块时，奖励为0.1分。

多回合工具的使用。 通过两个转弯，模型可以调用检查溶解工具，读取测试结果，并在最终分级前修改答案。 工具调用和结果以模型的本地聊天格式提供，所以模型的训练是相同的，用于

本文由 AI图鉴自动抓取并翻译，版权归原作者所有；如需查看原文请点下方链接。

Artificial Intelligence

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

Fine-tuning a large language model (LLM) today usually forces a tradeoff between training loop control and operational simplicity. With fully managed services you select a base model, provide data and hyperparameters, and the service owns everything else, which works well when a standard, prepackaged recipe fits your problem. Less managed services give you more control, but you also manage more of the underlying infrastructure and training frameworks. Sometimes a standard recipe isn’t enough: you need a custom loss or data curriculum for supervised fine-tuning (SFT), your own preference optimization, a reinforcement learning (RL) reward function, or an alignment algorithm published last week. As soon as you need to shape the training procedure yourself, you take on that infrastructure work to gain algorithmic freedom.

Today, serverless model customization in Amazon SageMaker AI expands with training session APIs, which remove that tradeoff by drawing the abstraction boundary one level lower. Instead of abstracting away the algorithm, they abstract away the compute. Operations like forward_backward(), optim_step(), and sample() become SageMaker API calls that you compose into whatever procedure your technique requires: SFT, preference optimization, RL, or something new. The calls target a live, stateful session, so the model stays loaded on the GPUs across steps, and you drive as many calls as you need against it. The service handles runtime setup, distributed execution, gradient synchronization, fault recovery, tenant isolation, and the operations of running GPUs at scale. You write the training loop, and SageMaker AI runs the infrastructure.

Alongside the service, we’re releasing the SageMaker training sessions cookbook, a set of ready-to-run SFT and RL recipes that you can use as blueprints for your own training loops.

In this post, we introduce the programming model and the three SDK clients, and show where training session APIs fit among the model training and customization options in SageMaker AI. We then take apart a minimal RL loop on grade-school math problems (GSM8K) to show how the operations fit together and finish with a real-world workload from the cookbook: reinforcement learning on competitive programming problems, where generated code is graded by running it against unit tests.

What are SageMaker AI training session APIs?

With SageMaker AI training session APIs, you can drive a live, GPU-backed session from your own Python process on your local machine. You call the operations, and SageMaker AI runs them on managed accelerators loaded with your chosen base model plus your Low-Rank Adaptation (LoRA) adapter. There’s no config file describing the whole job up front and no cluster creation step. You control what happens at each step, interactively, from code, so you can iterate and debug quickly.

The division of responsibilities looks like this:

The API is organized around three clients in the SageMaker Python SDK, created in sequence. A ServiceClient holds the platform-level settings (execution role, model package group, Amazon Simple Storage Service (Amazon S3) output location, and optionally an AWS Key Management Service (AWS KMS) key and virtual private cloud (VPC) configuration, which you can also override per session) and starts a session. The ServiceClient instantiates a TrainingClient (training and checkpointing operations such as forward, forward_backward, optim_step, and save_weights_for_sampler), which in turn instantiates a SamplingClient (sample and compute_logprobs, for generation and scoring during rollouts and evaluation). Each API call maps to one operation in your loop. An active session also serves an OpenAI-compatible chat completions endpoint, so existing OpenAI client code can sample from it.

For reinforcement learning, the loop has a characteristic shape. The following figure shows one training step and where each part runs: the steps in green call APIs on service-managed GPUs, and the steps in purple, your loop control and reward logic, run on your machine.

Each step loads a batch of prompts and (1) samples from the current model on the GPU. Your Python code then (2) scores those outputs: a verifier, a reward-model call, a test harness, anything you can express in code. (3) forward_backward and optim_step turn the scores into a gradient update on the GPU, and (4) save_weights_and_get_sampling_client() refreshes the sampler so the next step draws from the updated model.

You own the data, the loop control, and the reward; SageMaker AI runs the GPU-bound sampling, gradient, and checkpoint operations. The same primitives cover supervised fine-tuning (SFT) and other LLM training techniques. For SFT, drop the sampling and reward steps and call forward_backward with a cross_entropy loss over labeled examples.

Who is this for and where does it fit within Amazon SageMaker AI?

Training session APIs are the right starting point when you need fast iteration cycles and algorithmic freedom, such as a custom loss, a new RL reward signal, or a multi-stage pipeline, but don’t want to provision or operate GPU infrastructure yourself.

Usage is billed per token processed, so you’re charged only for the compute your loop consumes. If your reward function takes 10 seconds to score a batch on your machine, no GPU time is billed while it runs. Compared with hourly GPU billing, this has two practical advantages. First, there’s no idle-GPU cost: reward computation, data loading, logging, and debugging happen off the clock, and you’re billed only when you call a GPU-bound operation. Second, cost scales with your dataset and generation volume, not with wall-clock time or cluster size, which makes spend easier to forecast and experiments easier to compare. The savings over reserving GPUs for the full session can be significant for workloads with heavy local computation, such as complex reward models, multi-tool agentic rollouts, or human-in-the-loop evaluation.

Training session APIs are built for teams whose work is defined by the training algorithm, not only the dataset. Because you compose the loop yourself, the same primitives support supervised fine-tuning (SFT) with a cross_entropy loss over labeled data, reinforcement learning with ppo, importance_sampling, or cispo (GRPO, PPO, and other group-relative methods), and custom objectives, such as preference optimization or distillation, built from the same pieces. They fit:

- Post-training and alignment researchers implementing RL methods (GRPO, PPO, importance-sampling variants), custom advantage estimators, or new objectives.

- Teams doing supervised fine-tuning or preference optimization who need a non-standard loss, a data curriculum, or a multi-stage pipeline that config-driven recipes don’t express.

- Agentic and multi-turn RL builders who need custom rollout logic, reward shaping, and curricula.

- Machine learning (ML) engineers who’ve outgrown config-driven fine-tuning but don’t want to own the operational burden of GPU clusters.

- Researchers porting existing training-loop code, since the operation-based programming model maps directly onto loops written against low-level forward_backward, optim_step, and sample APIs.

The supported base models and the AWS Regions where each is available are listed in the training session APIs documentation.

Where training session APIs fit within Amazon SageMaker AI

Training session APIs sit at a distinct point among the model customization options in SageMaker AI: substantially more control over the training algorithm than fully managed Serverless Model Customization, and far less operational ownership than training on Amazon SageMaker Training Jobs or Amazon SageMaker HyperPod.

Anatomy of a training loop: GRPO on math problems

This section takes apart a minimal reinforcement learning loop to show how the operations fit together. The task is GSM8K, a benchmark of roughly 8,500 grade-school math word problems, each with a single numeric answer, as in the following example:

Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May? → 72

We train a LoRA adapter for OpenAI’s GPT-OSS 20B with Group Relative Policy Optimization (GRPO): for each question the model generates several responses, each response is scored for correctness, and responses that beat the group’s average are reinforced. The complete, runnable example is in the Amazon SageMaker examples repository.

Before you open a session, complete the prerequisites in the documentation.

Open a session

Create a session, then get the tokenizer and a SamplingClient from it:

Define the reward and the advantage

The task-specific logic is two short functions that run on your machine:

The loop

Each step follows the four stages shown in Figure 2. draw, build_prompt, and build_training_datum are short helpers in the sample:

Three details make this loop work:

Operations return futures. Every GPU-bound call returns an APIFuture immediately: the request is queued and sent by a background submitter, and .result() waits for the value. The loop submits all of a step’s sampling requests before waiting on any of them, so the whole batch generates concurrently. Training operations run in the order you call them, so optim_step can be submitted before forward_backward has finished.

The datum carries everything the loss needs. build_training_datum packs the prompt and response tokens, the per-token logprobs from sample(), the advantage, and a weights mask that keeps the prompt out of the loss.

Checkpoints are model packages. Every save registers a version in your model package group, including the per-step sampler refresh in the loop above. After training, two more saves add a merged model for deployment and a resumable state with weights and optimizer.

To resume, start a new session from the latest checkpoint with create_training_client_from_state_with_optimizer.

A running session counts against your concurrent-session quota until you call stop() or it expires, so stop it when you finish, including in error paths, or open the client in a with block. The merged checkpoint deploys to a real-time endpoint. To serve checkpoints as LoRA adapters instead, see the cookbook’s deployment notebook.

In our run of the sample, 10 training steps on 40 problems took held-out accuracy from 56 percent to 75 percent.

A real-world workload: RL on competitive programming

GSM8K keeps the mechanics visible, but a production RL workload adds more: longer generations, multi-turn interaction, an external grader, and enough throughput to keep the GPUs busy. The SageMaker training sessions cookbook packages that engineering into reusable training loops built on the same operations. For a new RL task you write an environment: what the model sees, how its answer is scored, and whether it gets another turn. The cookbook handles batching, chat templates, advantage estimation, overlapped sampling and training, validation, checkpointing with resume, and per-step metrics, and records each run’s configuration so it can be repeated. The code_rl recipe uses it to train a model on competitive programming problems.

The task. Problems come from the DeepCoder dataset (PrimeIntellect, TACO, and LiveCodeBench for training; Codeforces and LiveCodeBench for testing). The model writes a Python solution, which runs against the problem’s unit tests in a SandboxFusion container on your machine. The reward is 1 when all tests pass and 0 otherwise, with a 0.1 penalty when the response contains no code block.

Multi-turn tool use. With two turns, the model can call a check_solution tool, read the test results, and revise its answer before the final grade. The tool calls and results are rendered in the model’s native chat format, so the model is trained on the same token for

Auto-fetched and translated by AI Atlas. All rights belong to the original author.