Reinforcement Fine-Tuning on SageMaker AI: GRPO with Verifiable Rewards

Supervised fine-tuning teaches a model to imitate answers. It cannot teach a model to get things right when correctness is checkable — whether the JSON parses, whether the SQL returns the expected rows, whether the unit test passes. That is what reinforcement fine-tuning (RFT) with verifiable rewards is for, and since the DeepSeek-R1 era of reasoning models it has become the default second stage of post-training for tool-using and structured-output workloads.

This tutorial shows how to run GRPO (Group Relative Policy Optimization) on Amazon SageMaker AI as an ordinary training job: how the reward function works, how to size the GPUs, what to watch while it trains, and how to know whether the run actually helped. It assumes you already have a supervised fine-tune in hand — RFT is a refinement stage, not a replacement for SFT.

When RFT is worth it (and when it is not)

Reinforcement fine-tuning earns its cost only when you can write a program that scores an answer. Use it when:

• Correctness is checkable. Schema validation, SQL execution against a fixture database, pytest on generated code, a regex/exact match on a final answer, a deterministic tool-call trace.
• The base model is close but inconsistent. RFT sharpens behaviour the model can already produce sometimes. It does not install knowledge that isn't there.
• You have 500–5,000 prompts with checkable outcomes. RFT needs far less data than SFT, but the prompts must be hard enough that the model is neither always right nor always wrong.

Skip it when the only judge is taste (use SFT or DPO on preference pairs), when your grader is an LLM you haven't validated (you will optimize its bias), or when a prompt change plus constrained decoding gets you there for free. We tell clients to try structured-output constraints first; if the gap is reasoning quality rather than format, RFT is next.

Why GRPO rather than PPO

Classic PPO needs a separate value network — a second model in GPU memory, plus its own optimizer state and its own failure modes. GRPO drops it. For each prompt it samples a group of G completions, scores them all, and uses the group's mean reward as the baseline:

advantage_i = (reward_i - mean(rewards)) / (std(rewards) + eps)

That single substitution removes roughly a third of the memory footprint and most of the tuning surface. The practical consequence for a SageMaker budget: a 7–8B GRPO run fits comfortably where a PPO run of the same size would not, and there is one fewer learning rate to get wrong.

The cost you pay instead is generation. Every step samples G completions per prompt, so a step with G=8 and 512 new tokens is ~4,000 tokens of decoding per prompt before any gradient is computed. Generation, not backprop, will dominate your wall-clock — which is why vLLM-backed generation inside the training container is not optional at any serious scale.

Step 1: Shape the dataset

One JSONL row per prompt. Keep the ground truth in the row — the reward function needs it, and you want it versioned with the data, not hidden in code:

{"prompt": "Return the order total for customer 42 as JSON with keys total and currency.", "schema": "order_total_v1", "answer": {"total": 1284.5, "currency": "USD"}}

Filter the set before you train. Sample the SFT checkpoint 8 times per prompt and keep only prompts whose pass rate is between roughly 10% and 90%. Prompts the model always gets right contribute zero advantage (every completion in the group scores the same, so the normalized advantage collapses); prompts it never gets right contribute noise. This one filtering pass is the highest-leverage thing you can do to a GRPO run and it typically removes half the corpus.

Step 2: Write the reward function

Rewards should be a small sum of independently verifiable components, each cheap and deterministic:

import json, re

def reward(completion: str, row: dict) -> float:
    score = 0.0

    # 1. Format: did it emit a single fenced JSON object?
    m = re.search(r"```json\s*(\{.*?\})\s*```", completion, re.S)
    if not m:
        return -0.5
    score += 0.2

    # 2. Parseability
    try:
        obj = json.loads(m.group(1))
    except json.JSONDecodeError:
        return score - 0.3
    score += 0.2

    # 3. Correctness against the verified answer
    want = row["answer"]
    if set(obj.keys()) == set(want.keys()):
        score += 0.2
    if obj.get("currency") == want["currency"] and \
       abs(float(obj.get("total", 0)) - want["total"]) < 0.01:
        score += 1.0

    # 4. Mild brevity pressure to discourage rambling preambles
    if len(completion) > 4000:
        score -= 0.2

    return score

Four rules that matter more than the reward shape itself:

1. Never let the grader raise. A traceback in the reward path kills the job hours in. Wrap everything, return a floor value on failure, and log the exception.
2. Keep it fast. The grader runs batch_size × G times per step. A 200 ms grader at G=8 and batch 16 adds ~26 s per step. If you must execute code or SQL, do it in a pool of pre-warmed sandboxes with a hard timeout.
3. Make format a small reward, not most of it. If format is worth more than correctness, the model learns to emit beautiful empty answers.
4. Hold out a grader test set. Feed your grader 50 known-good and 50 known-bad completions and assert it scores them correctly before the training job launches. Most "RFT didn't work" post-mortems are grader bugs.

Step 3: The training script

TRL's GRPOTrainer is the shortest correct path, and it will use vLLM for generation when you ask it to. Save this as train_grpo.py under source_dir:

import json, os
from datasets import load_dataset
from trl import GRPOConfig, GRPOTrainer
from rewards import reward  # your graded reward, imported and unit-tested

train = load_dataset(
    "json", data_files=os.environ["SM_CHANNEL_TRAIN"] + "/train.jsonl"
)["train"]

def reward_fn(completions, **kwargs):
    rows = [dict(zip(kwargs, v)) for v in zip(*kwargs.values())]
    return [reward(c, r) for c, r in zip(completions, rows)]

cfg = GRPOConfig(
    output_dir="/opt/ml/checkpoints",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,
    num_generations=8,               # G
    max_prompt_length=1024,
    max_completion_length=512,
    learning_rate=1e-6,              # an order of magnitude below SFT
    beta=0.02,                       # KL penalty toward the reference policy
    bf16=True,
    gradient_checkpointing=True,
    use_vllm=True,
    vllm_mode="colocate",
    vllm_gpu_memory_utilization=0.35,
    logging_steps=1,
    save_steps=50,
    report_to="mlflow",
)

trainer = GRPOTrainer(
    model=os.environ["SM_CHANNEL_MODEL"],  # your SFT checkpoint
    args=cfg,
    train_dataset=train,
    reward_funcs=reward_fn,
)
trainer.train()
trainer.save_model("/opt/ml/model")

Two settings deserve emphasis. learning_rate=1e-6 is deliberately tiny: RFT at SFT learning rates diverges into gibberish within a few dozen steps. And beta is your leash — it penalizes KL divergence from the frozen reference policy. Too low and the model reward-hacks; too high and nothing moves. Start at 0.02 and only touch it if the KL curve tells you to.

Step 4: Launch it on SageMaker AI

import sagemaker
from sagemaker.pytorch import PyTorch

role = sagemaker.get_execution_role()

estimator = PyTorch(
    entry_point="train_grpo.py",
    source_dir="src",                 # includes rewards.py + requirements.txt
    role=role,
    framework_version="2.6",
    py_version="py312",
    instance_type="ml.g6e.12xlarge",  # 4x L40S, 48 GB each
    instance_count=1,
    keep_alive_period_in_seconds=1800,
    checkpoint_s3_uri="s3://my-bucket/grpo/ckpt/",
    checkpoint_local_path="/opt/ml/checkpoints",
    max_run=24 * 3600,
    environment={"MLFLOW_TRACKING_URI": "<your-mlflow-arn>"},
)

estimator.fit({
    "train": "s3://my-bucket/grpo/data/",
    "model": "s3://my-bucket/sft/model/",
})

Sizing guidance for an 8B policy with LoRA adapters and G=8:

Setup	Instance	Notes
8B, LoRA, colocated vLLM	ml.g6e.12xlarge	cheapest workable single-node start
8B, full-parameter	ml.p4d.24xlarge	needs FSDP + a dedicated generation server
32B, LoRA	ml.p5.48xlarge or HyperPod	reserve capacity; expect multi-day runs
Smoke test (1B)	ml.g5.2xlarge	run this first, always

Do the 1B smoke test on 50 prompts and 20 steps before anything else. It exercises the grader, the channels, the checkpoint path, and MLflow logging for a couple of dollars. Note that colocating vLLM with training means the memory split (vllm_gpu_memory_utilization) is a real tradeoff — too generous and training OOMs, too stingy and generation crawls. 0.3–0.4 is a sane starting band.

Spot instances are tempting here and they do work, but GRPO steps are long; make sure save_steps is frequent enough that an interruption costs you minutes, not hours, and set max_wait >= max_run.

Step 5: Read the curves

Four metrics tell you almost everything. Log them every step:

Metric	Healthy	Trouble
reward (mean)	slow, noisy climb	flat = advantages collapsed or LR too low; vertical jump = reward hack
reward_std within group	stays > 0	→ 0 means every completion scores alike; re-filter prompts
kl to reference	creeps up, plateaus	spikes = raise beta; pinned at 0 = leash too tight
completion_length	roughly stable	runaway growth = length exploitation; add brevity penalty

The classic failure is a reward curve that shoots up while your held-out eval gets worse. That is reward hacking, and it is not subtle once you look at samples. Which is the real discipline of RFT: read 20 completions by hand at step 0, mid-run, and at the end. No dashboard substitutes for that.

Step 6: Prove it before you ship it

RFT changes behaviour globally, so evaluate on more than the reward you trained against:

1. Task pass rate on a held-out set the grader never saw during training.
2. Regression suite — general instruction-following and safety prompts, scored with fmeval so a fluency collapse cannot sneak past. This is where alignment tax shows up.
3. Latency and token cost. If the RFT model now "thinks" for 400 extra tokens per answer, your per-request cost moved and your endpoint sizing is stale.

Gate the artifact through the Model Registry with those three results attached to the model card, and deploy the winner behind an inference component so you can shadow-test it against the incumbent SFT model on live traffic before shifting any of it.

A realistic cost picture

An 8B LoRA GRPO run on one ml.g6e.12xlarge — 1,500 filtered prompts, G=8, ~600 steps — lands in the 10–16 hour range, so tens of instance-hours plus the smoke tests and the two or three restarts you will need after the first grader bug. Budget a week of engineering, not an afternoon: the training job is the easy part, and the reward function is where the whole project actually lives.

The teams that get value out of RFT are the ones with a checkable objective and the patience to debug a grader. If that is your situation and you want the reward design, sizing, and eval gates reviewed before you burn GPU hours on it, get in touch.