# SageMaker AI Training Sessions: Custom Training Loops Without Managing GPUs

AWS now lets developers control each training step through APIs, changing where the work and cost of model customization sit.

By Paras Madan | 2026-10-23


AWS made SageMaker AI training sessions generally available on October 22. Developers can write their own model-training loop while AWS runs the GPUs. The useful change is control over individual training operations, including custom rewards, without taking responsibility for the cluster underneath them.

The [launch announcement](https://aws.amazon.com/blogs/machine-learning/announcing-new-ways-to-customize-models-with-sagemaker-ai-without-managing-gpus/) expands SageMaker AI's serverless model customization. AWS handles GPU allocation, distributed execution, recovery and checkpoint storage. Your code decides what the model learns from and how the next training step should work.

## What changes from managed fine-tuning

A managed fine-tuning recipe is convenient when the available settings describe your experiment. Training sessions go lower: a Python program calls operations against a model that stays loaded between requests.

AWS's [technical walkthrough](https://aws.amazon.com/blogs/machine-learning/write-your-own-training-loop-without-managing-gpus-with-amazon-sagemaker-training-sessions/) describes three clients. ServiceClient creates the session, TrainingClient handles training and checkpoints, and SamplingClient generates responses and computes log probabilities. An active session also exposes an OpenAI-compatible chat completions endpoint for sampling.

For reinforcement learning, the program generates responses, scores them, calls forward_backward to calculate gradients, then optim_step to update weights. It refreshes the sampler before generating the next batch. The scoring logic runs in your environment, where it can use a verifier or test harness.

**The decision about what counts as a good answer remains yours.** That matters for a coding agent: passing a test is measurable, but a weak test suite can still reward an implementation that fails outside the training examples. Moving the compute does not fix that problem.

## Supported models and Regions

The [training session documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/customizing-models-training-sessions.html) currently lists these combinations:

| Model | AWS Regions listed |
| --- | --- |
| GPT-OSS-20B | Northern Virginia, Oregon and Mumbai |
| Gemma-4-31B-it | Oregon and Mumbai |
| Qwen-3.6-27B | Oregon and Mumbai |

The launch announcement additionally names Qwen3.8-27B, which is not yet in that documentation table. Check the service's current model list before planning a run around it. Announcement-level support should not be read as confirmation of availability in every Region.

The documented default limit is two concurrent training sessions, with quota increases available. A session continues to consume that quota until stopped or expired, so cleanup belongs in the training program's error handling too.

## How training-session pricing works

The documentation separates billing into prefill, sample and train: input processing, generated rollout tokens and the learning operation. AWS says the GPU meter stops while your reward function scores a batch in your own environment.

That is useful when an experiment spends time waiting for something other than the model. But **token-based billing does not make the whole experiment free while it waits**. Your verifier, storage and any external model calls still need their own cost accounting.

There is no single useful price for a complete training run: its cost depends on the chosen model and the mix of these operations. Check the applicable rates and estimate all three token categories before comparing a session with an hourly training job.

## What you still need to build

AWS's [setup documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/training-sessions-prereqs.html) requires a caller identity, a session execution role and a model package group for checkpoints. Removing cluster management does not remove permissions and artifact management.

This is most useful when the training algorithm itself is changing. If an existing recipe already solves the problem, a custom loop adds code to maintain without necessarily adding value.

For teams experimenting with agent behavior, the next comparison should be cost per useful improvement on held-out tasks. A cheaper training operation is only helpful when the reward and evaluation can distinguish better behavior from better performance on the training exercise.

I'm Paras. I work at the intersection of AI, engineering and distribution.

If you're building an AI product, I help turn technical launches and product thinking into content developers and founders can understand and share.

[Book a 30-minute call](https://cal.com/paras-madan/30min)