View a markdown version of this page

Training session APIs - Amazon SageMaker AI
Who it is forKey differentiatorsSupported models, pricing, and RegionsTopics

Training session APIs

The SageMaker AI training session APIs run model customization on stateful, serverless GPU sessions where you orchestrate training operations directly. Existing high-level fine-tuning APIs constrain your training loop to predefined workflows such as supervised fine-tuning (SFT). You cannot write custom loss functions or run frontier reinforcement learning (RL) with elastic rollouts. You are also locked into the infrastructure decisions those APIs make for you.

Teams building production AI systems need to iterate on algorithms, experiment with novel training objectives, and scale RL across Regions without rewriting distributed training code. Training session APIs give you that level of control.

Note

Training session APIs require you to write Python code that drives the training loop. If you want a managed workflow that does not require programming, see Serverless model customization. If you want to submit a complete training script as a job instead, see SageMaker AI Training Jobs.

With the training session APIs, instead of submitting a complete training script, you open a session and call operations like forward_backward(), optim_step(), and sample() as API calls, and compose them into whatever training procedure your technique requires. The model stays loaded on GPUs across calls, and the service handles distributed execution, gradient synchronization, and tenant isolation.

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 reinforcement learning 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()

Use the training session APIs when you need:

  • Full control over a custom training loop (SFT, RL, or novel techniques).

  • Multi-stage pipelines that chain multiple procedures in one session.

  • Step-by-step debugging of loss divergence or reward hacking.

  • Multi-turn RL where rollout generation and policy updates share model state.

Who it is for

  • ML researchers experimenting with novel training techniques.

  • Advanced ML teams running production RL systems.

  • Teams building custom loss functions, such as group relative policy optimization (GRPO), DRO, and DAPO.

  • Organizations needing multi-Region training and elastic rollout inference.

  • Teams that own their data and want to own their models.

Key differentiators

Capability What it means
Operation-level API You orchestrate forward_backward(), sample(), and optim_step() individually rather than submitting a monolithic training script.
Stateful sessions The model stays loaded on GPUs across calls. No cold starts between rollout generation and policy updates.
Serverless No instance selection, no cluster management, no capacity planning. The service provisions and manages GPU infrastructure.
Technique-agnostic Implement SFT, proximal policy optimization (PPO), GRPO, multi-turn RL, or any custom procedure. You choose the algorithm; the service does not limit your technique.

Supported models, pricing, and Regions

Supported models

Model Region
openai/gpt-oss-20b us-east-1, us-west-2, ap-south-1
google/gemma-4-E2B-it us-west-2
google/gemma-4-31B-it us-west-2, ap-south-1
Qwen/Qwen3.6-27B us-west-2, ap-south-1
Qwen/Qwen3.8-27B us-west-2

Context length limits

Each supported model has a maximum context length. The context length is the total number of tokens in one sequence. For sampling, this is the prompt plus the requested maximum completion length (maxTokens). For training, this is the full datum. The limit applies to both sampling and training calls. Requests that exceed the limit fail with the error The input exceeds the model's maximum context length. Please reduce the length of your prompt or the requested max tokens. Plan your prompt and response lengths so that each sequence fits within the limit. This matters most for multi-turn rollouts, where the context grows with each turn. For an example of sampling and training calls, see End-to-end examples.

The following table lists the maximum context length for each supported model and whether you can adjust it.

Model Maximum context length (tokens) Adjustable
openai/gpt-oss-20b 65,536 No
google/gemma-4-E2B-it 32,768 No
google/gemma-4-31B-it 32,768 No
Qwen/Qwen3.6-27B 65,536 No
Qwen/Qwen3.8-27B 65,536 No

Pricing

For full pricing details, see the SageMaker AI pricing page. Charges are based on three dimensions:

  • Prefill – the cost of processing the input tokens fed into the model at the start of each training step. This includes the prompt, conversation history, and any context the model receives before generating a response.

  • Sample – the cost of the tokens the model generates during the training rollout. This is where the model produces its responses, which are then evaluated and used to compute the reward signal.

  • Train – the cost of the backward pass, where the model weights are updated based on the reward signal. This is the core learning step in the RL process.

Training session quotas

The following quotas apply to the training session APIs. These quotas are adjustable. To request an increase, open the Service Quotas console. For more information, see Requesting a quota increase.

Name Default Adjustable
Maximum number of concurrent fine-tuning training sessions 2 Yes
Note

Some accounts have a default quota of 0 for training sessions. If your account does, request a quota increase in the Service Quotas console.

For a list of all SageMaker AI quotas, see Supported Regions and Quotas.

Topics