GRPO Training

Group Relative Policy Optimization — a reinforcement learning method for training language models with verifiable reward functions.

Overview

Group Relative Policy Optimization (GRPO) is a reinforcement learning method that improves language models by generating multiple completions per prompt, scoring them with reward functions, and using the relative ranking within each group to compute advantage estimates. Unlike DPO, which requires pre-collected preference pairs, GRPO generates its own training data online and can work with any programmatic reward signal (math correctness, format compliance, code execution results, etc.).

Use GRPO when you have a task with a verifiable reward signal and want the model to discover solution strategies on its own. Use DPO when you already have human preference data. Use SFT when you have gold-standard completions to imitate directly.

Axolotl’s GRPO implementation builds on TRL. Synchronous GRPO runs on TRL’s GRPOTrainer; async GRPO (async_prefetch: true) runs on TRL’s AsyncGRPOTrainer, extended with zero-signal filtering, parallel reward workers, reward weighting, and a cap on the policy/sampler ratio. Both scale across GPUs with FSDP2; synchronous GRPO also supports DeepSpeed.

Architecture

GRPO training uses a two-process architecture: a vLLM server for fast generation and a trainer process for scoring and gradient updates.

Terminal 1 (GPU 0)                    Terminal 2 (GPU 1)
┌──────────────────────┐              ┌──────────────────────────────────┐
│  vLLM Server         │              │  Trainer                         │
│                      │   HTTP       │                                  │
│  Serves base model   │◄────────────►│  Send prompts to vLLM            │
│  + LoRA adapter      │  /v1/*       │  Score completions (rewards)     │
│                      │  LoRA load   │  Calculate advantages            │
│  Punica kernels for  │              │  Compute policy log-probs        │
│  LoRA inference      │              │  PPO-clip gradient update        │
│                      │              │  Sync weights to vLLM            │
└──────────────────────┘              └──────────────────────────────────┘

Data flow for each synchronous training step:

The trainer sends prompts to vLLM, which generates num_generations completions per prompt.
Completions are scored with your reward functions.
Advantages are computed within each prompt group (group-relative normalization).
Policy log-probabilities are computed by running a forward pass on the training model.
The PPO-clip loss is computed and gradients are applied.
Periodically, weights are synced back to vLLM so future generations reflect the updated policy.

Async GRPO moves steps 1-3 into a separate process that runs continuously, so the trainer never waits for a full batch to finish generating.

Quick Start

A GRPO training run requires three components: a YAML config, a reward module (Python file), and a running vLLM server.

1. Write a reward module

Create a file called rewards.py in your working directory:

# rewards.py
import re


def accuracy_reward(completions, answer, **kwargs) -> list[float]:
    """Check if the completion contains the correct numerical answer."""
    rewards = []
    for completion, correct in zip(completions, answer):
        text = completion[0]["content"]
        # Extract the last number from the completion
        numbers = re.findall(r"-?\d+(?:\.\d+)?", text)
        predicted = numbers[-1] if numbers else ""
        rewards.append(1.0 if predicted == str(correct) else 0.0)
    return rewards


def format_reward(completions, **kwargs) -> list[float]:
    """Reward completions that use a structured thinking format."""
    rewards = []
    for completion in completions:
        text = completion[0]["content"]
        has_think = "<think>" in text and "</think>" in text
        has_answer = "<answer>" in text and "</answer>" in text
        rewards.append(1.0 if has_think and has_answer else 0.0)
    return rewards


def prompt_transform(cfg, *args, **kwargs):
    """Convert GSM8K dataset rows into chat prompts."""
    def transform_fn(example, tokenizer=None):
        label = example["answer"].split("####")[-1].strip().replace(",", "")
        return {
            "prompt": [
                {"role": "system", "content": "Solve the math problem. Show your reasoning in <think> tags and your final numerical answer in <answer> tags."},
                {"role": "user", "content": example["question"]},
            ],
            "answer": label,
        }
    return transform_fn, {"remove_columns": ["question"]}

2. Write the config

Create config.yaml:

base_model: Qwen/Qwen2.5-1.5B-Instruct

rl: grpo
chat_template: tokenizer_default

vllm:
  host: 0.0.0.0
  port: 8000
  gpu_memory_utilization: 0.85
  dtype: auto
  max_model_len: 2048

adapter: lora
lora_r: 32
lora_alpha: 64
lora_target_linear: true

trl:
  use_vllm: true
  async_prefetch: true
  loss_type: cispo            # see Loss below; measured against dapo on LoRA runs
  vllm_server_host: 0.0.0.0
  vllm_server_port: 8000
  vllm_server_timeout: 300
  vllm_lora_sync: true
  num_generations: 8
  max_completion_length: 512
  temperature: 0.7
  reward_funcs:
    - rewards.accuracy_reward
    - rewards.format_reward
  reward_weights:
    - 1.0
    - 0.5

datasets:
  - path: openai/gsm8k
    name: main
    type: rewards.prompt_transform
    split: train

skip_prepare_dataset: true
val_set_size: 0.0
sequence_len: 512
micro_batch_size: 2
gradient_accumulation_steps: 4
max_steps: 200
learning_rate: 5.0e-6
optimizer: adamw_torch_fused
lr_scheduler: cosine
warmup_steps: 10

bf16: true
attn_implementation: flash_attention_2
activation_checkpointing: true

special_tokens:
  pad_token: "<|endoftext|>"

output_dir: ./grpo-output
logging_steps: 1

3. Start vLLM and train

# Terminal 1: Start vLLM server on GPU 0
CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve config.yaml

# Wait 30-90 seconds for model loading and CUDA graph capture

# Terminal 2: Train on GPU 1
CUDA_VISIBLE_DEVICES=1 axolotl train config.yaml

Tip

Use tmux or separate terminal sessions to manage the two processes. The vLLM server must remain running for the entire training duration.

Custom Reward Functions

Function signature

TRL calls reward functions with this signature:

def my_reward(completions, **kwargs) -> list[float]:

completions is a list of single-element lists, where each element is a dict {"role": "assistant", "content": "..."}. So completions[i][0]["content"] gives you the text of the i-th completion.
**kwargs contains all dataset columns that were not removed by the dataset transform. This is how you pass ground truth answers, metadata, or any other information to your reward function.
Return a list[float] with the same length as completions. You may return None for individual elements to exclude them from aggregation.

Example: accuracy reward with answer extraction

def accuracy_reward(completions, answer, **kwargs) -> list[float]:
    rewards = []
    for completion, correct_answer in zip(completions, answer):
        text = completion[0]["content"]
        # Extract answer from <answer>...</answer> tags
        match = re.search(r"<answer>(.*?)</answer>", text, re.DOTALL)
        predicted = match.group(1).strip() if match else ""
        rewards.append(1.0 if predicted == str(correct_answer) else 0.0)
    return rewards

Example: length penalty

def length_penalty(completions, **kwargs) -> list[float]:
    """Penalize very short or very long completions."""
    rewards = []
    for completion in completions:
        length = len(completion[0]["content"])
        if length < 50:
            rewards.append(-0.5)
        elif length > 2000:
            rewards.append(-0.2)
        else:
            rewards.append(0.0)
    return rewards

Multiple rewards and weighting

You can combine multiple reward functions with different weights:

trl:
  reward_funcs:
    - rewards.accuracy_reward
    - rewards.format_reward
    - rewards.length_penalty
  reward_weights:
    - 1.0    # accuracy is most important
    - 0.5    # format compliance
    - 0.1    # mild length preference

Rewards are combined by the multi_objective_aggregation strategy:

sum_then_normalize (default): weights and sums all rewards first, then normalizes across the group.
normalize_then_sum (GDPO): normalizes each reward independently, then sums. This prevents one reward from dominating and is recommended when using multiple reward functions with different scales.

trl:
  multi_objective_aggregation: normalize_then_sum

Dataset transforms

The dataset transform converts raw HuggingFace dataset rows into chat-format prompts:

def prompt_transform(cfg, *args, **kwargs):
    def map_fn(example, tokenizer=None):
        return {
            "prompt": [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": example["question"]},
            ],
            # Keep 'answer' column for the reward function
            "answer": example["answer"],
        }
    # Remove columns consumed by the transform; keep columns needed by rewards
    return map_fn, {"remove_columns": ["question"]}

The transform returns a tuple of (map_function, kwargs_dict). The remove_columns in the kwargs dict removes columns that are no longer needed. Columns that your reward functions reference via **kwargs (like answer) must not be removed.

Warning

The reward module must be importable from the directory where you run axolotl train. If your reward file is rewards.py, the import path is rewards.accuracy_reward. If it is inside a package my_rewards/scoring.py, use my_rewards.scoring.accuracy_reward.

Reward models (neural network rewards)

Instead of a Python function, you can pass a HuggingFace model path as a reward function. TRL will load it as a reward model and use its scalar output as the reward (synchronous GRPO only; async GRPO scores in a process without a GPU):

trl:
  reward_funcs:
    - OpenAssistant/reward-model-deberta-v3-large-v2
    - rewards.format_reward
  reward_weights:
    - 1.0
    - 0.3

Using math_verify

The math_verify library provides robust mathematical answer verification but uses signal.alarm() internally, which only works in the main thread. If you use math_verify in a reward function, set reward_num_workers to use subprocess workers:

trl:
  reward_num_workers: 4

Each worker runs in its own subprocess with its own main thread, so signal.alarm() works correctly.

vLLM Setup

GRPO generates with vLLM, either as a separate server (below) or, on one GPU, in colocate mode. For a complete guide on server modes, LoRA sync, weight synchronization, and restart procedures, see vLLM Serving.

The minimal setup:

vllm:
  host: 0.0.0.0
  port: 8000
  gpu_memory_utilization: 0.85

trl:
  use_vllm: true
  vllm_lora_sync: true         # Recommended with LoRA — faster sync, no NCCL contention
  weight_sync_steps: 5        # Sync weights every 5 steps

CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve config.yaml   # GPU 0: vLLM
CUDA_VISIBLE_DEVICES=1 axolotl train config.yaml         # GPU 1: training

Warning

vLLM must be restarted between experiments — stale weight syncs corrupt server state. See Restart Requirements.

LoRA weight sync precision

Without vllm_lora_sync, the adapter reaches vLLM merged into BF16 weights. At GRPO learning rates the BF16 cast drops most of the LoRA update: vLLM generates from a policy that kept about 24% of the update at learning_rate: 1e-5 and about 2.5% at 1e-6 (Qwen3-1.7B-Base, rank-16 LoRA on all linear layers, step 100). The trainer logs lora_merge/retained at the first sync and every trl.vllm_merge_check_interval steps, and warns when less than half the update survives. Keep the update with one of these options:

trl.vllm_lora_sync: true with async GRPO (async_prefetch), which runs in server mode. vLLM then applies the adapter unmerged.
lora_merge_aware: bf16 when weights must be merged: colocate mode, the standard trainer, or async GRPO without vllm_lora_sync. The trainer then optimizes the merged BF16 weights directly.

See Merged Sync Rounding for details and restrictions.

Async Training Features

Async GRPO decouples generation from training. It runs on TRL’s experimental AsyncGRPOTrainer:

A CUDA-free rollout process streams requests to the vLLM server (max_inflight_tasks at a time), scores each prompt group as soon as its num_generations completions finish, and computes the group’s advantages.
Scored samples wait in a queue tagged with the policy version that generated them. The trainer drops any sample more than max_staleness weight syncs old.
The trainer packs samples into micro-batches of at most token_budget tokens, forwarding the prompt shared by a group once (tree attention), with no padding.
The policy ratio is taken against vLLM’s own sampling log-probabilities, so there is no extra scoring forward pass on the trainer.
Weights are synced to vLLM every weight_sync_steps optimizer steps: only the adapter when the server has LoRA enabled, merged weights over NCCL otherwise.

Enabling async GRPO

trl:
  use_vllm: true
  async_prefetch: true
  vllm_lora_sync: true        # `axolotl vllm-serve` starts vLLM with LoRA enabled
  weight_sync_steps: 1       # weight sync every optimizer step
  max_staleness: 4            # drop rollouts more than 4 syncs old
  reward_num_workers: 4       # score rewards in 4 processes

Start the server with axolotl vllm-serve config.yaml as usual; with async_prefetch it reserves max_staleness + 2 LoRA slots, one per policy version a queued rollout may still name.

Requirements:

vllm_mode: server. Colocate mode is not supported.
FSDP2 or DDP. DeepSpeed and context_parallel_size > 1 are not supported.
Conversational prompts (a list of messages in the prompt column). Text-only: multimodal datasets are not supported.
Reward functions get each completion as an assistant message. TRL splits out reasoning and tool calls only for the chat templates it ships a response parser for (Qwen, Llama 3.1/3.2, Gemma 4 and a few others); with any other template, including axolotl’s chatml and llama3, the message content is the decoded completion.
No reference model: beta must be 0. The loss is loss_type: score_centering (default, as in TRL), dapo or cispo; see Loss below.
Reward functions run in a spawned process with no GPU, so they must be picklable (module-level functions, not lambdas). Reward models are not supported.
A quantized base model (QLoRA, FP8) needs vllm_lora_sync: true, since merged weights cannot be folded into quantized ones.
The trainer packs with TRL’s tree attention, which uses the kernels-community/flash-attn3 hub kernel, or kernels-community/flash-attn4 on Blackwell. Where TRL’s kernel cannot load (FA3 has builds for SM80 and SM90 only, and FA4 fails to import under nvidia-cutlass-dsl releases it was not built against), axolotl falls back to kernels-community/flash-attn2.
Models with recurrent or linear-attention layers share prompt prefixes too. Gated DeltaNet layers (Qwen3.5, Qwen3-Next, OLMo-Hybrid) follow the trie through TRL’s own Triton kernels (patch_tree_gdn); the tree-recurrence package makes the other mixers follow it: Mamba2 (Mamba2, Nemotron-H, Bamba, Granite-MoE-Hybrid) and LFM2’s short convs. Axolotl’s Kimi-Linear and Bailing models follow it through their own forwards. Mamba, Falcon-Mamba, Falcon-H1 and Qwen4-exp pack each rollout as its own sequence, seeing its boundaries through the same patches as sample_packing. Other recurrent models are rejected when the config or the model loads, as is a mixer whose forward use_kernels replaced.

Without max_steps, training stops after num_epochs passes over the prompts and max_steps becomes a loose upper bound, so use lr_scheduler: constant or set max_steps for a decaying schedule.

What axolotl adds to TRL’s trainer

Zero-signal filtering (skip_zero_advantage_batches, default true). A group whose completions all get the same reward has zero advantage everywhere: its tokens add nothing to the policy gradient but cost a full forward and backward. Async GRPO drops such samples before packing, which cut forwarded tokens per trained signal-bearing token by 20% with a quarter of the groups zero-signal, 42% with half, and 71% with three quarters (measured through TRL’s batcher and collator). As in TRL’s own loop, a group without signal is dropped in the rollout process and reaches the trainer as its id only; the trainer drops the remaining zero-advantage samples, such as completions no reward function scored. As in DAPO’s dynamic sampling, the dropped tokens also leave the loss’s token count. Dropped groups still count as consumed, so epochs and checkpoint resume stay keyed on the prompts. Watch frac_reward_zero_std (the share of groups without signal), sample/dropped_zero_adv_groups_total and sample/dropped_zero_adv_total. As in TRL, reward averages every scored completion and reward_kept those sent to the trainer.

Dropped rollouts and the end of a run. A rollout can be dropped before training for having zero advantage, being truncated under mask_truncated_completions, being more than max_staleness syncs old, or not fitting token_budget. Each drop counts its prompt group as consumed (sample/dropped_<reason>_total). An epoch-bounded run stops once num_train_epochs passes over the prompts are consumed, even when the last groups were all dropped; an optimizer step that stop interrupts is not applied. A run that trains no group during a full pass over the prompts fails with the number of groups dropped for each reason and what to change, whether num_train_epochs or max_steps bounds it.

Steps that train nothing. When rewards are scored on the trainer (EBFT), zero-advantage and truncated samples are taken out of the loss rather than dropped, so a whole optimizer step can hold no trained token. Such a step counts: global_step and the learning-rate schedule advance. It applies no update, though: Adam’s momentum and weight decay would move the weights even without a gradient, so the optimizer skips every parameter and keeps its state as it was. The weights being unchanged, the following weight sync is skipped and the policy version stays, so in-flight rollouts do not age. A group whose every sample is excluded counts toward the no-progress check like a dropped one, so a full pass over the prompts that trains nothing fails the run here too. Watch batch/empty_steps_total and perf/weight_sync_skipped_total.

Parallel reward workers (reward_num_workers). TRL scores one group at a time and runs synchronous reward functions on threads. With reward_num_workers > 1, rewards run in that many processes and as many groups score at once. On a reward costing ~10 ms of CPU per completion, 4 workers scored groups 4.3x faster and 8 workers 6.2x; on math_verify checks of short answers (~1 ms each), 4 workers gave 1.6x and more workers nothing further. Each worker runs rewards on its main thread, where signal.alarm-based timeouts (math_verify) work; on threads they must be switched off.

Reward options. reward_weights, multi_objective_aggregation (including normalize_then_sum for GDPO), scale_rewards (group or none) and mask_truncated_completions keep their GRPO meaning. Advantages are normalized per group as each group finishes, so scale_rewards: batch is not available.

Loss (loss_type). Three objectives, which differ in what happens to a token whose probability has moved since vLLM sampled it. score_centering is the default, as in TRL; EBFT and NeMo Gym multi-turn, which cannot score-center, default to dapo.

dapo is TRL’s clipped objective: the PPO surrogate with the ratio taken against the vLLM sampler’s log-probabilities, clipped at [1 - epsilon, 1 + epsilon_high]. It logs TRL’s clip_ratio/* metrics. Setting vllm_importance_sampling_correction, _mode or _cap with it is an error.
cispo (opt-in) keeps every token’s gradient. Each rollout is trained on once, so, as in synchronous GRPO, the PPO ratio is taken against the policy itself and never clips, and the policy/sampler ratio becomes a detached per-token weight: vllm_importance_sampling_mode: token_truncate (default) caps it at vllm_importance_sampling_cap (default 3.0), token_mask drops tokens beyond the cap. It logs sampling/capped_token_frac. Setting epsilon or epsilon_high with it is an error: unlike synchronous GRPO’s cispo, which caps at epsilon_high, async cispo caps at vllm_importance_sampling_cap.
score_centering (default) is TRL’s score centering with masked importance sampling, which TRL introduced for samplers running in lower precision than the trainer. Each token’s gradient is weighted by its policy/sampler ratio, and tokens whose ratio falls outside [vllm_importance_sampling_clip_min, vllm_importance_sampling_clip_max] (default 0.5 and 5.0) are dropped. From that score, its expectation under the sampler is subtracted, so that a constant advantage does not pull the policy toward the sampler. The expectation is taken over the sampler’s num_sampler_logprobs most likely tokens (default 8), which the rollout worker requests from vLLM for every generated token. The vLLM server’s --max-logprobs must allow that many. TRL reports that serializing them costs 15-20% of generation throughput, and recommends --api-server-count 8 on vLLM 0.32+. It logs sampling/masked_mean (the share of dropped tokens) and sampling/sampler_mass (how much of the sampler’s probability the top tokens cover). Setting epsilon, epsilon_high or cispo’s importance-sampling options with it is an error. It needs TRL’s own rollout loop, so EBFT and NeMo Gym multi-turn cannot use it. dapo stays the objective above: TRL’s clipped objective also drops tokens outside the importance-sampling bounds, but axolotl’s dapo does not.

top_p < 1 and min_p > 0 need top_k > 0 (e.g. top_k: 512), and with top_p < 1 or top_k > 0, a vLLM 0.28+ server must be started with --return-sampling-mask --enable-scale-out, which axolotl vllm-serve passes for such a config. The trainer then normalizes each token’s log-probability over the candidate set it was sampled from, as vLLM did, under every loss. An older vLLM cannot return those sets, so the log-probabilities stay normalized over the whole vocabulary, with a warning; so do EBFT and NeMo Gym rollouts, which call vLLM themselves. Synchronous GRPO with vLLM importance sampling (vllm_importance_sampling_correction, on by default) also needs top_k with top_p < 1 or min_p > 0.

The policy/sampler ratio grows with staleness and trainer/inference mismatch. Under dapo, a token whose probability moved past epsilon since sampling stops getting gradient, and those are the tokens being learned. That matters when the policy moves far per step, as with LoRA or higher learning rates. On Qwen2.5-1.5B LoRA (r 16, lr 2e-5, max_staleness: 4) with a GSM8K accuracy + answer-format reward, cispo reached a 0.9 format rate in 17-20 steps, like axolotl’s removed trainer (15-18), and stayed stable over 200 steps. dapo took 38-49 steps, and from about step 80 its completions ran to max_completion_length in both of two seeds: the model stopped ending them. On 128 GSM8K test prompts, the final dapo models truncated 75-78% of completions and scored 0.46-0.52 accuracy; cispo truncated 4% and scored 0.64, the removed trainer 0.63. With dapo, watch completions/clipped_ratio. Validation warns when a LoRA or QLoRA run keeps dapo with rollouts allowed to be more than 2 optimizer steps old (max_staleness × weight_sync_steps; TRL’s max_staleness defaults to 4).

Every loss logs ratio and kl (policy/sampler divergence) and sampling/ratio_min, the step’s smallest policy/sampler ratio.

Out-of-place merged sync. When the server has no LoRA support, TRL merges the adapter into the base weights for the send and unmerges it afterwards. In bf16 that round trip does not restore the weights: one sync changed 6% of a layer’s frozen weights, and 1000 syncs with a moving adapter left 58% changed and a 1.3% relative drift. Async GRPO instead streams W + ΔW computed in fp32 per parameter, so the trained model is never written to. DoRA is supported this way on a single GPU or DDP, but not under FSDP.

Adapter checkpoints in fp32. Under FSDP with mixed precision, TRL writes checkpoints in the bf16/fp16 compute dtype, and FSDP2 trains LoRA factors in that dtype unless lora_fp32_gradients: true. Async GRPO writes a LoRA adapter’s tensors in fp32 either way; a fully fine-tuned model is still written in the compute dtype, as TRL does.

Migrating from the data producer trainer

Axolotl’s own GRPO trainer, which ran behind use_data_producer and, before TRL shipped AsyncGRPOTrainer, behind async_prefetch, has been removed. trl.use_data_producer: true remains as a deprecated alias for async_prefetch: true with max_staleness: 0: rollouts of an older policy are dropped, so training stays close to on-policy (the dataloader fetches one micro-batch ahead, so a step’s first micro-batch can still be one weight sync old).

The alias is slower than the trainer it replaces: every weight sync discards the generation requests still in flight. Training Qwen2.5-1.5B with LoRA on GSM8K, it dropped about 85 samples per step and took 22.5 s per step, against 16.7 s on the old synchronous path. Unless every rollout must come from the current policy, set async_prefetch: true and keep the default max_staleness (4): the same run was 12% faster per step than the old trainer’s async mode.

Config validation rejects the trainer’s other options and names what replaces each:

Option	Replacement
`use_data_producer`	`async_prefetch: true` with the default `max_staleness`; `max_staleness: 0` keeps rollouts near-synchronous at the cost above
`prefetch_depth`	`max_staleness`: how many weight syncs a rollout may lag
`streaming_partial_batch`, `streaming_min_groups`	None needed: groups are scored as they finish
`replay_buffer_size`, `replay_recompute_logps`	None: each rollout is trained on once, and `skip_zero_advantage_batches` drops samples without signal
`reroll_start_fraction`, `reroll_max_groups`	None
`batch_flattening` (with `rl: grpo`)	`token_budget`: rollouts are packed padding-free up to a token budget
`off_policy_mask_threshold`, `use_bias_correction_kl`	Synchronous GRPO only
`vllm_importance_sampling_mode: sequence_*`	Synchronous GRPO only; async GRPO uses the token modes
`beta`, `sync_ref_model`, `loss_type` other than `dapo`, `cispo` or `score_centering`, `num_iterations > 1`	Synchronous GRPO only

skip_zero_advantage_batches and reward_num_workers apply to async GRPO. Plugins’ on_rollouts_scored hook is called once per optimizer step on rank 0, with the step’s prompts, completions, per-function rewards and advantages.

Importance Sampling and Off-Policy Correction

Synchronous GRPO recomputes log-probabilities on the trainer and weights each token by the trainer/sampler ratio, with sequence-level weights and OPSM as options. Async GRPO, by default (loss_type: score_centering), weights by that ratio and drops tokens outside vllm_importance_sampling_clip_min / _max, taking it from the training forward; it clips the ratio with loss_type: dapo and weights by it, capped, with loss_type: cispo (see Loss). Synchronous GRPO:

trl:
  vllm_importance_sampling_correction: true
  importance_sampling_level: token     # 'token' recommended (especially with Liger kernel)
  off_policy_mask_threshold: 0.5       # KL threshold — masks sequences that are too off-policy

Use token level IS. Sequence-level has numerical issues with Liger’s chunked computation. The off_policy_mask_threshold (OPSM, synchronous GRPO only) is a safety net that drops sequences where KL divergence exceeds the threshold — 0.5 is a reasonable starting point.

For detailed coverage of IS modes (token_mask, token_truncate, etc.), capping, and bias-corrected KL, see vLLM Serving — IS Correction.

Scaling

Fused LM head

TRL scores GRPO and GDPO tokens, and structured EBFT’s, with a fused LM head: per-token log-probabilities are computed in chunks from the hidden states, so the [batch, seq, vocab] logits are never built. The head bypasses the model’s own forward, so cut_cross_entropy and liger_fused_linear_cross_entropy are not applied to these trainers (a warning says so), and trl.use_liger_loss only applies Liger’s layer kernels. A LoRA adapter on lm_head is rejected; train the head with lora_modules_to_save: [lm_head].

Activation checkpointing and offloading

Both GRPO trainers, synchronous and async, take activation_checkpointing, including selective and offload.all / offload.hidden_states, as described in Activation Checkpointing & Offloading. TRL’s GRPO trainers have no activation offloading of their own, so axolotl runs their training step inside its offload context. Selective checkpointing keeps saving attention after the synchronous trainer turns checkpointing off for its scoring pass and back on. All three are off unless the config sets them; TRL’s AsyncGRPOConfig alone would turn checkpointing on.

FP8 training

FP8 quantization halves model VRAM usage with minimal impact on training quality. It does not significantly speed up computation for small models but allows larger models to fit in memory.

fp8: true
torch_compile: true

Warning

FP8 requires patching for zero-padding edge cases. The act_quant_kernel can produce NaN when input is all zeros (padding positions). If you see NaN in grad norms, check whether your padding token embedding is non-zero.

FSDP (Fully Sharded Data Parallel)

FSDP distributes model parameters across multiple GPUs for training while vLLM runs on a separate GPU:

fsdp_version: 2
fsdp_config:
  auto_wrap_policy: TRANSFORMER_BASED_WRAP
  transformer_layer_cls_to_wrap: Qwen2DecoderLayer
  reshard_after_forward: true
activation_checkpointing:
  reentrant: false

Launch with:

# GPU 0: vLLM
CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve config.yaml

# GPUs 0,1: Training (FSDP will use both visible GPUs)
CUDA_VISIBLE_DEVICES=0,1 axolotl train config.yaml

Note

Async GRPO runs no collectives off the main thread: rank 0 pulls samples from the rollout queue and the dataloader dispatches one packed row to each rank.

DeepSpeed ZeRO-3

deepspeed: deepspeed_configs/zero3_bf16.json
activation_checkpointing:
  reentrant: true   # Required -- non-reentrant causes CheckpointError with ZeRO-3

Note

DeepSpeed ZeRO-3 requires activation_checkpointing.reentrant: true. This is the opposite of the FSDP recommendation. Non-reentrant checkpointing causes tensor metadata mismatches during recomputation with ZeRO-3’s parameter partitioning.

Multi-GPU considerations

Concern	Recommendation
vLLM GPU allocation	Dedicate one or more GPUs to vLLM; do not share with trainer GPUs
Weight sync contention	Use `vllm_lora_sync: true` to avoid NCCL contention between training and vLLM
Generation bottleneck	Add vLLM nodes; see Multi-Node Serving
Async GRPO	FSDP2 or DDP; DeepSpeed is not supported
DeepSpeed + gradient checkpoint	Must use `use_reentrant: true`
OOM during scoring	Reduce `micro_batch_size` or `num_generations`. The logits tensor scales with `batch_size * vocab_size`

Monitoring and Debugging

For detailed metric ranges, failure diagnosis, and OOM debugging, see Training Stability & Debugging.

Quick health checks during GRPO training:

rewards/*/mean should be > 0.15 within 20 steps — if it stays at 0, test your reward function standalone
reward_std should be > 0 on most steps — all-zero means no learning signal
entropy in 0.05-0.5 — below 0.01 suggests mode collapse
grad_norm in 0.001-1.0 — > 10 is unstable, 0.0 is expected when zero-advantage skip fires
Async GRPO: perf/rollout_wait_s near 0 means training is the bottleneck, a large value means generation is; sample/staleness_mean shows how far behind the policy rollouts are; sample/dropped_frac is the share of samples filtered out

Tip

Pipe training output to a log file: axolotl train config.yaml 2>&1 | tee /tmp/training.log

Configuration Reference

All GRPO-specific options live under the trl: key in your config. Standard training options (learning_rate, micro_batch_size, etc.) are set at the top level as usual.

Core GRPO

Option	Type	Default	Description
`use_vllm`	bool	`false`	Enable vLLM for generation
`vllm_mode`	`"server"` or `"colocate"`	`null`	vLLM deployment mode. `colocate` hosts the engine on the training GPU (single-GPU); engine options come from the `vllm:` block, see Colocate Mode
`vllm_server_host`	str	`"0.0.0.0"`	vLLM server hostname
`vllm_server_port`	int	`8000`	vLLM server port
`vllm_server_timeout`	int	`null`	Timeout (seconds) for vLLM responses
`num_generations`	int	`null`	Completions generated per prompt
`generation_batch_size`	int	`null`	Number of unique prompts per generation step
`max_completion_length`	int	`null`	Maximum tokens per completion
`beta`	float	`null`	KL penalty coefficient
`num_iterations`	int	`null`	Iterations per batch (mu in the GRPO paper)
`epsilon`	float	`null`	PPO clipping lower bound (async GRPO: `loss_type: dapo` only)
`epsilon_high`	float	`null`	PPO clipping upper bound (async GRPO: `loss_type: dapo` only)
`loss_type`	str	`null`	Loss formulation: `grpo`, `bnpo`, `dr_grpo` and TRL’s others; async GRPO: `score_centering` (default), `dapo` or `cispo`
`scale_rewards`	bool or str	`true`	Reward std scaling: `group`, `batch`, or `none` (bool accepted: `true`->`group`, `false`->`none`)
`mask_truncated_completions`	bool	`false`	Exclude truncated completions from loss
`log_multimodal`	bool	`null`	Log multimodal content (images/videos) alongside completions (TRL >= 1.9)

Entropy regularization

Available in TRL >= 1.8. Adds an entropy bonus to encourage exploration.

Option	Type	Default	Description
`entropy_coef`	float	`null`	Static entropy regularization coefficient
`use_adaptive_entropy`	bool	`null`	Adaptively tune `entropy_coef` toward `entropy_target` (Skywork-OR1 style)
`entropy_target`	float	`null`	Target entropy when `use_adaptive_entropy` is enabled
`entropy_coef_delta`	float	`null`	Step size for adjusting `entropy_coef`
`entropy_coef_min`	float	`null`	Lower bound for adaptive `entropy_coef`
`entropy_coef_max`	float	`null`	Upper bound for adaptive `entropy_coef`

Reward functions

Option	Type	Default	Description
`reward_funcs`	list[str]	`null`	Import paths to reward functions or HF model IDs
`reward_weights`	list[float]	`null`	Relative weights for each reward function
`multi_objective_aggregation`	str	`null`	`"sum_then_normalize"` (GRPO) or `"normalize_then_sum"` (GDPO)
`rollout_func`	str	`null`	Import path to custom rollout function for OpenEnv-style tasks

Generation parameters

Option	Type	Default	Description
`temperature`	float	`null`	Sampling temperature
`top_p`	float	`null`	Nucleus sampling probability
`top_k`	int	`null`	Top-k sampling
`min_p`	float	`null`	Minimum probability threshold
`repetition_penalty`	float	`null`	Penalty for repeated tokens
`generation_kwargs`	dict	`null`	Additional vLLM SamplingParams (e.g., `stop_token_ids`); not with `async_prefetch`
`chat_template_kwargs`	dict	`null`	Chat template kwargs (e.g., `{enable_thinking: false}`)
`vllm_guided_decoding_regex`	str	`null`	Regex constraint for structured/guided decoding (mapped to TRL’s `vllm_structured_outputs_regex` on TRL >= 1.7); not with `async_prefetch`

Async pipeline

Option	Type	Default	Description
`async_prefetch`	bool	`false`	Async GRPO on TRL’s `AsyncGRPOTrainer` (also structured EBFT)
`weight_sync_steps`	int	`null` (1)	Sync weights to the rollout servers every N optimizer steps (`vllm_sync_interval` is a deprecated alias)
`vllm_lora_sync`	bool	`false`	Sync only the LoRA adapter; makes `axolotl vllm-serve` enable LoRA
`vllm_merge_check_interval`	int	`50`	Log how much of the LoRA update survives merged weight sync every N steps; `0` disables
`max_staleness`	int	`null` (4)	Async GRPO: drop rollouts more than N weight syncs old
`max_inflight_tasks`	int	`null`	Async GRPO: concurrent generation requests (default `max_staleness` x samples per step)
`token_budget`	int	`null`	Async GRPO: max tokens per rank’s micro-batch (default: vLLM `max_model_len`; <= 0 packs `micro_batch_size` samples)
`scoring_token_budget`	int	`null`	Async GRPO with rewards scored on the trainer (EBFT): max tokens per scoring forward; a rank’s scored sequences are split into as many forwards as needed (default: `token_budget`; <= 0 one forward per micro-batch)
`skip_zero_advantage_batches`	bool	`true`	Async GRPO: drop zero-advantage samples
`reward_num_workers`	int	`1`	Processes for synchronous reward functions
`queue_maxsize`	int	`null` (1024)	Async GRPO: rollout samples buffered for the trainer
`request_timeout`	int	`null` (600)	Async GRPO: seconds per generation request
`weight_sync_timeout`	int	`null` (1800)	Async GRPO: seconds per vLLM weight transfer before it raises
`heartbeat_stale_after_s`	float	`null` (300)	Async GRPO: abort when the rollout worker has not reported for this long
`fork_threshold_tokens`	int	`null` (1024)	Async GRPO multi-turn: re-tokenization drift realigned as context instead of forking a new row
`max_tool_calling_iterations`	int	`null`	Async GRPO: tool-calling turns per rollout (unlimited)
`router_aux_loss_coef`	float	`null`	Async GRPO with MoE: router load-balancing loss weight (default: the model config’s)
`use_data_producer`	bool	`false`	Deprecated: `async_prefetch: true` with `max_staleness: 0`, slower than the default staleness (migration)
`vllm_enable_sleep_mode`	bool	`null`	Offload vLLM weights when idle (colocate mode)

Importance sampling

Option	Type	Default	Description
`vllm_importance_sampling_correction`	bool	`null`	Enable IS correction for async distribution shift (async GRPO: implied by `loss_type: cispo`)
`importance_sampling_level`	`"token"` or `"sequence"`	`null`	Granularity of IS ratios. Use `token` with Liger
`vllm_importance_sampling_mode`	str	`null`	`token_mask`, `token_truncate`, `sequence_mask`, or `sequence_truncate` (async GRPO: token modes, `loss_type: cispo`)
`vllm_importance_sampling_cap`	float	`null`	Cap C for IS ratio clipping/masking (async GRPO: `loss_type: cispo`, default 3.0)
`vllm_importance_sampling_clip_min`	float	`null`	Lower clip bound for the IS ratio (TRL >= 1.6; async GRPO: `loss_type: score_centering`, default 0.5)
`vllm_importance_sampling_clip_max`	float	`null`	Upper clip bound for the IS ratio (TRL >= 1.6; async GRPO: `loss_type: score_centering`, default 5.0)
`num_sampler_logprobs`	int	`null`	Async GRPO, `loss_type: score_centering`: top log-probs requested from vLLM per generated token (default 8)
`off_policy_mask_threshold`	float	`null`	KL threshold for off-policy sequence masking (OPSM), synchronous GRPO
`use_bias_correction_kl`	bool	`null`	Apply IS correction to KL divergence term, synchronous GRPO

Reference model

Option	Type	Default	Description
`sync_ref_model`	bool	`false`	Periodically sync reference model with training model
`ref_model_mixup_alpha`	float	`0.9`	EMA coefficient for reference model sync
`ref_model_sync_steps`	int	`64`	Sync reference model every N steps

Logging

Option	Type	Default	Description
`log_completions`	bool	`false`	Log sample completions to W&B
`num_completions_to_print`	int	`null`	Number of completions to print per step
`use_liger_loss`	bool	`null`	Deprecated by TRL: GRPO scores tokens with TRL’s fused LM head, so this only applies Liger’s layer kernels