DEV Community

Mikhail Savchenko
Mikhail Savchenko

Posted on Originally published at inite.ai

AWS SageMaker adds training session APIs for custom LLM loops

AWS has expanded serverless model customization in Amazon SageMaker AI with training session APIs, according to a post on the AWS Machine Learning Blog. The aim is to remove a tradeoff in LLM fine-tuning: fully managed services are simple but fix the recipe, while less managed options give control at the price of running infrastructure.

What the APIs do

Operations such as forward_backward(), optim_step() and sample() become SageMaker API calls against a live session in which the model stays loaded on the GPUs. You compose them into whatever procedure you need. The service handles runtime setup, distributed execution, gradient synchronization, fault recovery and tenant isolation. The loop itself runs on your machine and trains a LoRA adapter on a chosen base model.

The SageMaker Python SDK exposes three clients, created in sequence: a ServiceClient for platform settings such as the execution role, S3 output location and optional KMS key and VPC configuration, a TrainingClient for training and checkpointing, and a SamplingClient for generation and scoring. An active session also serves an OpenAI-compatible chat completions endpoint.

Billing and where it fits

Usage is billed per token processed. According to AWS, no GPU time is billed while your own code scores a batch, loads data or logs. Cost follows dataset and generation volume rather than wall-clock time or cluster size.

AWS places the APIs between two existing options. They give far more control over the algorithm than managed serverless customization, and far less operational work than SageMaker Training Jobs or HyperPod. A running session counts against a concurrent-session quota until it is stopped or expires.

Supported techniques and the cookbook

The same primitives cover supervised fine-tuning with a cross-entropy loss, reinforcement learning with ppo, importance_sampling or cispo losses, and custom objectives such as preference optimization or distillation. AWS also released a training sessions cookbook of ready-to-run SFT and RL recipes. Supported base models and Regions are listed in the documentation, and the source does not repeat them.

Results AWS reports

In a GRPO example on the GSM8K math benchmark, with a LoRA adapter for GPT-OSS 20B, 10 training steps on 40 problems raised held-out accuracy from 56 percent to 75 percent. In the cookbook's competitive-programming recipe, 50 steps of 128 problems with 8 responses each, over two turns, raised accuracy on 200 held-out problems from 40 percent to 50.5 percent. These are AWS's own runs, not independent results.

Who it is for

AWS aims the APIs at post-training researchers, teams that need a non-standard loss or data curriculum, builders of agentic and multi-turn RL, and ML engineers who have outgrown config-driven fine-tuning.

Top comments (0)