Developers Can Now Write Their Own AI Training Code on SageMaker

AWS has opened SageMaker training sessions to custom AI training loops, giving developers algorithmic control while the service manages GPUs and failures.

AI News

AWS is expanding Amazon SageMaker AI with serverless training sessions that let developers write and run their own post-training loops without provisioning or operating GPU infrastructure. The feature is aimed at teams that need custom loss functions, reward signals, preference optimization, or multi-turn reinforcement learning rather than a fixed fine-tuning recipe.

The change moves SageMaker’s abstraction boundary. Developers still control the Python logic that determines how a model is sampled, scored, and updated, while AWS manages GPU allocation, distributed execution, fault recovery, and checkpoint storage. For AI builders, that creates a middle ground between highly prescriptive managed fine-tuning and operating a training cluster directly.

AWS announced the capability through two Machine Learning Blog posts. The AI Economy also reported the development, although the available wire item contains no additional technical detail or independent performance evidence.

What SageMaker training sessions change

Traditional managed customization services generally ask users to select a model, provide data, and configure a known training method. That approach reduces infrastructure work but can restrict researchers and engineering teams when a project requires a new objective or a bespoke training sequence.

SageMaker AI training sessions expose lower-level operations through the SageMaker Python SDK. AWS describes calls such as forward_backward(), optim_step(), and sample() as building blocks that developers can compose into supervised fine-tuning, reinforcement learning, preference optimization, or other procedures.

The session remains live and stateful, with the selected model loaded on managed accelerators across operations. A developer can run a Python process locally, issue operations against the session, inspect results, and continue iterating without first writing a complete cluster configuration or creating a separate distributed training job.

AWS says the service currently supports techniques including supervised fine-tuning, direct preference optimization, reinforcement learning with verifiable rewards, AI-feedback-based reinforcement learning, and multi-turn reinforcement learning. Users can also choose between Low-Rank Adaptation and full fine-tuning, then chain techniques through a continuous customization workflow.

The launch materials list support for models including gpt-oss-20b, gemma-4-31B-it, Qwen3.6-27B, and Qwen3.8-27B. Availability varies by model and AWS Region, according to the company’s documentation.

How the developer-controlled loop works

In AWS’s example, a reinforcement learning loop begins by sending prompts to the managed sampling service. The developer’s code evaluates the generated responses using a verifier, reward model, test harness, or another environment. It then sends the scores to the training service, where gradient calculation and the optimizer step occur on AWS-managed GPUs.

The code can refresh the sampler with updated weights and repeat the cycle. This structure is relevant to agent developers because the reward may depend on several tool calls or interactions rather than a single response. An agent that books travel, resolves a support case, or performs a research task may only be judged after multiple turns.

AWS illustrates the approach with Group Relative Policy Optimization, in which several responses to the same prompt are generated, compared using a reward function, and used to update the model. The company also describes cookbook examples involving GSM8K math problems and competitive programming, where generated code is evaluated against unit tests.

For supervised fine-tuning, the same primitives can be used without the sampling and reward stages. A developer can apply a cross-entropy loss to labeled examples. AWS says the model remains on the managed GPUs while the training session is active, while loop control, data handling, reward logic, and other application code can remain in the developer’s environment.

Evidence, pricing, and limits of the claims

The core product details come from AWS’s own technical announcements, so they establish what the vendor says the service does, not independent evidence that it improves model quality or lowers total costs in every workload.

AWS says customers are billed per token processed rather than for reserved GPU capacity. The meter pauses while a reward function scores a batch in the customer’s environment, according to the company. AWS presents this as a potential advantage for workloads with slow evaluation, complex tool use, or human review because developers are not paying for idle accelerators during those stages.

That pricing model could make experiments easier to compare, but actual economics will depend on token volume, generation length, reward computation, storage, data transfer, and the number of iterations required to reach a useful model. AWS does not provide a general cost comparison with self-managed GPUs in the supplied material.

The company also says sessions support checkpointing and resumption, tenant isolation, and failure recovery. Data and checkpoints remain in customer-controlled Amazon S3 locations, with optional AWS Key Management Service and virtual private cloud configuration. AWS states that customer data and customized weights are not used to train another model, but this is a vendor policy claim rather than an independently audited result in the source material.

AWS includes customer comments from ImagineSoftware and another unnamed customer example involving gpt-oss-20b and an Amazon Neptune knowledge graph. Those accounts describe fine-grained control and reliable checkpoint resumption, but they are customer evidence supplied by AWS, not a neutral benchmark or broad adoption study. The available AI Economy item provides no further corroboration.

Why the launch matters to AI builders

The most important audience is not every fine-tuning user. Teams using standard recipes may find a conventional managed job simpler. Training sessions are more relevant when the algorithm itself is the product: an alignment researcher testing a new advantage estimator, an agent team shaping rewards across tool calls, or an ML engineer migrating code from a low-level training framework.

The feature could reduce the operational tax associated with those experiments. Teams no longer need to decide how to place GPUs, synchronize gradients, recover failed workers, or maintain a separate inference service for every rollout. That may let smaller groups test more ambitious post-training methods while retaining control over the data curriculum and evaluation logic.

There is a tradeoff. Moving infrastructure out of the developer’s code does not remove the difficulty of designing a reliable reward function, preventing reward hacking, selecting evaluation data, or determining whether an improvement generalizes beyond a benchmark. It may also create new dependencies on SageMaker’s supported models, SDK operations, regions, and pricing model.

For enterprises, the governance details are significant. AWS says teams can use IAM execution roles, S3 storage, KMS keys, VPC settings, pre-deployment evaluation for accuracy, bias, and safety, and post-deployment monitoring through Amazon CloudWatch. Those controls may help organizations fit custom training into existing processes, but buyers will still need to validate data residency, access boundaries, auditability, and operational support for their own regions and workloads.

The competitive implication is equally specific. Cloud providers are increasingly trying to offer both ease of use and access to the training loop. SageMaker’s positioning is not simply “managed fine-tuning”; it is managed execution for custom algorithms. That places it between fully managed model customization and services such as Amazon SageMaker Training Jobs or Amazon SageMaker HyperPod, where customers assume more infrastructure responsibility.

What to watch next

The first signal will be broader model and Region support. The launch names four supported base models, but the usefulness of the service for production teams will depend on how quickly that list expands and whether customers can bring the open-weight models they already use.

Pricing transparency will also matter. AWS describes token-based billing, but detailed examples across supervised fine-tuning, long agent rollouts, and reward functions with substantial local computation would help buyers compare the service with reserved or on-demand GPU infrastructure.

Builders should watch for independent evaluations of training quality, convergence, failure recovery, and end-to-end cost. AWS’s examples demonstrate the programming model, but they do not establish that training sessions produce better models than existing frameworks.

Finally, the practical test will be whether teams can move from experimentation to deployment without losing observability and safety controls. Support for custom loops is useful only if the resulting models can be evaluated, governed, versioned, and monitored as reliably as models trained through more conventional SageMaker workflows.

Creati.ai perspective

SageMaker training sessions address a real gap in the AI development stack: many teams want to change the learning procedure, but do not want to become GPU infrastructure operators. By exposing training operations while retaining managed execution, AWS is making custom post-training more accessible without pretending that algorithm design has become simple.

The launch should be judged less by the novelty of writing a Python loop and more by operational outcomes: reproducible experiments, predictable token costs, robust reward evaluation, and a credible path from custom training to governed production deployment. AWS has established the service’s capabilities through its own technical material; the market still needs independent evidence on those outcomes.

Ads