Amazon Web Services (AWS) has announced the release of training session APIs for Amazon SageMaker AI, a development designed to offer machine learning engineers more control over model training loops without the need to manage underlying GPU infrastructure. The new service seeks to bridge the gap between fully managed, configuration-based training and the operational overhead associated with self-managed cluster environments.
Customizing the Training Procedure
Previously, developers choosing fully managed services were often restricted to standard, pre-packaged training recipes. While less managed options provided greater algorithmic freedom, they required users to handle complex tasks such as GPU provisioning, fault recovery, and distributed execution. The new training session APIs shift the abstraction boundary; while the service continues to manage compute resources, developers can now write their own Python-based training loops.
By utilizing specific API calls such as forward_backward(), optim_step(), and sample(), users can define custom procedures for supervised fine-tuning (SFT), reinforcement learning (RL) reward functions, or novel alignment algorithms. These operations target a live, stateful session, allowing the model to remain loaded on GPUs across multiple training steps.
Programming Model and Efficiency
The system utilizes a hierarchy of three SDK clients: a ServiceClient to configure platform-level settings, a TrainingClient to handle gradient and checkpointing operations, and a SamplingClient for generation and evaluation tasks. Because the training loop runs on a local machine while GPU-bound operations are processed by SageMaker, developers only incur costs for the compute consumed during GPU operations. This model eliminates idle GPU costs and allows for easier cost forecasting compared to traditional hourly billing models.
Practical Implementation
To support adoption, AWS has released the “SageMaker training sessions cookbook,” which provides blueprints for reinforcement learning and fine-tuning workloads. Examples include using Group Relative Policy Optimization (GRPO) on grade-school math benchmarks and multi-turn reinforcement learning on competitive programming tasks. The latter involves integrating local sandboxes to grade model-generated code against unit tests, demonstrating how users can maintain complex local logic while offloading compute-intensive model operations to the cloud.
This release targets post-training and alignment researchers, as well as engineering teams requiring bespoke loss functions or multi-stage pipelines that exceed the capabilities of standard configuration-driven recipes.
Source: aws.amazon.com
Original article: https://aws.amazon.com/blogs/machine-learning/write-your-own-training-loop-without-managing-gpus-with-amazon-sagemaker-training-sessions/
