Artificial Intelligence

Amazon SageMaker Introduces Training Session APIs for LLMs

Amazon SageMaker AI has introduced new training session APIs that enable developers and researchers to write custom large language model training loops in Python while the service manages the underlying GPU infrastructure.

Radwan AbdulhadiEditor and technology writer

Published: October 22, 2026 · Updated: October 22, 2026 · 3 min read

Amazon SageMaker Introduces Training Session APIs for LLMs — Image: AWS Machine Learning Blog

Bridging Algorithmic Control and Operational Simplicity

Fine-tuning large language models traditionally forces a difficult tradeoff between training loop control and operational simplicity. Fully managed services provide prepackaged recipes, whereas less managed services grant control at the expense of infrastructure and framework management. To solve this dilemma, Amazon SageMaker AI has expanded its serverless model customization options with new training session APIs that abstract away the compute layer rather than the algorithm.

Instead of pre-describing an entire job via static configuration files or managing cluster creation steps, developers can drive live, stateful sessions directly from a local Python process. Operations like forward_backward(), optim_step(), and sample() become straightforward API calls that can be composed into any required procedure, ranging from supervised fine-tuning and preference optimization to advanced reinforcement learning techniques.

SageMaker Python SDK client hierarchy: ServiceClient creates a TrainingClient, which in turn creates a SamplingClient — Image related to the report from AWS Machine Learning Blog · Source: AWS Machine Learning Blog

Understanding the Three-Client SDK Hierarchy

The training session architecture is organized around three sequential clients available in the SageMaker Python SDK. First, a ServiceClient manages platform-level settings such as execution roles, Amazon S3 output locations, and optional VPC configurations while initiating the session. This client then instantiates a TrainingClient, which handles training and checkpointing tasks like forward, forward_backward, optim_step, and save_weights_for_sampler.

Finally, the TrainingClient instantiates a SamplingClient to manage generation and scoring tasks through sample and compute_logprobs functions during evaluations and rollouts. Furthermore, active sessions expose an OpenAI-compatible chat completions endpoint, allowing existing OpenAI client code to seamlessly sample from the loaded model.

One reinforcement learning training step, showing which operations run on managed GPUs and which run on the local client — Image related to the report from AWS Machine Learning Blog · Source: AWS Machine Learning Blog

Architecture of a Reinforcement Learning Loop

Reinforcement learning workflows follow a distinct pattern utilizing both managed infrastructure and local compute. Each training step loads a batch of prompts, samples outputs from the model running on managed GPUs, and passes those outputs back to the developer's local machine for scoring through a reward function, test harness, or verifier.

Once scored, the local code invokes forward_backward and optim_step operations to convert the evaluation scores into gradient updates on the GPU. Finally, a checkpointing call refreshes the active sampler so that subsequent iterations draw generations from the newly updated model parameters.

Chart plotting control over the training algorithm against infrastructure you manage, showing training session APIs at high control and low infrastructure relative to other SageMaker AI options — Image related to the report from AWS Machine Learning Blog · Source: AWS Machine Learning Blog

Token-Based Billing and Economic Advantages

A notable aspect of the training session APIs is their token-based billing structure. Rather than paying for continuous hourly GPU reservations, users are billed exclusively for the specific tokens processed during active GPU operations. If local code takes ten seconds to compute rewards or process data, no GPU costs accrue during that downtime.

This model eliminates idle-GPU expenses, making expenditures easier to forecast and compare across experiments. The financial efficiency is especially pronounced for workloads characterized by heavy local computation, including complex reward models, human-in-the-loop evaluations, and multi-tool agentic rollouts.

Code RL setup: a local training loop and SandboxFusion container on the left, and managed GPUs running GPT-OSS 20B with a LoRA adapter on the right — Image related to the report from AWS Machine Learning Blog · Source: AWS Machine Learning Blog

Resources and Pre-Packaged Blueprints

To help developers get started, AWS released the SageMaker training sessions cookbook alongside the service. This repository offers ready-to-run recipes and blueprints covering supervised fine-tuning and reinforcement learning workloads, such as reinforcement learning applied to competitive programming where generated code is evaluated against strict unit tests.

Additional concrete examples and community resources can also be found across related repositories, including implementations focused on grade-school math problems and group-relative policy optimization techniques.