AWS details SageMaker AI training sessions for custom model loops

AWS documentation describes SageMaker AI training sessions, which give developers APIs to assemble custom fine-tuning and reinforcement-learning loops while the service provisions the underlying GPU compute.

The key difference is the unit of control. A conventional SageMaker Training Job receives a complete training script. A training session stays active as the developer calls individual operations such as forward and backward passes, optimizer steps and sampling. AWS documentation says the model remains loaded on GPUs between calls, with gradient state persisting within a training step.

Session management and model work run through separate interfaces. Control-plane APIs create, inspect, list, stop and delete sessions. A runtime endpoint handles training steps and sampling requests against an active session, including an OpenAI-compatible chat-completions interface for sampling from the model being trained.

AWS says the serverless service handles GPU provisioning, distributed execution, gradient synchronization and tenant isolation. That removes the need for customers to select instances, manage clusters or plan capacity. Developers retain control over the operation sequence and loss function. Documented loss options include cross-entropy, importance sampling, PPO and CISPO, alongside a custom-loss path. AWS also lists supervised fine-tuning, PPO, GRPO and multi-turn reinforcement learning among the workflows teams can assemble.

Checkpointing remains part of the developer’s loop. The save_state operation stores sharded weights and optimizer state; save_weights exports inference-ready weights. AWS advises users to save state periodically because a session cannot be reattached by name. After a fault, the documented recovery path starts from the latest checkpoint. Restoring with load_state loads weights but starts the optimizer fresh, while load_state_with_optimizer restores both.

The CreateTrainingSession API currently provisions compute for foundation-model fine-tuning with Low-Rank Adaptation. Customers provide the base model, LoRA rank and alpha, an output location, an execution role and a session name. The API also accepts optional VPC and customer-managed encryption configurations.

AWS says customers are charged for tokens processed during prefill, sampling and training. Its documentation lists a default adjustable quota of two concurrent fine-tuning sessions and warns that supported base models vary by AWS Region.