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

Amazon SageMaker AI is expanding its serverless model customization with training session APIs, which let developers drive a live, GPU-backed session from their own Python process instead of handing the whole training job to a managed service.

• The APIs abstract away the compute rather than the algorithm, so operations like forward_backward(), optim_step() and sample() become SageMaker calls that developers compose into supervised fine-tuning, preference optimization, reinforcement learning or new alignment procedures.
• The service handles runtime setup, distributed execution, gradient synchronization, fault recovery, tenant isolation and GPU operations, while the model stays loaded on the GPUs across steps in a stateful session.
• The SageMaker Python SDK organizes the work around three clients created in sequence: a ServiceClient for platform settings and session start, a TrainingClient for training and checkpointing, and a SamplingClient for generation and scoring during rollouts and evaluation.
• An active session also serves an OpenAI-compatible chat completions endpoint, so existing OpenAI client code can sample from it.
• Amazon is releasing a training sessions cookbook with ready-to-run SFT and RL recipes, including reinforcement learning on competitive programming problems where generated code is graded against unit tests.
• Usage is billed per token processed, so no GPU time is charged while reward computation, data loading, logging or debugging run on the developer's machine.