SageMaker lets teams write their own training loops without managing GPUs

Abstracting the compute, not the algorithm

AWS describes a familiar problem in LLM fine-tuning. Fully managed services work well when a standard recipe fits, but they leave little room for a custom loss, a data curriculum, your own preference optimization or a reinforcement learning (RL) reward function. Less managed options give more control, but the user then runs more of the infrastructure and training frameworks.

The new training session APIs draw the boundary one level lower. Operations such as forward_backward(), optim_step() and sample() become SageMaker API calls that developers compose into supervised fine-tuning (SFT), preference optimization, RL or a new technique. The calls target a live, stateful session, so the model stays loaded on the GPUs across steps.

According to AWS, the service handles runtime setup, distributed execution, gradient synchronization, fault recovery, tenant isolation and the operation of GPUs at scale. The user writes the loop. AWS also released the SageMaker training sessions cookbook, a set of ready-to-run SFT and RL recipes meant as blueprints.

How the programming model works

Developers drive a GPU-backed session from their own Python process on a local machine. The session runs a chosen base model plus a Low-Rank Adaptation (LoRA) adapter. There is no up-front config file describing the whole job and no cluster creation step.

The SageMaker Python SDK exposes three clients, created in sequence:

• ServiceClient holds platform settings such as the execution role, model package group and Amazon S3 output location, and can include an AWS KMS key and VPC configuration. It starts a session.
• TrainingClient covers training and checkpointing operations such as forward, forward_backward, optim_step and save_weights_for_sampler.
• SamplingClient offers sample and compute_logprobs 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.

In an RL step, the model samples outputs on the GPU, the user's code scores them with a verifier, reward-model call or test harness, and forward_backward plus optim_step turn the scores into a gradient update. The call save_weights_and_get_sampling_client() then refreshes the sampler. For SFT, the sampling and reward steps are dropped and forward_backward is called with a cross_entropy loss over labeled examples.

Per-token billing and who it targets

AWS says usage is billed per token processed. Time spent on reward computation, data loading, logging or debugging on the user's machine is not billed as GPU time. The company argues this removes idle-GPU cost and makes spend easier to forecast, since cost scales with dataset and generation volume rather than wall-clock time or cluster size. It says savings can be significant for workloads with heavy local computation, such as complex reward models, multi-tool agentic rollouts or human-in-the-loop evaluation.

The same primitives support SFT, RL with ppo, importance_sampling or cispo, and custom objectives such as preference optimization or distillation. AWS lists its intended users as post-training and alignment researchers, teams needing non-standard losses or multi-stage pipelines, agentic and multi-turn RL builders, ML engineers who have outgrown config-driven fine-tuning, and researchers porting existing training-loop code.

Where it sits in SageMaker AI

AWS positions the APIs between its existing options. Serverless Model Customization with managed APIs offers low control over training. The training session APIs give high control while remaining serverless with no instance selection. SageMaker Training Jobs and SageMaker HyperPod offer full control, but the user selects instance types or runs persistent clusters.

For teams that need algorithmic freedom but do not want to provision GPUs, this gives a middle path. Supported base models and Regions are listed in the training session APIs documentation.