Write your own training loop — without managing GPUs — with Amazon SageMaker training session APIs
aws ·
Amazon SageMaker AI has introduced training session APIs to allow developers to write custom training loops without managing GPU infrastructure. This service abstracts away compute resources, enabling users to control algorithmic details like custom loss functions or reinforcement learning reward signals while SageMaker handles distributed execution and fault recovery. The programming model relies on three SDK clients: ServiceClient for setup, TrainingClient for operations like forward_backward and optim_step, and SamplingClient for generation. Billing is based on tokens processed rather than hourly GPU usage, which reduces costs for workloads with heavy local computation. The article demonstrates this approach through a minimal GRPO loop on the GSM8K math benchmark and a more complex Code RL recipe for competitive programming. These examples show how researchers can iterate quickly on new techniques while leveraging managed accelerators. The release includes a cookbook with ready-to-run recipes for supervised fine-tuning and various reinforcement learning methods.
aws