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

Amazon SageMaker AI has expanded its serverless model customization offering with training session APIs, a new programming model that lets developers write their own training loops in Python while the service handles GPU infrastructure. The APIs, announced alongside a companion cookbook of ready-to-run supervised fine-tuning (SFT) and reinforcement learning (RL) recipes, target teams that need algorithmic freedom beyond what config-driven fine-tuning allows but do not want to provision or operate clusters themselves.

The core idea is a lower abstraction boundary. Where fully managed customization abstracts away the algorithm, the new APIs abstract away the compute. Operations such as `forward_backward()`, `optim_step()`, and `sample()` become SageMaker API calls that developers compose into whatever procedure their technique requires, whether that is SFT, preference optimization, RL, or something newly published. The calls target a live, stateful session, so the model stays loaded on the GPUs across steps and can be driven with as many calls as needed.

Who this is for

The APIs are aimed at post-training and alignment researchers implementing methods such as GRPO, PPO, or importance-sampling variants; teams doing SFT or preference optimization with a non-standard loss, data curriculum, or multi-stage pipeline; agentic and multi-turn RL builders who need custom rollout logic and reward shaping; ML engineers who have outgrown config-driven fine-tuning; and researchers porting existing training-loop code. Supported base models and AWS Regions are listed in the training session APIs documentation.

The three clients

The programming model is organized around three clients in the SageMaker Python SDK, created in sequence. A `ServiceClient` holds platform-level settings — execution role, model package group, Amazon S3 output location, and optionally an AWS KMS key and VPC configuration, which can also be overridden per session — and starts a session. It instantiates a `TrainingClient` for training and checkpointing operations such as `forward`, `forward_backward`, `optim_step`, and `save_weights_for_sampler`, which in turn instantiates a `SamplingClient` for `sample` and `compute_logprobs`, used for generation and scoring during rollouts and evaluation. An active session also serves an OpenAI-compatible chat completions endpoint, so existing OpenAI client code can sample from it.

How one RL step runs

In a reinforcement learning loop, each step loads a batch of prompts and samples from the current model on the GPU. Python code on the user's machine then scores those outputs — a verifier, a reward-model call, a test harness, anything expressible in code. `forward_backward` and `optim_step` turn the scores into a gradient update on the GPU, and `save_weights_and_get_sampling_client()` refreshes the sampler so the next step draws from the updated model. The user owns the data, loop control, and reward; SageMaker AI runs the GPU-bound sampling, gradient, and checkpoint operations. For SFT, the sampling and reward steps are dropped and `forward_backward` is called with a cross-entropy loss over labeled examples.

Billing and cost profile

Usage is billed per token processed, so charges apply only to the compute the loop consumes. If a reward function takes 10 seconds to score a batch locally, no GPU time is billed while it runs. Compared with hourly GPU billing, this removes idle-GPU cost — reward computation, data loading, logging, and debugging happen off the clock — and makes spend scale with dataset and generation volume rather than wall-clock time or cluster size, which the company says makes costs easier to forecast and experiments easier to compare. The savings over reserving GPUs for a 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.

Where it sits among SageMaker AI options

The new APIs occupy a distinct point in the model customization lineup. Fully managed Serverless Model Customization offers low control over the training algorithm: users pick a model, configure available hyperparameters, and supply data, with no instance selection. Training session APIs keep the serverless model but raise control to high, since the user writes the loop and calls primitives. SageMaker Training Jobs and SageMaker HyperPod both offer full control, but require selecting instance types for managed, ephemeral jobs or operating managed, persistent clusters with advanced fault recovery, respectively. In short, the new option sits between managed recipes and full infrastructure ownership.

A minimal GRPO loop on GSM8K

The announcement walks through a minimal RL 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 — for 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?" with the answer 72. The example trains 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.

Opening a session takes a few lines. A `ServiceClient` is created with a role and region, then `create_lora_training_client` is called with the base model name from `list_base_models()`, a session name, LoRA rank 32 and alpha 64, a model package group ARN, and an S3 output path. The training client supplies a tokenizer, and `create_sampling_client()` produces the sampler.

The task-specific logic is two short functions that run locally. `compute_reward` parses the number after "#### " in the response and returns 1.0 for a match with the ground truth, 0.0 otherwise. `compute_advantages` centers rewards within the group — subtracting the mean, the GRPO baseline — and normalizes by the standard deviation plus 1e-8.

Each training step follows four stages. First, for each question the loop encodes a prompt and issues one `sampling_client.sample` request with `num_samples=K`, `max_tokens=512`, `temperature=0.7`, and `top_p=0.9`; the call returns a future immediately so the batch stays in flight. Second, once results arrive, each response is scored and rewards are centered within the group. Third, `forward_backward` and `optim_step` apply the gradient update on the GPU. Fourth, `save_weights_and_get_sampling_client()` refreshes the sampler for the next step. Data loading, prompt building, and datum construction are short local helpers in the sample.

A real-world workload: RL on competitive programming

The post closes with a workload from the cookbook that goes beyond toy math: reinforcement learning on competitive programming problems, where generated code is graded by running it against unit tests. That setup stresses the same division of labor — the test harness runs locally as the reward, while sampling and gradient operations run on managed GPUs — and illustrates why per-token billing matters when scoring is expensive.

What to watch

The practical question for teams is where their bottleneck sits. If a standard recipe fits, fully managed customization remains the simplest path. If the algorithm itself is the work — a custom loss, a new reward signal, a multi-stage pipeline — the new APIs offer a middle ground that removes cluster operations without taking away control of the loop. The cookbook's SFT and RL recipes provide starting blueprints, and the documentation lists supported base models and regional availability, which will determine how quickly teams can adopt the model they want.