SageMaker AI Training Session APIs Are GA: Write Your Own Training Loop, Serverless!

Hi everyone, it's Shii-chan! SageMaker AI now lets you write your own training loop!

AWS What's New aws.amazon.com

What was announced?

AWS What's New announced training session APIs for SageMaker AI serverless model customization. They give you full control over your post-training algorithm without provisioning or managing GPU infrastructure, and they are generally available.

The story so far

Until now, managed techniques such as SFT, DPO, and other reinforcement learning methods let you pick a technique and a model, and SageMaker AI ran the training. Easy, but the training itself stayed within the technique.

What changes

You can write the training loop yourself. For example:

• Reproduce SFT or DPO on your own terms
• Implement a custom loss
• Use a new RL reward signal
• Build a data curriculum
• Implement an algorithm straight from the latest research

Dive Deep

Operations such as forward_backward(), optim_step(), and sample() become API calls against a live, stateful session. You call them from your own Python code to drive reinforcement learning, supervised fine-tuning, or custom training objectives. Checkpointing and resume are built in.

SageMaker AI handles:

• GPU allocation
• Model and LoRA loading
• Distributed execution
• Fault recovery
• Checkpoint storage
• Security

Billing is per token processed, so local computation like reward scoring isn't billed as GPU time.

AWS also released the SageMaker training sessions cookbook, a set of ready-to-run SFT and RL recipes you can use as blueprints for your own loops.

Wrap-up

• Training session APIs are GA
• Call forward_backward(), optim_step(), and sample() to build your own loop
• GPUs, distributed execution, fault recovery, and checkpoints are handled for you
• Per-token billing; local computation is not billed
• A cookbook of SFT and RL recipes is available

This one is for ML engineers who want to research and prototype post-training methods!