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.

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.

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:

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:

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 format it uses at inference.

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 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.

Figure 2: One training step, with sampling, gradient, and checkpoint operations on managed GPUs and loop control and reward logic on your machine