Amazon’s SageMaker Gives Developers More Control Over How They Train AI Models

Amazon Web Services (AWS) now lets machine learning engineers, applied researchers, and data scientists write the training code for their AI models while the company supplies the computing power. On Thursday, AWS made a feature called training sessions generally available in Amazon SageMaker, expanding a service that lets customers fine-tune open-weight models without managing hardware. Now, customers write the training loop and the code that grades the model’s answers, and SageMaker handles the rest, including GPU allocation, distributed execution, fault recovery, and checkpoint storage. 

In the announcement, Sumedha Swamy, a senior manager of product management overseeing several areas of SageMaker; Jason Berkowitz, head of product marketing for model training and inference; and Pooja Karadgi, a senior product manager, acknowledged the hurdles customers face. “Model customization is difficult,” they wrote. “You must choose the right technique, set up and manage GPU infrastructure, and, for agents, build environments for multi-turn reinforcement learning.”

Until now, SageMaker customers customized models by choosing from AWS’s menu of training methods. They could train a model from examples of good answers (supervised fine-tuning), learn which of two answers people prefer (direct preference optimization), or improve through several forms of reinforcement learning, where the model gets scored on its attempts and adjusts. Customers can also choose to update a small add-on layer of the model through low-rank adaptation or the entire model through full fine-tuning, and they can chain techniques together. Training sessions let them go beyond that menu by writing their own training process, including how the model’s answers are graded.

“We heard from customers who wanted to define their own loss functions, shape reward signals, or RL stages in one loop,” Swamy, Berkowitz, and Karadgi said. “Training sessions give you that flexibility.”

But what is a training loop? It’s a repeating cycle behind every training method, in which a model tries something, gets measured against a standard, adjusts, and tries again. Whoever writes the loop dictates how the model learns. They choose what the model practices on, what counts as a good answer, and how strongly to reward or penalize each result. Previously, SageMaker provided a standard loop for customers. Now, customers can write their own, setting the grading standards and the training routine themselves.

According to AWS, the training loop can get complicated quickly: “You may train the model on one set of infrastructure, run inference on another to generate responses, score those responses, and feed the results into the next update. Each stage needs different compute, and connecting them across separate tools and clusters adds work.”

SageMaker training sessions eliminate the heavy infrastructure work, letting customers focus on ensuring the model is trained to their specifications. Karini AI, an agentic AI foundation platform, said it used this tool to fine-tune gpt-oss-20b for agentic, multi-turn tool calling against its Amazon Neptune knowledge graph. “Setup just worked, with no manual tuning and no out-of-memory wrangling,” chief executive Nitin Wagh said in AWS’s announcement. “Checkpoint-and-resume was a standout: resuming from an exact step was completely reliable. It gave our developers the flexibility they wanted while taking the operational burden off their plate.”

Here’s how training sessions work: The customer’s code runs the training loop and the grading in its own environment. Then, it connects the AI agent to the tool or simulator it needs to practice with. AWS SageMaker handles the GPU-intensive work, using one component to generate the model’s answers and another to calculate and apply updates. Customers pay for the amount of text the model processes, and the “meter” stops while their code grades the responses.

To help customers get started, AWS released what it calls a cookbook on GitHub with examples and use cases, along with AI agent skills for model customization. Developers describe their use case, and the agent handles data preparation, technique selection, training, evaluation, and deployment. The agent works inside SageMaker Studio and coding tools such as Kiro, Claude Code, Codex, and GitHub Copilot. 

AWS stated that every training session uses an AWS Identity and Access Management execution role that the customer creates. Customer data and saved versions of the model stay in the customer’s own Amazon S3 storage in the AWS Region they choose. The company added that it doesn’t use customer data or customized weights to train another model.

Customers can start using training sessions in SageMaker today, and at launch, it supports four open models: gpt-oss-20b, gemma-4-31B-it, Qwen3.6-27B, and Qwen3.8-27B. Training sessions can be used either in their preferred coding environment—just download the SageMaker Python SDK and clone the cookbook—or through SageMaker Studio.