Write your own training loop — without managing GPUs — with Amazon SageMaker training session APIs

This article introduces Amazon SageMaker AI training session APIs, which allow developers to write custom training loops in Python while SageMaker handles GPU infrastructure, distributed execution, and fault recovery.

• Write custom training loops with operations like forward_backward(), optim_step(), and sample() without managing GPU infrastructure
• Supports supervised fine-tuning, reinforcement learning (GRPO, PPO), and custom objectives with algorithmic freedom
• Billing per token processed eliminates idle GPU costs; reward computation and debugging run off-clock
• Three SDK clients (ServiceClient, TrainingClient, SamplingClient) provide a hierarchical programming model
• Minimal GRPO loop on GSM8K math problems improved accuracy from 56% to 75% in 10 training steps
• Code RL recipe demonstrates multi-turn tool use with sandboxed grading on competitive programming problems
• Checkpoints stored as model packages; supports resume from saved state with weights and optimizer

Training session APIs bridge the gap between fully managed services and self-managed infrastructure, enabling researchers and ML engineers to implement custom training algorithms while offloading operational complexity to AWS.