SageMaker Training Sessions let you write your own RL loop without provisioning GPUs

October 22, 2026. AWS makes the SageMaker AI Training Session APIs generally available, letting you customize a model with your own training loop. October 22, 2026. The starting catalog is four base models: GPT-OSS 20B, Gemma 4 31B, Qwen3.6 27B, and Qwen3.8 27B. October 22, 2026. Billing shifts to tokens — prefill, sample, train — rather than provisioned instances. The core shift: for the first time, custom post-training steps out of GPU cluster management.

For an ML or platform team, this is not cosmetic. It removes a ceiling: until now, anything beyond SFT or DPO — a GRPO, a PPO, a custom loss function — sent you back to your own clusters or to third-party post-training platforms.

A ceiling of the old model

SageMaker’s high-level fine-tuning APIs constrained the loop to predefined workflows, typically supervised fine-tuning (SFT). You could not write a custom loss function, run frontier RL with elastic rollouts, or escape the infrastructure decisions those APIs made for you. For any team iterating on new training objectives — GRPO, DRO, DAPO — or debugging a loss divergence, the only exit was to take infrastructure back into your own hands.

Training Sessions invert that relationship: instead of submitting a complete training script, you open a session and call unit operations — forward_backward(), optim_step(), sample() — composing them into whatever procedure your technique requires. The model stays loaded on the GPUs between calls, and the service handles distributed execution, gradient synchronization, and tenant isolation.

An RL loop in a few lines

The heart of the offer is this example from the documentation:

from sagemaker.train.training_session import ServiceClient, AdamParams

svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::111122223333:role/MyRole")
tc = svc.create_lora_training_client(base_model="openai/gpt-oss-20b", rank=32, alpha=64)

# Custom RL loop with on-policy sampling
sc = tc.create_sampling_client()
for step in range(num_steps):
    ops = [sc.sample(p, num_samples=4, sampling_params=params) for p in prompts]
    rollouts = [op.result() for op in ops]
    tc.forward_backward(build_batch(rollouts), loss_fn="ppo").result()
    tc.optim_step(AdamParams(learning_rate=5e-5)).result()
    sc = tc.save_weights_and_get_sampling_client()

tc.stop()

The point is that you own the algorithm — here an on-policy PPO with four samples per prompt — while AWS carries the GPUs, distribution, and fault recovery. The sampling client is recreated on every weight save so rollouts stay aligned with the current policy.

AWS lists four target use cases: full control over a custom loop (SFT, RL, or novel techniques), multi-stage pipelines that chain several procedures in one session, step-by-step debugging of loss divergence or reward hacking, and multi-turn RL where rollout generation and policy updates share the same model state.

Three doors, three levels of control

SageMaker now offers three paths to customization, and the choice matters. The first, Serverless model customization, is the managed workflow with no programming: you supply a dataset and the API applies a predefined SFT. The second, the Training Session APIs, hands you the loop — you write forward_backward(), optim_step(), and sample() — while the service carries the GPUs. The third, SageMaker Training Jobs, remains the door for complete scripts: you submit a full training job when you need control beyond what a session offers — pretraining, architecture changes, massive proprietary data.

The right selection criterion is control over the algorithm, not the hardware. As long as your technique fits a standard workflow, stay on the managed path. The moment you touch the loss function or the reward signal, Training Sessions are the equilibrium point: you keep the algorithm, AWS keeps the infrastructure. Training Jobs only matter again for what falls outside a stateful session.

Billing that follows RL compute

The move to three meters — prefill, sample, train — deserves attention because it changes the economics. In SFT, most cost comes from prefill (the inputs) and train (the backward pass). In RL, sample becomes dominant: every rollout generates tokens that get evaluated to produce the reward signal, and a multi-turn GRPO can multiply samples by orders of magnitude. That is exactly the line item teams underestimate when migrating from per-instance-hour cost. The operational advice is simple: instrument the tokens generated per rollout before launching, and set a sample ceiling per experiment.

Four models, three meters, one quota

The initial release is deliberately narrow. Only four base models: openai/gpt-oss-20b (us-east-1, us-west-2, ap-south-1), google/gemma-4-31B-it (us-west-2, ap-south-1), Qwen/Qwen3.6-27B (us-west-2, ap-south-1), and Qwen/Qwen3.8-27B (us-west-2). Customization runs through LoRA (rank, alpha), which bounds the size of the updates.

Billing changes nature: three token-based dimensions — prefill (input tokens at the start of each step), sample (tokens generated during the rollout, evaluated for the reward signal), and train (the backward pass that updates the weights). It is a cost model aligned to the real compute of RL, not to reserved instance-hours. The default quota is 2 concurrent training sessions, adjustable through Service Quotas.

Verdict

If your post-training is limited to SFT or DPO, do not migrate: the high-level APIs remain simpler, and the value of Training Sessions only materializes once you write your own loop.

If you run custom RL — GRPO, PPO, homegrown reward functions, multi-turn RL — or need to debug reward hacking, this removes the GPU-ops tax: no instance selection, no cluster management, no capacity planning. Two caveats belong in the calculation: the four-model catalog limits your base today, and per-token billing (prefill + sample + train) can surprise if your rollouts run long — calibrate the budget before launching, and treat the 2-session quota as a bottleneck to raise explicitly.

References

• AWS Docs — Training session APIs, Amazon SageMaker AI
• AWS What’s New — Amazon SageMaker AI announces training session APIs (October 22, 2026)
• ExplainX — SageMaker Training Sessions: Custom RL Loops, No GPUs