Amazon SageMaker Training Sessions: Write Your Own RL Loop Without Managing GPUs

Amazon has moved the line between "managed fine-tuning" and "do it yourself" down one level. On October 22, 2026, AWS announced that SageMaker AI serverless model customization now includes training sessions, generally available. You write your own post-training loop, with custom loss functions, reward signals or multi-turn reinforcement learning, and SageMaker handles GPU allocation, distributed execution, fault recovery and checkpoint storage. A companion AWS post, Write your own training loop without managing GPUs, walks through the programming model.

If you have followed the 2026 shift toward reinforcement learning for LLMs, this is AWS's answer to a problem that TRL's async GRPO work, Prime Intellect's sandboxes for agentic RL and Perplexity's JOLT have all been attacking: how do you get algorithmic freedom without running a GPU cluster?

TL;DR

What problem does this solve?

AWS frames the problem as a tradeoff. With fully managed fine-tuning you pick a base model, supply data and hyperparameters, and the service owns the rest, which works when a prepackaged recipe fits. Less managed services give you control but also infrastructure and training frameworks to run. If you need a custom loss, a data curriculum, your own preference optimization, an RL reward function, or "an alignment algorithm published last week," you used to take on the infrastructure work to gain algorithmic freedom.

Training sessions draw the abstraction boundary one level lower. In AWS's words, instead of abstracting away the algorithm, they abstract away the compute. The primitives that matter, forward_backward(), optim_step() and sample(), become API calls against a session where the model stays loaded on the GPUs across steps.

How does the programming model work?

The SageMaker Python SDK exposes three clients, created in sequence:

An active session also serves an OpenAI-compatible chat completions endpoint, so existing OpenAI client code can sample from it. There is no job config file describing the whole run up front and no cluster creation step, so you can iterate and debug interactively.

The RL loop has a recognizable shape. Per AWS, each step: (1) samples from the current model on the GPU; (2) your code scores the outputs, with a verifier, a reward-model call or a test harness; (3) forward_backward and optim_step turn scores into a gradient update; (4) save_weights_and_get_sampling_client() refreshes the sampler so the next step draws from the updated model. The GPU-bound steps run on AWS; loop control and reward logic run on your machine.

For supervised fine-tuning, drop sampling and rewards and call forward_backward with a cross_entropy loss over labeled examples. Our fine-tuning explainer covers when SFT versus other methods is the right tool.

The worked example: GRPO on GSM8K

AWS trains a LoRA adapter (rank 32, alpha 64 in the sample) for OpenAI's GPT-OSS 20B using Group Relative Policy Optimization. For each GSM8K question, about 8,500 grade-school word problems each with one numeric answer, the model generates several responses, each is scored for correctness, and responses that beat the group average are reinforced. The session opens with a call like this, adapted from the AWS post:

from sagemaker.train.training_session import ServiceClient

svc = ServiceClient(role_arn=role, region=REGION)
training_client = svc.create_lora_training_client(
    "openai/gpt-oss-20b",
    training_session_name=SESSION_NAME,
    rank=32, alpha=64,
    model_package_group_arn=MODEL_PACKAGE_GROUP_ARN,
    s3_output_path=OUTPUT_S3_PATH,
)
tokenizer = training_client.get_tokenizer()

AWS's second workload is RL on competitive programming problems, where generated code is graded by running it against unit tests. That example shows why the split matters: running test harnesses is CPU work you may want on your own infrastructure, and with this model it is not billed as GPU time. The full example lives in the Amazon SageMaker examples repository referenced from the AWS post.

How is it billed, and does the saving hold up?

AWS says usage is billed per token processed. Two advantages follow, per the post: no idle-GPU cost, because reward computation, data loading, logging and debugging happen off the clock; and cost that scales with dataset and generation volume rather than wall-clock time or cluster size, which makes spend easier to forecast. AWS notes the savings over reserved GPUs "can be significant" for heavy local computation such as complex reward models, multi-tool agentic rollouts or human-in-the-loop evaluation.

That is a vendor statement, and the post does not publish prices or a head-to-head cost comparison. A fair test is your own: count tokens for a pilot run, check the pricing page for your Region, and compare against an hourly instance for the same job. The model favors loops with slow reward functions; a loop with a tiny reward check and constant GPU work will see less benefit.

Where does it sit among SageMaker's options?

AWS positions training sessions for post-training and alignment researchers, teams needing a non-standard loss or multi-stage pipeline, agentic and multi-turn RL builders, and ML engineers who have outgrown config-driven fine-tuning. It also says the operation-based model maps onto loops already written against low-level forward_backward, optim_step and sample APIs, which should ease porting.

What are the limits and open questions?

• The examples use LoRA. The worked example trains a LoRA adapter, though the launch post says you can choose LoRA or full fine-tuning. Check the documentation for which models support which.
• The model list is short. The launch post names gpt-oss-20b, gemma-4-31B-it, Qwen3.6-27B and Qwen3.8-27B; Regions are in the documentation.
• You own the algorithm. RL is easy to get subtly wrong. A reward function that can be gamed will be gamed; see our coverage of how Anthropic trained an alignment auditor with RL for how much care verifiers need.
• Lock-in. The API names resemble other hosted training services, but sessions run on AWS; moving a loop elsewhere means swapping the client layer.
• Small demo results. AWS's companion post reports that 10 steps on 40 GSM8K problems raised held-out accuracy from 56% to 75%, and that its code-RL recipe lifted accuracy on 200 held-out problems from 40% to 50.5%. These are AWS's own runs of small demos, meant to show the mechanics rather than prove the service beats a baseline.

How would you try it this week?

A sensible pilot looks like this, using only what the AWS posts describe:

1. Complete the prerequisites in the training session documentation: execution role, a model package group, and an S3 output path. The ServiceClient takes all three, plus an optional KMS key and VPC configuration if your security team requires them.
2. Start from the cookbook. AWS ships ready-to-run SFT and RL recipes as blueprints. Run the GSM8K GRPO example unchanged first, to confirm access and to measure tokens per step.
3. Swap in your own reward. Replace the numeric-answer check with your verifier: unit tests for code, a schema check for structured output, or a call to a judge model. Because reward time is unbilled, a slow but accurate verifier costs you wall-clock time rather than GPU spend.
4. Watch for reward hacking. Log sampled outputs next to their scores every few steps. If scores climb while outputs get worse, fix the reward before scaling up.
5. Checkpoint and evaluate off-session. Sessions save weights for the sampler and the adapter lands in your S3 location, so you can evaluate the adapter against a held-out set before deciding to keep training.

Because an active session exposes an OpenAI-compatible chat completions endpoint, you can also point an existing evaluation harness at the live model mid-run instead of writing new sampling code.

Who should try it first?

If you already have a reward function or verifier and want to experiment with GRPO on an open-weight model without renting GPUs by the hour, this is the shortest path AWS offers. If you only need standard SFT on a dataset, the managed serverless API is simpler. If you need full-parameter control or custom hardware, stay on Training Jobs or HyperPod. For the decision between prompting, RAG and tuning in the first place, our prompt engineering vs fine-tuning vs RAG guide is the place to start, and Physical Intelligence's fine-tuning API shows the same "you bring the data, we run the GPUs" pattern in robotics. The wider reinforcement learning debate is in our piece on David Silver's Ineffable.

Related reading

• What is fine-tuning an LLM? Complete guide
• TRL async GRPO: 26x claim vs 3x measured
• Prime sandboxes for agentic RL
• Perplexity JOLT versus GRPO
• Anthropic's RL-trained alignment auditor
• Prompt engineering vs fine-tuning vs RAG

Details are accurate as of October 22, 2026; supported models, Regions and pricing may change.