AWS adds training session APIs to SageMaker for custom agent training loops

AWS has added training session APIs to serverless model customization in Amazon SageMaker AI, so developers can write the training loop in Python while the service runs the GPU work. Operations such as forward_backward, optim_step and sample are SageMaker API calls against a live session that keeps the model loaded across steps. The service handles distributed execution, gradient synchronization, fault recovery and tenant isolation.

Billing is per token processed rather than by the hour, so time spent on local reward computation, data loading or debugging does not incur GPU charges. AWS also released the SageMaker training sessions cookbook with ready-to-run recipes. In a GRPO run on GSM8K math problems using a LoRA adapter for GPT-OSS 20B, AWS reports held-out accuracy rose from 56 percent to 75 percent over 10 steps on 40 problems.

A Code RL recipe trains a model on competitive programming problems, running generated Python against unit tests in a SandboxFusion container and supporting a two-turn tool call. In AWS's run, 50 steps raised accuracy on 200 held-out problems from 40 percent to 50.5 percent.

Developers who need custom losses, reward functions or multi-turn rollouts can run that logic in their own code without provisioning GPU clusters. Reward scoring and sandboxed test runs happen on your machine and are not billed as GPU time, which lowers cost for workloads with slow graders. Check the documentation for supported base models, Regions and concurrent-session quotas, since a running session counts against the quota until it is stopped.