# SageMaker training sessions put custom Python loops on managed GPUs

- Source: MechNext — https://mechnext.eu/en/article/53-sagemaker-training-sessions-put-custom-python-loops-on-managed-gpus
- Published: 2026-10-26
- Updated: 2026-10-26
- Category: Industry News
- Language: en (Markdown in other languages: ?lang=de, ?lang=sk, ?lang=pl, ?lang=zh, ?lang=ru)

> AWS describes stateful training APIs for SFT, reinforcement learning and custom objectives, with local Python controlling the loop and SageMaker managing GPU execution.

Amazon SageMaker AI training session APIs let developers compose training loops in Python while the service manages GPU execution, according to AWS. Operations including `forward_backward`, `optim_step` and `sample` act on a live, stateful session that keeps the model loaded across steps. The service handles runtime setup, distributed execution, gradient synchronization, fault recovery and tenant isolation.

### Local control of the training loop

A local Python process can drive a session with a chosen base model and LoRA adapter loaded on managed accelerators. The programming model requires neither cluster creation nor a configuration file specifying the entire job upfront.

The SDK uses three clients in sequence. `ServiceClient` holds platform settings and starts a session; `TrainingClient` supplies training and checkpoint operations; `SamplingClient` provides `sample` and `compute_logprobs`. Active sessions also expose an OpenAI-compatible chat completions endpoint.

In the illustrated reinforcement learning loop, managed GPUs generate responses and local Python scores them. Calls to `forward_backward` and `optim_step` update the model, followed by `save_weights_and_get_sampling_client` to refresh sampling. Supervised fine-tuning omits sampling and reward calculation and applies `cross_entropy` loss to labeled examples.

### Execution and billing

GPU-bound calls immediately return `APIFuture` objects, with requests queued and submitted in the background. Calling `result()` waits for completion. Sampling requests can run concurrently, while training operations execute in submission order. The example training datum includes prompt and response tokens, sampled token log probabilities, an advantage and a mask excluding prompt tokens from the loss.

AWS says billing is per token processed during GPU-bound operations, with no GPU time billed during local reward computation, data loading, logging or debugging. Supported objectives include SFT, RL losses such as `ppo`, `importance_sampling` and `cispo`, and custom preference optimization or distillation. AWS positions sessions between the algorithm control of Serverless Model Customization and the operational ownership of Training Jobs or HyperPod.

### Checkpoints and reusable recipes

Every checkpoint save registers a model package version. The example saves both a merged model for deployment and weights plus optimizer state for resumption through `create_training_client_from_state_with_optimizer`. Merged checkpoints can deploy to real-time endpoints. Sessions count toward the concurrent-session quota until stopped or expired.

The accompanying cookbook handles batching, chat templates, advantage estimation, overlapping sampling and training, validation, checkpoint resumption, per-step metrics and run configurations. Recipes include chat SFT with Tulu 3 and UltraChat, full-scale GSM8K RL and KTO preference alignment using binary feedback.

### Reported experiments

The GSM8K example trains a GPT-OSS 20B LoRA adapter with GRPO, reinforcing responses that beat their group's average correctness score. AWS reports that 10 training steps on 40 problems raised held-out accuracy from 56% to 75%.

The `code_rl` recipe uses DeepCoder competitive programming problems and tests generated Python solutions in a local SandboxFusion container. It awards 1 if all tests pass and 0 otherwise, with a 0.1 penalty for missing a code block. A two-turn setup lets the model call `check_solution`, inspect results and revise its answer.

This recipe overlaps generation with training, bounds response staleness, uses `cispo` to correct policy differences and computes gradients in sub-batches. AWS reports a rank-32 GPT-OSS 20B LoRA run of 50 steps, using 128 problems, eight responses per problem, two turns and a three-step staleness bound. Training reward rose from 0.46 to 0.60; accuracy on 200 held-out problems rose from 40% to 50.5%. Metrics can optionally go to SageMaker managed MLflow.

These are AWS-reported results for specific runs, not a guarantee for other workloads. The supplied material does not enumerate supported base models, AWS Regions, numerical quotas or per-token prices.