trl Documentation

Technical Documentation: huggingface/trl

ℹ️ Provenance: Multi-Tier Ingestion · CodeWiki Reference

Overview (`README.md`) Copy Section

TRL - Transformers Reinforcement Learning

A comprehensive library to post-train foundation models

🎉 What's New Copy Chapter

🚀 Training on AMD GPUs: TRL's trainers now run on AMD Instinct GPUs through PyTorch's ROCm build.

🔥 Beam: Fine-tune Beam with SFTTrainer on a single node of 8 B300s. See the LoRA example.

🔥 Muse Spark Open: Fine-tune Muse Spark Open with SFTTrainer, sharded across your GPUs at load time. See the 8-node LoRA example.

🚀 Training on AWS Trainium: TRL's trainers now run on AWS Trainium through TorchNeuron, PyTorch's native Neuron backend.

Overview Copy Chapter

TRL is a cutting-edge library designed for post-training foundation models using advanced techniques like Supervised Fine-Tuning (SFT), Group Relative Policy Optimization (GRPO), and Direct Preference Optimization (DPO). Built on top of the 🤗 Transformers ecosystem, TRL supports a variety of model architectures and modalities, and can be scaled-up across various hardware setups.

Highlights Copy Chapter

• Trainers: Various fine-tuning methods are easily accessible via trainers like SFTTrainer, GRPOTrainer, DPOTrainer, KTOTrainer and more.

• Efficient and scalable:

  • Leverages 🤗 Accelerate to scale from single GPU to multi-node clusters using methods like DDP and DeepSpeed.
  • Full integration with 🤗 PEFT enables training on large models with modest hardware via quantization and LoRA/QLoRA.
  • Integrates 🦥 Unsloth for accelerating training using optimized kernels.

• Command Line Interface (CLI): A simple interface lets you fine-tune with models without needing to write code.

Installation Copy Chapter

Python Package Copy Section

Install the library using pip:

pip install trl

From source Copy Section

If you want to use the latest features before an official release, you can install TRL from source:

pip install git+https://github.com/huggingface/trl.git

Repository Copy Section

If you want to use the examples you can clone the repository with the following command:

git clone https://github.com/huggingface/trl.git

Quick Start Copy Chapter

For more flexibility and control over training, TRL provides dedicated trainer classes to post-train language models or PEFT adapters on a custom dataset. Each trainer in TRL is a light wrapper around the 🤗 Transformers trainer and natively supports distributed training methods like DDP, DeepSpeed ZeRO, and FSDP.

`SFTTrainer` Copy Section

Here is a basic example of how to use the SFTTrainer:

from trl import SFTTrainer
from datasets import load_dataset

dataset = load_dataset("trl-lib/Capybara", split="train")

trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B",
    train_dataset=dataset,
)
trainer.train()

`GRPOTrainer` Copy Section

GRPOTrainer implements the Group Relative Policy Optimization (GRPO) algorithm that is more memory-efficient than PPO and was used to train Deepseek AI's R1.

from datasets import load_dataset
from trl import GRPOTrainer
from trl.rewards import accuracy_reward

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = GRPOTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    reward_funcs=accuracy_reward,
    train_dataset=dataset,
)
trainer.train()

[!NOTE]
For reasoning models, use the reasoning_accuracy_reward() function for better results.

`DPOTrainer` Copy Section

DPOTrainer implements the popular Direct Preference Optimization (DPO) algorithm that was used to post-train Llama 3 and many other models. Here is a basic example of how to use the DPOTrainer:

from datasets import load_dataset
from trl import DPOTrainer

dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")

trainer = DPOTrainer(
    model="Qwen/Qwen3-0.6B",
    train_dataset=dataset,
)
trainer.train()

`KTOTrainer` Copy Section

KTOTrainer implements the Kahneman-Tversky Optimization (KTO) algorithm, which aligns models from simple binary (desirable / undesirable) feedback rather than paired preferences. Here is a basic example of how to use the KTOTrainer:

from datasets import load_dataset
from trl import KTOTrainer

dataset = load_dataset("trl-lib/kto-mix-14k", split="train")

trainer = KTOTrainer(
    model="Qwen/Qwen3-0.6B",
    train_dataset=dataset,
)
trainer.train()

`RewardTrainer` Copy Section

Here is a basic example of how to use the RewardTrainer:

from trl import RewardTrainer
from datasets import load_dataset

dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")

trainer = RewardTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    train_dataset=dataset,
)
trainer.train()

Command Line Interface (CLI) Copy Chapter

You can use the TRL Command Line Interface (CLI) to quickly get started with post-training methods like Supervised Fine-Tuning (SFT) or Direct Preference Optimization (DPO):

SFT:

trl sft --model_name_or_path Qwen/Qwen2.5-0.5B \
    --dataset_name trl-lib/Capybara \
    --output_dir Qwen2.5-0.5B-SFT

DPO:

trl dpo --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
    --dataset_name argilla/Capybara-Preferences \
    --output_dir Qwen2.5-0.5B-DPO

KTO:

trl kto --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
    --dataset_name trl-lib/kto-mix-14k \
    --output_dir Qwen2.5-0.5B-KTO

Read more about CLI in the relevant documentation section or use --help for more details.

Development Copy Chapter

If you want to contribute to trl or customize it to your needs make sure to read the contribution guide and make sure you make a dev install:

git clone https://github.com/huggingface/trl.git
cd trl/
pip install -e .[dev]

Experimental Copy Chapter

A minimal incubation area is available under trl.experimental for unstable / fast-evolving features. Anything there may change or be removed in any release without notice.

Example:

from trl.experimental.new_trainer import NewTrainer

Read more in the Experimental docs.

Citation Copy Chapter

@software{vonwerra2020trl,
  title   = {{TRL: Transformers Reinforcement Learning}},
  author  = {von Werra, Leandro and Belkada, Younes and Tunstall, Lewis and Beeching, Edward and Thrush, Tristan and Lambert, Nathan and Huang, Shengyi and Rasul, Kashif and Gallouédec, Quentin},
  license = {Apache-2.0},
  url     = {https://github.com/huggingface/trl},
  year    = {2020}
}

License Copy Chapter

This repository's source code is available under the Apache-2.0 License.

Overview (`examples/README.md`) Copy Section

Examples

Each example lives in its own self-contained folder. The full index and layout conventions are documented at https://huggingface.co/docs/trl/example_overview.

Adding a new example Copy Chapter

• Name the folder after the method and what makes the example distinctive: the task (grpo_wordle), the model (sft_gpt_oss), or the technique (grpo_qlora). Load folder-local assets relative to the script (Path(__file__).parent / ...), not the working directory.
• An example tells a story. A bare single-trainer training script is not an example: stable trainers already have one as a CLI command (trl/scripts), and every trainer's documentation page has a runnable snippet.
• Scripts declare their dependencies in a # /// script header and document the exact run command(s) in the module docstring (python examples/<folder>/<script>.py ...), including multi-GPU variants where relevant.
• Add a row to the Index table in docs/source/example_overview.md, with a Colab badge if the example has a notebook that runs on free Colab.

Overview (`examples/async_grpo_harbor/README.md`) Copy Section

AsyncGRPO with Harbor

An installed agent runs each task through OpenEnv. TRL trains on the engine's captured token IDs,
logprobs, and loss masks, then updates the same vLLM instance.

Local setup Copy Chapter

Use Python 3.12 and this TRL checkout. The typed capture API was merged in
OpenEnv #1280. Install the pinned upstream revision:

git clone https://github.com/huggingface/OpenEnv.git OpenEnv
git -C OpenEnv checkout --detach 86a180ede21e044f7929b9a7783ad83aa67d83a3
pip install -e . trackio datasets ./OpenEnv
export PYTHONPATH="$PWD/OpenEnv/envs${PYTHONPATH:+:$PYTHONPATH}"

The source path is required because this revision's environment wheel does not package
harbor_env.harness. Both the core package and the harness above use the same revision.

Run the OpenEnv server with its sandbox credentials, then start vLLM and the trainer using the commands
in async_grpo_harbor.py. The server additionally needs pip install './OpenEnv[harbor]';
the trainer can use an existing server. Use the same model and vLLM endpoint for rollouts and weight updates.
Set the server's MAX_CONCURRENT_ENVS to at least --max-inflight + 1 (9 for the defaults); the factory
keeps one connection for task metadata. The Jobs launcher sets this capacity automatically.

Training saves a checkpoint every 50 steps, keeps the latest three, and saves the final model and
tokenizer in the output directory on successful completion, including runs shorter than 50 steps.

Hugging Face Jobs Copy Chapter

launcher.py starts the server, vLLM, and trainer on one two-GPU job. It installs the pinned
OpenEnv runtime and checks out the matching harness source. Pass HF_TOKEN and the sandbox credential
as job secrets. A bucket mounted at /data preserves checkpoints and logs.

Use --train-script-url with the raw URL of the TRL revision being tested before the example is merged.
Use --tunnel-check-s 0 to disable tunnel supervision. A server restart interrupts existing sessions;
subsequent sessions receive the new proxy URL.

Capture contract Copy Chapter

• OpenEnv supplies a validated TrainingTrace with required masks. Prompt and completion IDs come
  from the inference engine; TRL does not re-tokenize captured prompts or select agent turns.
• loss_mask covers prompt plus completion. Prompt positions are zero; completion positions may mix
  zeros and ones. TRL maps the completion span to its generic TurnRecord.output_mask.
• For example, prompt [10, 11], completion [12, 13, 14], and mask [0, 0, 1, 0, 1] train on tokens
  12 and 14, while retaining 13 as context.
• The worker constructs the session factory with its sampling policy. Pass a callable such as
  partial(HarborSessionFactory, server_url, ...); do not configure sampling separately on the factory.
  OpenEnv applies the policy to captured requests.
• Invalid captures stop the worker. Transport failures remain unscorable; an agent timeout retains
  valid captured turns and the verifier's score.

Execution status and structured tool outcomes are follow-up OpenEnv API work. Today, TRL recognizes
timeouts through TimeoutError and reconstructs conversation and tool diagnostics from captures.
Set agent turn and time budgets on the session factory; the worker only forwards sampling.

Rewritten histories can produce multiple training rows. Token retention does not guarantee equal
rollout weighting or a suitable memory budget for every harness.

Run the packaged contract and reconciliation tests with this checkout on PYTHONPATH:

python -m pytest tests/experimental/test_openenv_tito.py \
    tests/experimental/test_async_grpo_trainer.py::TestReconciler -q

Overview (`scripts/generate_tiny_models/README.md`) Copy Section

Tiny model generation

This directory contains one script per tiny model used by the TRL test suite. Each script builds a random-weight, minimally-sized model on top of a real tokenizer/processor and pushes it to the trl-internal-testing organization on the Hub.

Layout Copy Chapter

generate_tiny_models/
├── _common.py                    # shared helpers (push_to_hub, smoke_test, print_config_diff, ...)
├── for_causal_lm/                # *ForCausalLM + GPT-2 LM head + small/PEFT variants
├── for_sequence_classification/  # *ForSequenceClassification (reward models)
├── for_conditional_generation/   # *ForConditionalGeneration (VLMs + T5 + Bart encoder-decoder)
└── for_block_diffusion/          # *ForBlockDiffusion (DiffusionGemma)

Running Copy Chapter

From the repo root, invoke a script by its module path:

python -m scripts.generate_tiny_models.for_causal_lm.qwen3_for_causal_lm

Each script:

1. Checks that the installed transformers version matches the one pinned in the script (fails otherwise).
2. Calls set_seed to seed the RNGs that model construction draws from.
3. Builds the tiny model with random weights.
4. Runs smoke_test — a minimal forward pass to catch config misspecification and NaNs.
5. Runs check_dtype_pattern — reads the reference safetensors header via the Hub API and flags any tensor whose dtype diverges from the reference (catches e.g. fp32 norms kept inside a bf16 checkpoint).
6. Runs print_config_diff — prints every flat-key difference between the reference Hub config and the tiny model's config (for debugging scale-downs).
7. Pushes the model, tokenizer/processor, generation config, and model card to the Hub in a single commit.

If the repo already exists on the Hub, the push is skipped by default. Pass --create-pr to open a PR against the existing repo instead:

python -m scripts.generate_tiny_models.for_causal_lm.qwen3_for_causal_lm --create-pr

Direct overwrites of main aren't supported — update via --create-pr and merge the PR on the Hub.

Version pinning Copy Chapter

Each script calls check_transformers_version(...), which enforces:

max(version that introduced the model, TRL's transformers floor)

By default — check_transformers_version() with no argument — the expected version is the transformers>= floor read from pyproject.toml (currently 4.56.2). Scripts for models introduced after the floor pass an explicit pin, e.g. check_transformers_version("4.57.0") for Qwen3-VL or check_transformers_version("5.6.0") for Gemma4. The check is an exact match via packaging.version.Version; install the matching version before running.

Why exact? transformers is backward-compatible (a checkpoint saved by X loads on any ≥ X) but not forward-compatible. TRL CI runs against the floor, so tiny models must be saved with the oldest version that supports them — any newer save risks using config fields the floor can't parse. The exact-match check prevents accidental drift.

Seeding Copy Chapter

Every script calls set_seed() right after the version check, before the model is built. The shared SEED lives in _common.py.

This matters because HfApi.create_commit drops files that are unchanged, so a re-push of an unmodified tiny model should be a no-op and open no PR. Without a seed the weights differ on every run, so model.safetensors always changed and every regeneration produced a diff whether or not anything real had changed. Seeding also makes a config-only change produce a config-only Hub diff.

The call has to come before the model is constructed rather than before init_weights_tiny_model: most scripts never call that helper, and by then the randomness has already happened inside the model's own initialization.

Note that the seed fixes the RNG stream, not the library that produces it. check_transformers_version pins transformers exactly, but nn.init belongs to torch, which is pinned nowhere, so a torch upgrade can still change the weights.

Adding a new tiny model Copy Chapter

1. Pick the right subfolder based on the model class suffix (ForCausalLM, ForSequenceClassification, ForConditionalGeneration).
2. Copy an existing script with the closest shape and adapt it — reference model id, config class, model class, special kwargs.
3. If the model requires a release newer than the TRL floor, pass it explicitly: check_transformers_version("X.Y.Z"). Otherwise leave the call argument-less to default to the floor.
4. Keep the set_seed() call that the copied script already has, directly after the version check.
5. Run it. Inspect the [smoke_test] and [config_diff] output before letting it push.

Overview (`tests/invariant/README.md`) Copy Section

Training invariant tests

Catches silent training bugs that don't fail unit tests but shift the training trajectory. Runs on real models, opt-in only.

How it works Copy Chapter

Configs are grouped into equivalence classes: configs in the same class must produce the same trajectory (e.g. PDB=1×GAS=8 must equal PDB=8×GAS=1, FA2 must equal eager). Each class has one canonical config (the first one) that owns the class's reference snapshot. Every config in the class — canonical included — is asserted to match the saved reference. This catches both invariant breakage (a non-canonical config drifting away from the canonical's pinned trajectory) and numerical regressions (the canonical itself drifting from its committed snapshot across versions).

Recording the references is a separate concern from testing them, so it's a separate entry point (python -m tests.invariant.test_invariant).

Each config is a trl <method> CLI invocation with a fixed set of args. The harness shells out (subprocess.run(["trl", method, ...])), the CLI runs end to end, writes trainer_state.json to its --output_dir, and the harness parses the log_history into a Trajectory.

This means the suite tests the actual user-facing entry point, not the Python API. Catches CLI-only bugs (arg parsing, defaults, dispatch) for free. Distributed runs are an additive change: prepend accelerate launch --config_file <strategy>.yaml to the same command.

Scope (initial) Copy Chapter

• Trainers: trl sft, trl dpo
• Model: Qwen/Qwen2.5-0.5B-Instruct (pinned revision)
• Equivalence classes:
  • sft: sft_default (canonical), sft_pdb1_gas8 (gradient accumulation), sft_attn_fa2_kernels (FA2 via kernels)
  • dpo: dpo_default (canonical), dpo_pdb1_gas8 (gradient accumulation)
• Single GPU, fp32, fixed seed, ~50 optimizer steps.

Other axes (sharding, DDP, more trainers) are deferred and will be additive.

Trajectory Copy Chapter

Per optimizer step: loss, grad_norm. One JSON per equivalence class in references/ (sft.json, dpo.json):

{
  "config": {"name": "sft_default", "method": "sft", "args": {...}},
  "env":    {"accelerate": "...", "torch": "...", "transformers": "...", "trl": "...", "gpu": "H100-80GB"},
  "steps":  [{"step": 1, "loss": 1.234, "grad_norm": 0.567}, ...]
}

Comparison Copy Chapter

Scalar series with absolute tolerance + zero-mean-residual. The residual check is what flags bugs like GAS-dropping — they show up as a one-sided systematic shift in the loss curve, not as point-wise outliers.

Hardware Copy Chapter

Reference snapshots are recorded on H100 80GB, except sft_fa2.json (H200). The recording hardware is in each snapshot's env block.

Running Copy Chapter

# test
pytest tests/invariant/ -m invariant

# record references
python -m tests.invariant.test_invariant      # all classes
python -m tests.invariant.test_invariant sft  # one class

Snapshot updates must be justified in the PR description.

Overview (`trl/chat_templates/README.md`) Copy Section

Chat Templates

Jinja2 chat templates stored here serve two purposes:

1. Identity comparison: detecting which model is being used (by comparing tokenizer.chat_template against known templates) to add the appropriate response schema (add_response_schema) or swap in a training template (get_training_chat_template).
2. Training patches: modified templates that fix training-specific issues (prefix-preservation for GRPO, {% generation %} markers for SFT assistant-only loss).

Why prefix-preserving? The GRPO tool call loop extracts tool response formatting tokens by comparing tokenizations with and without tool messages appended (_get_tool_suffix_ids). This requires the chat template to be prefix-preserving: appending messages must not change how earlier messages are rendered.

Why generation-tagged? SFT with assistant_only_loss=True requires the chat template to include {% generation %} / {% endgeneration %} markers around assistant output, so return_assistant_tokens_mask=True can produce correct masks. Most model templates don't include these markers natively.

Why _v2? Identity comparison is exact, so when a model repo changes its chat template, the new revision is stored next to the original as <name>_v2.jinja (then _v3, ...) and both are recognized. The suffix only orders revisions as they are added; each section below says where the revision comes from.

How is drift detected? A weekly CI job runs scripts/check_chat_templates.py, which fails when one of its reference Hub repos ships a chat template that is not stored here. When adding a template for a new model, add the repo it comes from to REPOS in that script.

Original templates Copy Chapter

Used for identity comparison only.

`cohere.jinja` Copy Section

Original Cohere chat template (as shipped by CohereLabs/aya-expanse-8b).

`cohere2.jinja` Copy Section

Original Cohere2 chat template (as shipped by CohereLabs/tiny-aya-earth).

`deepseek_r1_distill.jinja` Copy Section

Original DeepSeek-R1-Distill chat template, shipped byte-identically by every distill checkpoint (both the Qwen and the Llama ones). It prefills the assistant turn with <｜Assistant｜>\n, so the model's output starts inside the reasoning block; response parsing uses deepseek_r1_distill_template. Assistant turns are re-rendered without their reasoning (the template drops everything up to and including ``), and tool_calls are silently dropped whenever content is set, so the template does not support tool calling.

`deepseekv3.jinja` Copy Section

Original DeepSeek-V3 chat template.

`diffusion_gemma.jinja` Copy Section

Original DiffusionGemma chat template (as shipped by google/diffusiongemma-26B-A4B-it).

`gemma.jinja` Copy Section

Original Gemma chat template. Used by both Gemma (v1) and Gemma2, which ship identical templates.

`gemma3.jinja` Copy Section

Original Gemma 3 chat template (as shipped by google/gemma-3-* checkpoints). Unlike gemma.jinja, this template supports system messages and renders multimodal content blocks (text and image) for VLM training.

`gemma4.jinja` Copy Section

Original Gemma 4 chat template (as shipped by google/gemma-4-E2B-it and the other checkpoints in that generation, at 6b7e72c). Renders reasoning (<|channel>thought\n...<channel|>), assistant tool_calls (<|tool_call>call:<name>{...}<tool_call|> blocks), and multimodal content blocks. Response parsing uses gemma4_template.

`gemma4_v2.jinja` Copy Section

Gemma 4 chat template as shipped by google/gemma-4-E2B-it at 4742fe8: the revision the model was released with. Unlike gemma4.jinja, it does not render reasoning in assistant turns that carry tool_calls. Response parsing reuses gemma4_template.

`gemma4_v3.jinja` Copy Section

Gemma 4 chat template as shipped by google/gemma-4-E2B-it at b446025: the revision between the release and gemma4.jinja. Unlike gemma4_v2.jinja, it renders reasoning in assistant turns that carry tool_calls. Response parsing reuses gemma4_template.

`gemma4_v4.jinja` Copy Section

Gemma 4 chat template as shipped by google/gemma-4-E2B-it at 905e84b: the revision after gemma4.jinja. Differs from it only in emitting <|image|>, <|audio|> and <|video|> placeholders for multimodal tool responses. Response parsing reuses gemma4_template.

`gemma4_v5.jinja` Copy Section

Gemma 4 chat template as shipped by google/gemma-4-E2B-it at 8995379: the current revision. Unlike the earlier ones, it also renders the reasoning of a final assistant turn. Response parsing reuses gemma4_template.

`glm4moe.jinja` Copy Section

Original GLM-4-MoE chat template.

`gptoss.jinja` Copy Section

Original GPT-OSS chat template.

`idefics3.jinja` Copy Section

Original Idefics3 chat template (as shipped by HuggingFaceM4/Idefics3-8B-Llama3). Does not support tool calling.

`lfm2.jinja` Copy Section

Original LFM2 chat template (as shipped by LiquidAI/LFM2-* checkpoints). ChatML-style. Renders tools into the system prompt wrapped in <|tool_list_start|> / <|tool_list_end|>, and wraps tool message content in <|tool_response_start|> / <|tool_response_end|>. It never reads message['tool_calls'] though, so assistant tool calls are silently dropped: the model is trained to emit <|tool_call_start|> / <|tool_call_end|> as plain text inside content.

`lfm2_v2.jinja` Copy Section

Later revision of the LFM2 chat template (as shipped by LiquidAI/LFM2-1.2B since 40f3da0). Unlike lfm2.jinja, it renders assistant tool_calls, in the same <|tool_call_start|>[name(key=value, ...)]<|tool_call_end|> format as lfm2_2_5.jinja, and tolerates a missing content alongside them. Response parsing reuses lfm2_2_5_template.

`lfm2_2_5.jinja` Copy Section

Original LFM2.5 chat template (as shipped by LiquidAI/LFM2.5-230M and the other checkpoints in that generation). Unlike lfm2.jinja, it renders assistant tool_calls — as a single <|tool_call_start|>[name(key=value, ...)]<|tool_call_end|> block holding a comma-separated list of Python-style calls — supports a `` block (read off message.thinking), and already carries {% generation %} markers, so no training patch is needed. Response parsing uses lfm2_2_5_template.

`lfm2_2_5_v2.jinja` Copy Section

Later revision of the LFM2.5 chat template (as shipped by LiquidAI/LFM2.5-230M since 13a5383). Differs from lfm2_2_5.jinja only in reading the `` block off message.thinking, falling back to message.reasoning and message.reasoning_content, and rendering it only when it is a non-empty string. Response parsing reuses lfm2_2_5_template.

`lfm2_2_5_v3.jinja` Copy Section

Later LFM2.5 chat template (as shipped by LiquidAI/LFM2.5-24B-A2B and LiquidAI/LFM2.5-2.6B). Same tool-call format as lfm2_2_5_v2.jinja, with two rendering differences: the generation prompt pre-fills ``, so the model's output starts inside the reasoning block, and string arguments are escaped (\\, \', \n, \r) while lists render as JSON. Already carries {% generation %} markers, so no training patch is needed. Response parsing uses lfm2_2_5_v3_template.

`lfm2_2_5_vl.jinja` Copy Section

Original LFM2.5-VL chat template (as shipped by LiquidAI/LFM2.5-VL-3B). Same rendering as lfm2_2_5.jinja — including the {% generation %} markers, so no training patch is needed — plus multimodal content blocks (each image item renders as ``) and more defensive handling of non-string content and flat tool_calls. Response parsing reuses lfm2_2_5_template.

`llama3.jinja` Copy Section

Original Llama 3 chat template.

`llama3_1.jinja` / `llama3_2.jinja` Copy Section

Original Llama 3.1 / 3.2 chat templates. Both render tool calls as a single bare JSON object using the key parameters (instead of arguments) and support at most one tool call per assistant turn.

`llava_next.jinja` Copy Section

Original Llava-Next chat template (as shipped by llava-hf/llava-v1.6-mistral-7b-hf). Renders multimodal content blocks in the LLaVA / Mistral [INST] ... [/INST] format. Does not support tool calling.

`muse_glimmer.jinja` Copy Section

Original Muse Glimmer chat template (as shipped by meta-models/Muse-Glimmer-30B). Renders the ATEM protocol: one <|start|>assistant to=<recipient><|message|>... block per channel, where the recipient is self for reasoning, a tool name for tool calls (<atem:function_calls> / <atem:invoke> / <atem:parameter> XML), and user for content. The checkpoint ships its own response_template, so add_response_schema leaves it alone.

`nemotron_3_nano.jinja` Copy Section

Original Nemotron Nano chat template (as shipped by nvidia/NVIDIA-Nemotron-3-Nano-* checkpoints). Renders tool calls in the same Hermes-style <function=...> / `` format as Qwen3.5, so it reuses qwen3_5_schema for response parsing.

`nemotron_3_super.jinja` Copy Section

Original Nemotron Super chat template (as shipped by nvidia/NVIDIA-Nemotron-3-Super-* checkpoints). Same as nemotron_3_nano.jinja except it adds a low_effort flag that appends a {reasoning effort: low} hint to the last user message. Tool calls use the same Hermes-style format, so it also reuses qwen3_5_schema for response parsing.

`nemotron_3_ultra.jinja` Copy Section

Original Nemotron Ultra chat template (as shipped by nvidia/NVIDIA-Nemotron-3-Ultra-* checkpoints). Same as nemotron_3_nano.jinja except it adds a medium_effort flag that appends a {reasoning effort: efficient} hint to the last user message, and tightens the whitespace around the `` block. Tool calls use the same Hermes-style format, so it also reuses qwen3_5_schema for response parsing.

`nemotron_3_5_lightning.jinja` Copy Section

Original Nemotron 3.5 Lightning chat template (as shipped by nvidia/NVIDIA-Nemotron-3.5-Lightning-* checkpoints). Same as nemotron_3_ultra.jinja except it drops the medium_effort flag. Tool calls use the same Hermes-style format, so it also reuses qwen3_5_schema for response parsing.

`phi3.jinja` Copy Section

Original Phi-3 chat template.

`phi3_5.jinja` Copy Section

Original Phi-3.5 chat template.

`qwen2_5.jinja` Copy Section

Original Qwen2.5 chat template.

`qwen2_5_vl.jinja` Copy Section

Original Qwen2.5-VL chat template. Also matches Qwen2-VL, which ships a byte-identical template. Does not support tool calling.

`qwen3.jinja` Copy Section

Original Qwen3 chat template.

`qwen3_instruct_2507.jinja` Copy Section

Original Qwen3-Instruct-2507 chat template (as shipped by Qwen/Qwen3-4B-Instruct-2507). A simpler Qwen3 variant: unlike qwen3.jinja, it has no reasoning_content / `` handling, no multi_step_tool tracking and no enable_thinking flag. Response parsing uses qwen3_template.

`qwen3_vl.jinja` Copy Section

Original Qwen3-VL chat template. Unlike text-only Qwen3, this template is already prefix-preserving (no conditional thinking blocks), so no training patch is needed.

`qwen3_5_think.jinja` / `qwen3_5_nothink.jinja` Copy Section

Original Qwen3.5 chat templates. The two differ only in the default value of the enable_thinking flag in the generation-prompt block:

• qwen3_5_think.jinja — defaults to thinking enabled. Shipped by Qwen3.5-4B and larger.
• qwen3_5_nothink.jinja — defaults to thinking disabled. Shipped by Qwen3.5-2B and smaller.

`qwen3_6.jinja` Copy Section

Original Qwen3.6 chat template (shared across Qwen3.6-27B, Qwen3.6-35B-A3B, and their FP8 variants). Differs from qwen3_5_think.jinja by adding a preserve_thinking flag and tweaking how non-string tool-call argument values are stringified.

`qwen3_8.jinja` Copy Section

Original Qwen3.8 chat template (as shipped by Qwen/Qwen3.8-27B and its FP8 variant). Differs from qwen3_6.jinja by adding a reasoning_effort flag (xhigh — the default — medium or low) that prepends a reasoning-effort instruction to the system prompt, and by defaulting preserve_thinking to enabled. Tool calls use the same Hermes-style format, so it also reuses qwen3_5_schema for response parsing.

`smolvlm.jinja` Copy Section

Original SmolVLM chat template. Also matches SmolVLM2, which ships a byte-identical template. Does not support tool calling.

Training templates Copy Chapter

Patched templates that fix training-specific issues. Swapped in at init when tools are enabled (GRPO) or when assistant_only_loss=True (SFT).

`cohere_training.jinja` Copy Section

Patched Cohere template. Diff vs cohere.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`cohere2_training.jinja` Copy Section

Patched Cohere2 template. Diff vs cohere2.jinja:

Move the trailing <|END_OF_TURN_TOKEN|> from after the role-dispatch {% endif %} into each role branch, so it can be wrapped together with the assistant content. Wrap the assistant branch (<|START_RESPONSE|>...<|END_RESPONSE|><|END_OF_TURN_TOKEN|>) with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`deepseek_r1_distill_training.jinja` Copy Section

Patched DeepSeek-R1-Distill template. Diff vs deepseek_r1_distill.jinja:

• Removes the '' split on assistant content, so the reasoning block is always rendered. The original drops everything up to and including `` — fine for inference history, but it would silently remove the reasoning from the SFT training target.
• Wraps assistant message output with {% generation %} / {% endgeneration %} markers for SFT assistant-only loss.

`deepseekv3_training.jinja` Copy Section

Patched DeepSeek-V3 template. Diff vs deepseekv3.jinja:

• Uses | tojson on tool['function']['arguments'] so that arguments can be passed as a dict (the documented format per transformers docs). The original template uses raw string concatenation, which crashes on dict inputs.
• Wraps assistant message output with {% generation %} / {% endgeneration %} markers for SFT assistant-only loss.

`diffusion_gemma_training.jinja` Copy Section

Patched DiffusionGemma template. Diff vs diffusion_gemma.jinja:

Wrap the parts of a model turn that the model emits (thinking channel, tool calls, content, and the closing <turn|>) with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss. Tool responses embedded in the turn come from the environment and stay outside the generation block.

`gemma_training.jinja` Copy Section

Patched Gemma template, shared by Gemma (v1) and Gemma2 (which ship identical chat templates). Diff vs gemma.jinja:

Split the unified message output line into role-specific branches, so the <start_of_turn>model\n prompt cue sits outside the generation block (it is not generated by the model), while the assistant's content and <end_of_turn>\n (which the model must learn to produce and to stop on) sit inside. Wrap the assistant content with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`gemma3_training.jinja` Copy Section

Patched Gemma 3 template. Same diff as gemma_training.jinja (split the unified output line into role-specific branches so the <start_of_turn>model\n prompt cue sits outside the generation block, and wrap the assistant content with {% generation %} / {% endgeneration %}), applied to the Gemma 3 base template that supports system messages and multimodal content blocks.

`glm4moe_training.jinja` Copy Section

Patched GLM-4-MoE template. Diff vs glm4moe.jinja:

Require both and to be present before parsing, to avoid incorrect splitting when the model generates only one tag:

- {%- if '</think>' in content %}
+ {%- if '<think>' in content and '</think>' in content %}

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`gptoss_training.jinja` Copy Section

Patched GPT-OSS template. Diff vs gptoss.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`idefics3_training.jinja` Copy Section

Patched Idefics3 template. Diff vs idefics3.jinja:

Split the assistant message into its own branch so the {% generation %} / {% endgeneration %} markers wrap the assistant content. This enables return_assistant_tokens_mask=True to produce correct masks for SFT assistant-only loss.

`lfm2_training.jinja` Copy Section

Patched LFM2 template. Diff vs lfm2.jinja:

Split the unified message output line into role-specific branches, so the <|im_start|>assistant\n prompt cue sits outside the generation block (it is not generated by the model), while the assistant's content and <|im_end|>\n (which the model must learn to produce and to stop on) sit inside. Wrap the assistant content with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`lfm2_v2_training.jinja` Copy Section

Patched LFM2 template, later revision. Diff vs lfm2_v2.jinja:

Wrap the assistant output with {% generation %} / {% endgeneration %}, in both the branch that renders tool_calls and the plain-content branch, so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss. The <|im_start|>assistant\n prompt cue stays outside the generation block, while the content, the tool calls and <|im_end|>\n sit inside.

`llama3_training.jinja` Copy Section

Patched Llama 3 template. Diff vs llama3.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`llava_next_training.jinja` Copy Section

Patched Llava-Next template. Diff vs llava_next.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`muse_glimmer_training.jinja` Copy Section

Patched Muse Glimmer template. Diff vs muse_glimmer.jinja:

Wrap the whole assistant branch — every ATEM channel of the turn — with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss. Same treatment as gptoss_training.jinja, which shares the <|start|>assistant generation cue and the one-block-per-channel layout.

`nemotron_3_nano_training.jinja` Copy Section

Patched Nemotron Nano template. Diff vs nemotron_3_nano.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`nemotron_3_super_training.jinja` Copy Section

Patched Nemotron Super template. Diff vs nemotron_3_super.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`nemotron_3_ultra_training.jinja` Copy Section

Patched Nemotron Ultra template. Diff vs nemotron_3_ultra.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`nemotron_3_5_lightning_training.jinja` Copy Section

Patched Nemotron 3.5 Lightning template. Diff vs nemotron_3_5_lightning.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`phi3_training.jinja` Copy Section

Patched Phi-3 template. Diff vs phi3.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`phi3_5_training.jinja` Copy Section

Patched Phi-3.5 template. Diff vs phi3.5.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`qwen2_5_training.jinja` Copy Section

Patched Qwen2.5 template. Diff vs qwen2_5.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`qwen2_5_vl_training.jinja` Copy Section

Patched Qwen2.5-VL template (also used for Qwen2-VL, which ships a byte-identical template). Diff vs qwen2_5_vl.jinja:

Split the assistant message into its own branch so the {% generation %} / {% endgeneration %} markers wrap the assistant content. This enables return_assistant_tokens_mask=True to produce correct masks for SFT assistant-only loss.

`qwen3_training.jinja` Copy Section

Patched Qwen3 template. Diff vs qwen3.jinja:

Require both and to be present before parsing, to avoid incorrect splitting when the model generates only one tag:

- {%- if '</think>' in content %}
+ {%- if '<think>' in content and '</think>' in content %}

Always include the thinking block regardless of message position. The original conditionally omits it based on loop.last, which changes the assistant rendering when a tool message is appended — breaking prefix-preservation:

- {%- if loop.index0 > ns.last_query_index %}
-     {%- if loop.last or (not loop.last and reasoning_content) %}
-         {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
-     {%- else %}
-         {{- '<|im_start|>' + message.role + '\n' + content }}
-     {%- endif %}
- {%- else %}
-     {{- '<|im_start|>' + message.role + '\n' + content }}
- {%- endif %}
+ {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`qwen3_instruct_2507_training.jinja` Copy Section

Patched Qwen3-Instruct-2507 template. Diff vs qwen3_instruct_2507.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`qwen3_vl_training.jinja` Copy Section

Patched Qwen3-VL template. Diff vs qwen3_vl.jinja:

Wrap assistant message output (both content and tool_calls) with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`qwen3_5_think_training.jinja` / `qwen3_5_nothink_training.jinja` Copy Section

Patched Qwen3.5 templates. Same diff as qwen3_training.jinja (require both and before parsing, drop the loop.index0 > ns.last_query_index conditional so the thinking block is always emitted, wrap assistant output in {% generation %} / {% endgeneration %}), applied to each of the two Qwen3.5 base templates. The two training variants differ only in the default value of the enable_thinking flag, inherited from their respective base templates.

`qwen3_6_training.jinja` Copy Section

Patched Qwen3.6 template. Same diff as qwen3_training.jinja (require both and before parsing, drop the loop.index0 > ns.last_query_index conditional so the thinking block is always emitted, wrap assistant output in {% generation %} / {% endgeneration %}), applied to the Qwen3.6 base template.

`qwen3_8_training.jinja` Copy Section

Patched Qwen3.8 template. Diff vs qwen3_8.jinja: drop the preserve_thinking / loop.index0 > ns.last_query_index conditional so the thinking block is always emitted (prefix-preservation holds even when the caller passes preserve_thinking=False), and wrap assistant output in {% generation %} / {% endgeneration %} for SFT assistant-only loss.

`smolvlm_training.jinja` Copy Section

Patched SmolVLM template (also used for SmolVLM2, which ships a byte-identical template). Diff vs smolvlm.jinja:

Split the assistant message into its own branch so the {% generation %} / {% endgeneration %} markers wrap the assistant content. This enables return_assistant_tokens_mask=True to produce correct masks for SFT assistant-only loss.

Overview (`examples/grpo_harbor/harnesses/bash/README.md`) Copy Section

`bash` harness

The minimal harness: a single shell tool. This is the built-in HarborBashEnv (trl.experimental.harbor.HarborBashEnv); this folder just documents it and re-exports it as BashEnv.

Tools Copy Chapter

Submission Copy Chapter

No submit tool — write the answer to /workdir/answer.txt, e.g. echo -n "<value>" > /workdir/answer.txt. The task's verifier reads that file.

Use it Copy Chapter

from trl.experimental.harbor import HarborSpec
spec = HarborSpec(dataset, agent="bash")   # built-in name

Overview (`examples/grpo_harbor/harnesses/jupyter/README.md`) Copy Section

`jupyter` harness

A stateful Python kernel harness. Variables, imports, and side-effects persist across cells, so the model builds up state like a notebook. Implemented by JupyterEnv in env.py.

How it works Copy Chapter

On reset, the harness (_setup) uploads two helpers into the sandbox and starts a tiny kernel server:

• kernel_server.py — a local HTTP server (127.0.0.1:8765) holding one persistent Python kernel.
• run_cell.py — sends a base64-encoded cell to the kernel server and prints its JSON result.

Each add_and_execute_code_cell call runs python3 /opt/run_cell.py against that kernel.

Tools Copy Chapter

Submission Copy Chapter

Write the answer to /workdir/answer.txt — e.g. add_and_execute_code_cell(code="open('/workdir/answer.txt','w').write(str(ans))").

Use it Copy Chapter

from trl.experimental.harbor import HarborSpec
spec = HarborSpec(dataset, agent="examples/grpo_harbor/harnesses/jupyter/env.py:JupyterEnv")

Overview (`examples/grpo_harbor/harnesses/terminal_notes/README.md`) Copy Section

`terminal_notes` harness

A richer shell harness: 6 shell tools (including background processes) plus a 4-tool persistent note toolkit. Implemented by TerminalNotesEnv in env.py.

Tools Copy Chapter

Shell (backed by the sandbox) Copy Section

Notes (in-env state, persist across turns of a rollout) Copy Section

Unlike the original SETA agent, notes are not auto-injected into the prompt each turn (TRL owns the prompt under environment_factory); recall them on demand with read_note / list_note.

Submission Copy Chapter

Write the answer to /workdir/answer.txt, e.g. shell_write_content_to_file(path="/workdir/answer.txt", content="<value>").

Use it Copy Chapter

from trl.experimental.harbor import HarborSpec
spec = HarborSpec(dataset, agent="examples/grpo_harbor/harnesses/terminal_notes/env.py:TerminalNotesEnv")

Source Example Overview (`docs/source/example_overview.md`) Copy Section

Examples

The examples/ directory contains a collection of self-contained examples that demonstrate how to use the TRL library for various applications. Each example lives in its own folder named after the method and the task it demonstrates (e.g. grpo_wordle, sft_gpt_oss), and holds everything the example needs: scripts, notebooks, prompts, chat templates, and evaluation code.

Basic single-trainer training scripts are not examples: they live in trl/scripts and are exposed through the command line interface (trl sft, trl dpo, trl grpo, …). Each trainer's documentation page also contains a complete runnable snippet.

Shared resources sit at the root of examples/:

• examples/accelerate_configs: 🤗 Accelerate configuration files for multi-GPU, DeepSpeed ZeRO, FSDP, and context-parallel setups, used by many examples.
• examples/datasets: the scripts used to generate the trl-lib datasets used across the examples.

Getting Started

Install TRL and additional dependencies as follows:

pip install --upgrade trl[quantization]

Check for additional optional dependencies here. Notebook-based examples are self-contained and can run on free Colab; script-based examples run on single-GPU, multi-GPU, or DeepSpeed setups (see Distributed Training below).

Index Copy Chapter

Distributed Training Copy Chapter

You can run the example scripts on multiple GPUs with 🤗 Accelerate:

accelerate launch --config_file=examples/accelerate_configs/multi_gpu.yaml --num_processes {NUM_GPUS} path_to_script.py --all_arguments_of_the_script

For DeepSpeed ZeRO-{1,2,3}:

accelerate launch --config_file=examples/accelerate_configs/deepspeed_zero{1,2,3}.yaml --num_processes {NUM_GPUS} path_to_script.py --all_arguments_of_the_script

Adjust NUM_GPUS and --all_arguments_of_the_script as needed.

Source Experimental Overview (`docs/source/experimental_overview.md`) Copy Section

Experimental

This directory contains a minimal, clearly separated space for fast iteration on new ideas.

[!WARNING]
Stability contract: Anything under trl.experimental may change or be removed in any release (including patch versions) without prior deprecation. Do not rely on these APIs for production workloads.

Promotion Path (Simple) Copy Chapter

1. Prototype outside the main repo: Start development in your own fork or a separate repository to iterate quickly.
2. Experimental inclusion: Once it’s ready for early users, move the idea into trl.experimental.<feature>.
3. Improve: Add tests, a short doc/example, and demonstrate the usage.
4. Promote: Once the API proves stable and there is clear interest or adoption from the community, move it into trl.<feature> (stable module).

Removal Copy Chapter

Experimental is a staging area, not long-term storage: a feature is either promoted or removed. Removal needs no deprecation cycle, as stated in the stability contract above. The code remains available in the git history, and the paper index entry stays, pointing at the last release that shipped the implementation.

No threshold triggers a removal. It is a judgment call, made in the removal pull request with the numbers in front of everyone. What that judgment weighs:

• Usage. Telemetry, Hub tags, and replies on the pull request or on a discussion thread. None of them is conclusive on its own, but all of them quiet is a strong signal.
• External issues and pull requests that come from running the feature, as opposed to reports produced by scanning the code.
• Whether a stable trainer already covers it. If one does, that outweighs usage.
• Cost. CI time, maintenance, triage load.
• Downstream consumers. Libraries that import it.
• Owner. Someone who wants it promoted and is willing to do the work.
• Age. Whether it has been available long enough for the community to find it and try it.

FAQ Copy Chapter

Why not just use branches?
Because branches are not shipped to users; experimental code inside the package lets early adopters try things and give feedback.

Can these APIs change or vanish without warning?
Yes. Anything inside trl.experimental can change or disappear in any release.

Should I use this in production?
Only if you are fine with updating your code quickly when things change.

Will maintainers promptly fix issues in trl.experimental?
Not necessarily. The experimental module is a playground for new ideas, and maintainers may not prioritize bug fixes or feature requests there. Issues may remain unresolved until (or unless) the feature graduates to the stable API.

I contributed a paper implementation. Can it be removed?
Yes. Anything under trl.experimental can be removed, paper implementations included. If that happens, the paper index entry stays and points at the last release that shipped the code.

How to silence the runtime notice?

Use: export TRL_EXPERIMENTAL_SILENCE=1.

Source Index (`docs/source/index.md`) Copy Section

TRL - Transformers Reinforcement Learning

TRL is a full stack library where we provide a set of tools to train transformer language models with methods like Supervised Fine-Tuning (SFT), Group Relative Policy Optimization (GRPO), Direct Preference Optimization (DPO), Reward Modeling, and more.
The library is integrated with 🤗 transformers.

🎉 What's New Copy Chapter

🚀 Training on AMD GPUs: TRL's trainers now run on AMD Instinct GPUs through PyTorch's ROCm build.

🔥 Beam: Fine-tune Beam with [SFTTrainer] on a single node of 8 B300s. See the LoRA example.

🔥 Muse Spark Open: Fine-tune Muse Spark Open with [SFTTrainer], sharded across your GPUs at load time. See the 8-node LoRA example.

🚀 Training on AWS Trainium: TRL's trainers now run on AWS Trainium through TorchNeuron, PyTorch's native Neuron backend.

Taxonomy Copy Chapter

Below is an overview of TRL trainers, organized by maturity and method type.

Online methods

• GRPOTrainer
• RLOOTrainer

Reward modeling

• RewardTrainer

Offline methods

• SFTTrainer
• DPOTrainer
• KTOTrainer

Knowledge distillation

• DistillationTrainer

Experimental Copy Section

Online methods

• A2POTrainer
• AsyncGRPOTrainer
• GMPOTrainer
• OnlineDPOTrainer

Offline methods

• BEMA for Reference Model
• CPOTrainer
• ORPOTrainer
• TPOTrainer

Knowledge distillation

• AsyncDistillationTrainer
• GKDTrainer
• GOLDTrainer
• IWOPDTrainer
• SDFTTrainer
• SDPOTrainer

You can also explore TRL-related models, datasets, and demos in the TRL Hugging Face organization.

Learn Copy Chapter

Learn post-training with TRL and other libraries in 🤗 smol course.

Contents Copy Chapter

The documentation is organized into the following sections:

• Getting Started: installation and quickstart guide.
• Conceptual Guides: dataset formats, training FAQ, and understanding logs.
• How-to Guides: reducing memory usage, speeding up training, distributing training, etc.
• Integrations: DeepSpeed, PEFT, etc.
• Examples: example overview, community tutorials, etc.
• API: trainers, utils, etc.

Blog posts Copy Chapter

<a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/delta-weight-sync">
  
  Published on May 27, 2026
  Shipping a Trillion Parameters With a Hub Bucket: Delta Weight Sync in TRL
</a>
<a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/trl-v1">
  
  Published on March 31, 2026
  TRL v1: Post-Training Library That Holds When the Field Invalidates Its Own Assumptions
</a>
<a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/async-rl-training-landscape">
  
  Published on March 10, 2026
  Keep the Tokens Flowing: Lessons from 16 Open-Source RL Libraries
</a>
<a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/ulysses-sp">
  
  Published on March 9, 2026
  Ulysses Sequence Parallelism: Training with Million-Token Contexts
</a>
<a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/openenv">
  
  Published October 23, 2025
  Building the Open Agent Ecosystem Together: Introducing OpenEnv
</a>
<a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/trl-vlm-alignment">
  
  Published on August 7, 2025
  Vision Language Model Alignment in TRL ⚡️
</a>
<a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/vllm-colocate">
  
  Published on June 3, 2025
  NO GPU left behind: Unlocking Efficiency with Co-located vLLM in TRL
</a>
<a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/liger-grpo">
  
  Published on May 25, 2025
  🐯 Liger GRPO meets TRL
</a>
<a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/open-r1">
  
  Published on January 28, 2025
  Open-R1: a fully open reproduction of DeepSeek-R1
</a>
<a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/dpo_vlm">
  
  Published on July 10, 2024
  Preference Optimization for Vision Language Models with TRL
</a>
<a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/putting_rl_back_in_rlhf_with_rloo">
  
  Published on June 12, 2024
  Putting RL back in RLHF
</a>
<a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/unsloth-trl">
  
  Published on January 10, 2024
  Make LLM Fine-tuning 2x faster with Unsloth and 🤗 TRL
</a>
<a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/trl-ddpo">
  
  Published on September 29, 2023
  Finetune Stable Diffusion Models with DDPO via TRL
</a>
<a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/dpo-trl">
  
  Published on August 8, 2023
  Fine-tune Llama 2 with DPO
</a>
<a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/stackllama">
  
  Published on April 5, 2023
  StackLLaMA: A hands-on guide to train LLaMA with RLHF

Published on March 9, 2023
  Fine-tuning 20B LLMs with RLHF on a 24GB consumer GPU
</a>
<a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/rlhf">
  
  Published on December 9, 2022
  Illustrating Reinforcement Learning from Human Feedback
</a>

<a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/datasets/trl-lib/documentation-images/resolve/main/Fine%20tuning%20with%20TRL%20(Oct%2025).pdf">
  
  Talk given on October 30, 2025
  Fine tuning with TRL
</a>

📜 Paper: https://huggingface.co/papers/2402.03300

Introduces Group Relative Policy Optimization (GRPO) and shows strong math-reasoning gains from math-centric pretraining plus group-relative PPO-style optimization. Used in TRL via [GRPOTrainer].

from trl import GRPOConfig, GRPOTrainer

# The paper doesn't specify its hyperparameters, so here we provide hyperparameters from "DeepSeek-R1 incentivizes reasoning in LLMs through reinforcement learning" instead.
training_args = GRPOConfig(
    loss_type="grpo",
    beta=0.001,  # "the KL coefficient to 0.001"
    epsilon=10.0, # "the GRPO clip ratio ϵ to 10"
    num_generations=16,  # "For each question, we sample 16 outputs..."
    max_completion_length=32_768,  # "...with a maximum length of 32,768"
    steps_per_generation=16,  # "To accelerate training, each rollout generates 8,192 outputs, which are randomly split into 16 minibatches"
    # "resulting in a training batch size of 512". One way to achieve this setting with 1 device is per_device_train_batch_size=4, gradient_accumulation_steps=128
    per_device_train_batch_size=4,
    gradient_accumulation_steps=128,  
)
trainer = GRPOTrainer(
    ...,
    args=training_args,
)

DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning Copy Section

📜 Paper: https://huggingface.co/papers/2501.12948

DeepSeek-R1 achieves reasoning performance comparable to OpenAI-o1 through a multi-stage pipeline that transitions from pure reinforcement learning (RL) to a refined, human-aligned model. Unlike its predecessor, DeepSeek-R1-Zero, which used pure RL on a base model, R1 follows a structured four-stage evolution:

1. Cold Start: The base model is fine-tuned on a small set of high-quality, long Chain-of-Thought (CoT) data to provide a stable starting point.
2. Reasoning-Oriented RL: Large-scale RL is applied to enhance performance in math, coding, and logic, using rule-based rewards and a language consistency reward to reduce language mixing.
3. Rejection Sampling & SFT: The RL checkpoint generates 600k reasoning samples via rejection sampling, which are combined with 200k non-reasoning (general) samples to create a new dataset for a second round of Supervised Fine-Tuning.
4. RL for all Scenarios: A final RL stage aligns the model with human preferences (helpfulness and harmlessness) across all domains while maintaining reasoning strength.

Distillation: Empowering Small Models

A key contribution of the paper is demonstrating that reasoning patterns can be distilled from a large model (DeepSeek-R1) into smaller dense models (e.g., Qwen and Llama series). Distillation was found to be more effective for small models than training them with pure RL from scratch.

You can use the GRPOTrainer to replicate the reasoning-heavy stages of this pipeline.

from trl import GRPOConfig, GRPOTrainer

# Example configuration for a reasoning-oriented GRPO stage
# Based on the Open-R1 recipe for Qwen-7B
training_args = GRPOConfig(
    learning_rate=4.0e-5,
    max_completion_length=32768, # Support for long Chain-of-Thought
    num_generations=16,          # Sample 16 outputs per prompt for group relative advantage
    beta=0.001,                  # KL coefficient
    use_vllm=True,               # Use vLLM backend for accelerated rollout generation
)

trainer = GRPOTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    reward_funcs=[accuracy_reward, format_reward], # R1-Zero used rule-based rewards
)

trainer.train()

Group Sequence Policy Optimization Copy Section

📜 Paper: https://huggingface.co/papers/2507.18071

GSPO is a GRPO variant that computes importance sampling weights at the sequence level instead of per-token. To reproduce the paper's setting, use this configuration:

from trl import GRPOConfig

training_args = GRPOConfig(
    importance_sampling_level="sequence",
    loss_type="grpo",
    beta=0.0,  # GSPO set KL regularization to zero: https://github.com/volcengine/verl/pull/2775#issuecomment-3131807306
    epsilon=3e-4,  # GSPO paper (v2), section 5.1
    epsilon_high=4e-4,  # GSPO paper (v2), section 5.1
    gradient_accumulation_steps=1,
    steps_per_generation=4,  # partition rollout batch into 4 mini-batches. GSPO paper (v2), section 5.1. Must be 4 times gradient_accumulation_steps
)

Note that this method only has an effect when training goes slightly off-policy—for example, when steps_per_generation > gradient_accumulation_steps or num_iterations > 1. Otherwise, it is effectively equivalent to no modification.

TRL shipped an experimental GSPO-token implementation up to v1.13; it is no longer part of the library and remains available in the git history. It differed from the sequence-level objective above only when the advantage varies with \( t \), which TRL does not currently produce.

Policy ratio: GRPO vs. GSPO

In GSPO, the policy ratio is defined at the sequence-level. In other words, it is the ratio between the probability of the current policy generating a sequence over the old policy generating that same sequence.

The sequence likelihood is defined as:

$$
\pi_\theta (o_i | q) = \prod_{t=1}^{|o_i|} \pi_\theta (o_{i,t} | q, o_{i, < t} ),
$$

where \( \pi_\theta \) is the policy \( \pi \) with parameters \(\theta\), \( o_i \) is the \( i \)-th output sequence \( o \) and \(o_{i,t}\) is the \( t \)-th token in this sequence, \( q \) is the input query. The sequence likelihood ratio \( s_i (\theta) \) is defined as:

$$
s_i (\theta) = \left(\frac{\pi_\theta (o_i | q)}{\pi_{\theta_{old}} (o_i | q)} \right)^{\frac{1}{|o_i|}}
$$

The exponent \( \frac{1}{|o_i|} \) represents a sequence-length normalization, minimizing the influence of sequence length in sequence likelihood. In other terms, it computes the geometric mean of token probabilities, ensuring a fair comparison across sequences of varying lengths.

While GSPO defines the policy ratio at the sequence level, GRPO operates at the token level. Specifically, GRPO computes an importance ratio for each token in the sequence:

$$
w_{i,t}(\theta) = \frac{\pi_\theta (o_{i,t} | q, o_{i,< t})}{\pi_{\theta_{\text{old}}} (o_{i,t} | q, o_{i,< t})}
$$

This token-level ratio is then combined with a shared advantage \( \hat{A}_i \), and the GRPO objective clips and optimizes each token independently across the sequence.

Geometric-Mean Policy Optimization Copy Section

📜 Paper: https://huggingface.co/papers/2507.20673

Geometric-Mean Policy Optimization (GMPO) is a GRPO variant that maximizes the geometric mean of the token-level importance ratios instead of the arithmetic mean. The geometric mean is far less sensitive to outlier ratios, so the policy update is more stable and tolerates a much wider clipping range. Clipping is applied per token, in log space, and is one-sided per the advantage sign (the standard PPO trust region) — crucially, before the geometric mean is taken. This is what distinguishes GMPO from GSPO (importance_sampling_level="sequence"), which clips the sequence-level ratio after averaging.

The GMPO objective replaces GRPO's per-token arithmetic mean with the geometric mean of the (clipped) token ratios:

$$
\mathcal{J}{\text{GMPO}}(\theta) = \mathbb{E}{q, {o_i}} \left[ \frac{1}{G} \sum_{i=1}^{G} \left( \prod_{t=1}^{|o_i|} \min\left[ w_{i,t}(\theta)^{\operatorname{sgn}(\hat{A}i)},\ \operatorname{clip}\left(w{i,t}(\theta)^{\operatorname{sgn}(\hat{A}_i)}, \epsilon_1, \epsilon_2\right) \right]^{\operatorname{sgn}(\hat{A}_i)} \right)^{\frac{1}{|o_i|}} \hat{A}_i \right]
$$

where \( w_{i,t}(\theta) \) is the per-token importance ratio. In practice the product and clipping are computed in log space for numerical stability, and the clip range \( (\epsilon_1, \epsilon_2) = (e^{-0.4}, e^{0.4}) \) is markedly wider than GRPO/DAPO to encourage exploration.

TRL provides an experimental implementation, see Experimental - GMPO:

from trl.experimental.gmpo import GMPOConfig, GMPOTrainer

training_args = GMPOConfig(
    epsilon=0.4,  # log-space clip range -> ratios clipped to (exp(-0.4), exp(0.4)); paper, Sec. 4
    beta=0.0,
)

DAPO: An Open-Source LLM Reinforcement Learning System at Scale Copy Section

📜 Paper: https://huggingface.co/papers/2503.14476

The DAPO algorithm includes 5 key components:

• Overlong Filtering
• Clip-Higher
• Soft Overlong Punishment
• Token-level Loss
• Dynamic Sampling (⚠️ Not supported in TRL)

To reproduce the paper's setting, use this configuration:

from trl import GRPOConfig, GRPOTrainer

training_args = GRPOConfig(
    # Overlong Filtering
    mask_truncated_completions=True,
    # Token-level Loss
    loss_type="dapo",
    # Clip-Higher
    epsilon_high=0.28, # DAPO paper: section 4.1
    epsilon=0.2, # DAPO paper: section 4.1
    # Other parameters used
    per_device_train_batch_size=512, # mini-batch size for training in the paper, DAPO paper: section 4.1
    num_generations=16, # number of sample responses in the paper, DAPO paper: section 4.1
    max_completion_length=20480, #  maximum number of tokens for generation in the paper, DAPO paper: section 4.1
    beta=0.0, # section 2.3, DAPO paper

)
# Soft Overlong Punishment
sop_reward = get_soft_overlong_punishment(max_completion_len=20480, soft_punish_cache=4096) # DAPO paper: section 4.1
trainer = GRPOTrainer(
    ...,
    args=training_args,
    reward_funcs=[..., sop_reward],
)

Demystifying Long Chain-of-Thought Reasoning in LLMs Copy Section

📜 Paper: https://huggingface.co/papers/2502.03373

This paper studies long chain-of-thought RL and introduces two complementary rule-based rewards:

• A cosine length-scaled reward ([~rewards.get_cosine_scaled_reward], Appendix C.1) that scales the correctness reward by completion length: a correct completion is rewarded more when it is shorter, while a wrong completion is penalized less when it is longer (preserving exploration).
• An n-gram repetition penalty ([~rewards.get_repetition_penalty_reward], Appendix C.2) that discourages the degenerate, repetitive completions that emerge as a reward-hacking strategy under length shaping.

The two are designed to be used together:

from trl import GRPOTrainer
from trl.rewards import get_cosine_scaled_reward, get_repetition_penalty_reward

# max_len should match the generation budget (in tokens)
cosine_scaled_reward = get_cosine_scaled_reward(max_len=4096)
# Penalize repeated 3-grams, down to -1.0 for fully repetitive completions
repetition_penalty = get_repetition_penalty_reward(ngram_size=3, max_penalty=-1.0)
trainer = GRPOTrainer(
    ...,
    reward_funcs=[cosine_scaled_reward, repetition_penalty],
)

INTELLECT-2: A Reasoning Model Trained Through Globally Decentralized Reinforcement Learning Copy Section

📜 Paper: https://huggingface.co/papers/2505.07291

INTELLECT-2 is the first globally distributed reinforcement learning training run of a 32 billion parameter language model using fully asynchronous RL across a dynamic, heterogeneous swarm of permissionless compute contributors. The authors propose modifications to the standard GRPO training recipe, including two-sided GRPO clipping for increased training stability. To reproduce the paper's setting, use this configuration:

from trl import GRPOConfig

training_args = GRPOConfig(
    delta=4,  # δ in section 4.1 of the paper
    epsilon=0.2,  # ε in section 4.1 of the paper
    beta=0.001,  # KL divergence coefficient in section 4.1 of the paper
    num_generations=16,  # responses per prompt in section 4.1 of the paper
    learning_rate=3e-7,  # section 4.1 of the paper
)

Skywork-OR1: Open Reasoning Models Copy Section

📜 Paper: https://huggingface.co/papers/2505.22312

Skywork-OR1 is a family of open reasoning models trained with GRPO. The paper introduces adaptive entropy control: an entropy regularization term −α·H(π_θ) is added to the GRPO objective, and the coefficient α is automatically adjusted each optimizer step. When the model's mean per-token entropy falls at or below a target, α is incremented to encourage more exploration; otherwise it is decremented. The bonus is only applied while entropy is at or below the target. To replicate this adaptive entropy control, use the following configuration:

from trl import GRPOConfig, GRPOTrainer

training_args = GRPOConfig(
    use_adaptive_entropy=True,   # enable adaptive entropy control (Section 3.3 of the paper)
    entropy_coef=0.01,           # initial entropy regularization coefficient
    entropy_target=5.0,          # target mean per-token entropy (nats); tune for your model
    entropy_coef_delta=0.005,    # step size for coefficient updates per optimizer step
)
trainer = GRPOTrainer(
    ...,
    args=training_args,
)

Beyond the 80/20 Rule: High-Entropy Minority Tokens Drive Effective Reinforcement Learning for LLM Reasoning Copy Section

📜 Paper: https://huggingface.co/papers/2506.01939

A minority of tokens with high entropy act as reasoning "forks" in the CoT path, driving exploration and performance gains for RLVR, while low-entropy majority tokens contribute little or even impede learning. RLVR mainly adjusts high-entropy tokens, largely preserving the base model’s overall entropy patterns. Thus landing on the 80/20 rule, training on only 20% of the tokens with the highest entropy is comparable or supasses full-gradient updates for Qwen3 models.

The paper's main results use vanilla DAPO (⚠️ Dynamic Sampling is not supported in TRL). To replicate the main results, use the following configuration:

from trl import GRPOConfig, GRPOTrainer
from trl.rewards import get_soft_overlong_punishment

training_args = GRPOConfig(
    # --- vanilla DAPO parameters (80/20 rule: section 5.2) --- #
    # Overlong Filtering
    mask_truncated_completions=True,
    # Token-level Loss
    loss_type="dapo",
    # Clip-Higher
    epsilon_high=0.28, # DAPO paper: section 4.1
    epsilon=0.2, # DAPO paper: section 4.1
    # Other parameters used
    per_device_train_batch_size=512, # mini-batch size for training in the paper, DAPO paper: section 4.1
    num_generations=16, # number of sample responses in the paper, DAPO paper: section 4.1
    max_completion_length=20480, #  maximum number of tokens for generation in the paper, DAPO paper: section 4.1
    beta=0.0, # section 2.3, DAPO paper
    # --- Gradients on the highest entropy tokens --- #
    top_entropy_quantile=0.2
)
# Soft Overlong Punishment
sop_reward = get_soft_overlong_punishment(max_completion_len=20480, soft_punish_cache=4096) # DAPO paper: section 4.1
trainer = GRPOTrainer(
    ...,
    args=training_args,
    reward_funcs=[..., sop_reward],
)

Dr. GRPO: Understanding R1-Zero-Like Training: A Critical Perspective Copy Section

📜 Paper: https://huggingface.co/papers/2503.20783

A study of R1-Zero training identifies pretraining effects on RL performance and proffers Dr. GRPO to enhance token efficiency, achieving superior accuracy on AIME 2024. To reproduce the paper's setting, use this configuration:

from trl import GRPOConfig

training_args = GRPOConfig(
    loss_type="dr_grpo",
    per_device_train_batch_size=1, # train_batch_size_per_device in the Training section of the repository
    num_generations=8, #  num_samples in the Training section of the repository
    max_completion_length=3000, # generate_max_length in the Training section of the repository
    beta=0.0, # β in the Training section of the repository
)

Part I: Tricks or Traps? A Deep Dive into RL for LLM Reasoning (Lite PPO) Copy Section

📜 Paper: https://huggingface.co/papers/2508.08221

The authors of this paper find that the combination of:

1. scaling rewards by the standard deviation computed over the entire batch and
2. aggregating loss over the total number of tokens

can unlock the learning capability of critic-free policies using vanilla PPO loss. Their results demonstrate that this simple combination consistently improves performance, surpassing strategies like GRPO and DAPO.

TRL supports using these learnings to train a GRPO model by:

from trl import GRPOConfig

training_args = GRPOConfig(
    ...
    scale_rewards="batch",
    loss_type="dapo",
    # Other parameters used
    beta=0.0,  # = init_kl_coef in the paper
    top_p=0.99,
    top_k=100,
    temperature=0.99,
    num_generations=8, # = num_return_sequences in the paper
    num_iterations=1,  # = ppo_epochs in the paper
    per_device_train_batch_size=4,
    gradient_accumulation_steps=32,
    steps_per_generation=8,  # (rollout_batch_size*num_return_sequences) / (per_device_train_batch_size*gradient_accumulation_steps)
)

Note that when using gradient accumulation, the loss is aggregated over the total number of tokens in the batch, but not over the accumulated batch. For more details, see the GRPO Trainer - Loss types.

Truncated Importance Sampling Copy Section

📰 Blog: https://fengyao.notion.site/off-policy-rl

📜 Paper: https://huggingface.co/papers/1606.02647

Online policy learning methods commonly use an optimized inference framework for rollout generation (e.g vLLM) that is separate from the training backend. This introduces a rollout-training mismatch, exemplified in the following PPO objective:

$$
\small{
\mathbb{E}{a\sim\textcolor{red}{\pi{\text{inference}}}(\theta_{\mathrm{old}})}
\Bigl[
\min\Bigl(
\frac{\textcolor{blue}{\pi_{\text{training}}}(a, \theta)}{\textcolor{blue}{\pi_{\text{training}}}(a, \theta_{\mathrm{old}})},\hat A,
;\mathrm{clip}\bigl(\frac{\textcolor{blue}{\pi_{\text{training}}}(a, \theta)}{\textcolor{blue}{\pi_{\text{training}}}(a, \theta_{\mathrm{old}})},,1-\epsilon,,1+\epsilon\bigr),\hat A
\Bigr)
\Bigr]
}
$$

Despite \( \textcolor{red}{\pi_{\text{inference}}} \) and \( \textcolor{blue}{\pi_{\text{training}}} \) sharing the same model parameters \( \theta \), they can produce significantly different token probabilities. This unexpected behavior implicitly breaks the on-policy assumption, and silently turns training off-policy.

Truncated Importance Sampling (TIS) addresses this issue by adapting the model update via importance-sampling correction. The gradient computation of the aforementioned PPO objective becomes

$$
\small{
\mathbb{E}{a\sim\textcolor{red}{\pi{\text{inference}}}(\theta_{\mathrm{old}})}
\Bigl[
\underbrace{\text{clip}\bigl(\frac{\textcolor{blue}{\pi_{\text{training}}}(a, \theta_{\mathrm{old}})}{\textcolor{red}{\pi_{\text{inference}}}(a, \theta_{\mathrm{old}})}, C_{\min}, C_{\max}\bigr)}{\text{truncated importance ratio}} \cdot
\nabla\theta
\min\Bigl(
\frac{\textcolor{blue}{\pi_{\text{training}}}(a, \theta)}{\textcolor{blue}{\pi_{\text{training}}}(a, \theta_{\mathrm{old}})},\hat A,
;\mathrm{clip}\bigl(\frac{\textcolor{blue}{\pi_{\text{training}}}(a, \theta)}{\textcolor{blue}{\pi_{\text{training}}}(a, \theta_{\mathrm{old}})},,1-\epsilon,,1+\epsilon\bigr),\hat A
\Bigr)
\Bigr]
}
$$

where \( C_{\min} \) and \( C_{\max} \) are hyper-parameters. TIS is implemented in GRPO, and is enabled by selecting a vllm_importance_sampling_mode variant that includes the term truncate, such as "sequence_truncate" or "token_truncate".

from trl import GRPOConfig

training_args = GRPOConfig(
    ...
    use_vllm=True,
    vllm_importance_sampling_correction=True, # default True
    vllm_importance_sampling_mode="sequence_truncate", # or "token_truncate"
    vllm_importance_sampling_clip_max=2.0, # hyper-parameter C_max
)

Masked Importance Sampling Copy Section

📰 Blog: https://ringtech.notion.site/icepop

📜 Paper: https://huggingface.co/papers/2510.18855

📰 Blog: https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda

Masked Importance Sampling (MIS) addresses the same issue as Truncated Importance Sampling but replaces clipping with masking. MIS takes a more decisive stance by discarding updates whose discrepancy falls outside the range \( [C_{\min}, C_{\max}] \).

$$
\small{
\mathbb{E}{a\sim\textcolor{red}{\pi{\text{inference}}}(\theta_{\mathrm{old}})}
\Bigl[
\underbrace{\mathbf{1}\left[
C_{\min} \le \frac{\pi_{\text{training}}(a, \theta_{\mathrm{old}})}
{\pi_{\text{inference}}(a, \theta_{\mathrm{old}})}
\le C_{\max}
\right]
\cdot
\frac{\pi_{\text{training}}(a, \theta_{\mathrm{old}})}
{\pi_{\text{inference}}(a, \theta_{\mathrm{old}})}}{\text{masked importance ratio}} \cdot
\nabla\theta
\min\Bigl(
\frac{\textcolor{blue}{\pi_{\text{training}}}(a, \theta)}{\textcolor{blue}{\pi_{\text{training}}}(a, \theta_{\mathrm{old}})},\hat A,
;\mathrm{clip}\bigl(\frac{\textcolor{blue}{\pi_{\text{training}}}(a, \theta)}{\textcolor{blue}{\pi_{\text{training}}}(a, \theta_{\mathrm{old}})},,1-\epsilon,,1+\epsilon\bigr),\hat A
\Bigr)
\Bigr]
}
$$

MIS is implemented for GRPO, and is enabled by selecting a vllm_importance_sampling_mode variant that includes the term "mask", such as "sequence_mask" or "token_mask".

from trl import GRPOConfig

training_args = GRPOConfig(
    ...
    use_vllm=True,
    vllm_importance_sampling_correction=True, # default True
    vllm_importance_sampling_mode="sequence_mask", # or "token_mask"
    vllm_importance_sampling_clip_max=2.0, # hyper-parameter C_max
    vllm_importance_sampling_clip_min=0.5, # hyper-parameter C_min
)

Sequence-level Importance Sampling Copy Section

📰 Blog: https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda

The theoretically principled way to correct for the training-inference distribution shift is importance sampling, as introduced in the two papers above Truncated Importance Sampling and Masked Importance Sampling. However, the choice of formulation is crucial for keeping the gradient unbiased and ensuring stable training.

This work shows that sequence-level importance sampling is the sound approach for addressing the training–inference mismatch. Although token-level importance sampling achieves lower variance than a sequence-level ratio, it introduces bias and is therefore argued to be unsuitable for autoregressive models. The token-level gradient estimator is

$$
\mathbb{E}{x\sim\mathcal{D},, y\sim \pi^{\text{inference}}\theta(\cdot|x)}
\Bigg[
R(x,y),\cdot,
\sum_{t=0}^{|y|-1}
\frac{\pi^{\text{training}}\theta(y_t,|,x, y{<t})}
{\pi^{\text{inference}}\theta(y_t,|,x, y{<t})}
,\nabla_\theta \log \pi^{\text{training}}\theta(y_t,|,x, y{<t})
\Bigg]
$$
The correct, unbiased policy gradient estimator applies a single importance ratio over the entire generated sequence (trajectory) \( y \), The Sequence-Level IS estimator looks like:

$$
\mathbb{E}{x\sim\mathcal{D},, y\sim \pi^{\text{inference}}\theta(\cdot|x)}
\Bigg[
\frac{\pi^{\text{training}}\theta(y|x)}
{\pi^{\text{inference}}\theta(y|x)}
, R(x,y),
\nabla_\theta \log \pi^{\text{training}}_\theta(y|x)
\Bigg]
$$

TRL exposes the Importance Sampling granularity level through the vllm_importance_sampling_mode configuration parameter where "sequence_*" modes implement a sequence-level importance sampling ratio and "token_*" a per-token ratio.

The Art of Scaling Reinforcement Learning Copy Section

📜 Paper: https://huggingface.co/papers/2510.13786

A systematic study that defines a framework for analyzing and predicting reinforcement learning scaling in large language models, identifies key design choices that affect compute efficiency and propose a best-practice recipe called ScaleRL.

You can partially reproduce the ScaleRL recipe using the [GRPOTrainer] with the following configs:

from trl import GRPOConfig

training_args = GRPOConfig(
    loss_type="cispo",
    epsilon_high=5.0,
    num_generations=16,
    scale_rewards="batch",
    cast_lm_head_to_fp32=True
)

It Takes Two: Your GRPO Is Secretly DPO Copy Section

📜 Paper: https://huggingface.co/papers/2510.00977

Shows that GRPO's effectiveness stems from an implicit contrastive objective rather than accurate advantage estimation via large group sizes. This establishes a formal connection between GRPO and DPO, where group size only affects Monte Carlo estimators of the contrastive objective. The authors introduce 2-GRPO — using just two rollouts — which matches the performance of 16-GRPO at significantly lower training cost. Used in TRL via [GRPOTrainer] with num_generations=2. To reproduce the paper's setting, use this configuration:

from trl import GRPOConfig, GRPOTrainer

training_args = GRPOConfig(
    num_generations=2,  # minimal two-rollout case (2-GRPO) from the paper
    loss_type="grpo",
)
trainer = GRPOTrainer(
    ...,
    args=training_args,
)

Soft Adaptive Policy Optimization Copy Section

📜 Paper: https://huggingface.co/papers/2511.20347

Soft Adaptive Policy Optimization (SAPO), replaces hard clipping with a smooth, temperature-controlled gate that adaptively attenuates off-policy updates while preserving useful learning signals. Compared with GSPO and GRPO, SAPO is both sequence-coherent and token-adaptive. Like GSPO, SAPO maintains sequence-level coherence, but its soft gating forms a continuous trust region that avoids the brittle hard clipping band used in GSPO.

To reproduce the paper's setting, use this configuration:

from trl import GRPOConfig

training_args = GRPOConfig(
    loss_type="sapo",
    sapo_temperature_pos=1.0,  # default value
    sapo_temperature_neg=1.05,  # default value
    scale_rewards="group",
    ...
)

DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models Copy Section

📜 Paper: https://huggingface.co/papers/2512.02556

DeepSeek-V3.2 technical report introduces several techniques to enhance the performance of GRPO. In TRL we implement:

• The Unbiased KL Estimate, which corrects the K3 estimator (as used in the original GRPO implementation) to obtain an unbiased KL estimate using the importance-sampling ratio between the current policy \( \pi_\theta \) and the behavior policy \( \pi_{\text{old}} \).

$$ \mathrm{D}_{\mathrm{KL}}\!\left(\pi_\theta(o_{i,t}) \,\|\, \pi_{\text{ref}}(o_{i,t})\right) = \textcolor{red}{\frac{\pi_\theta(o_{i,t}\mid q, o_{i, Copy Chapter

\log \frac{\pi_{\text{ref}}(o_{i,t}\mid q, o_{i,<t})}{\pi_\theta(o_{i,t}\mid q, o_{i,<t})}

• 1
  \right).
  $$

This feature is enabled by default (use_bias_correction_kl=True in the [GRPOConfig]), so you only need to set beta > 0:

from trl import GRPOConfig

training_args = GRPOConfig(
    ...,
    beta=0.001,  # the paper doesn't specify the value used, so we use the value from "DeepSeek-R1 incentivizes reasoning in LLMs through reinforcement learning"
)

• The Off-Policy Masking, which stabilizes training by ignoring sequences where the policy performs poorly (negative advantage) and has drifted significantly from the old policy (high KL divergence).

The off-policy binary mask \(\textcolor{red}{M_{i,t}}\) is defined as:

$$
\textcolor{red}{M_{i,t}} = \begin{cases}
0 & \text{if } \hat{A}{i,t} < 0 \quad \text{and} \quad \frac{1}{|o_i|} \sum{t=1}^{|o_i|} \log \frac{\pi_{\theta_{\text{old}}}(o_{i,t} \mid q, o_{i,<t})}{\pi_\theta(o_{i,t} \mid q, o_{i,<t})} > \textcolor{blue}{\delta} \
1 & \text{otherwise}
\end{cases}
$$

This mask is then applied to the GRPO loss as follows:

$$
\mathcal{L}{\text{GRPO}}(\theta) = -\frac{1}{G} \sum{i=1}^G \frac{1}{|o_i|} \sum_{t=1}^{|o_i|} \left[ \min \left( \frac{\pi_\theta(o_{i,t} \mid q, o_{i,< t})}{\pi_{\theta_{\text{old}}}(o_{i,t} \mid q, o_{i,< t})} \hat{A}{i,t}, , \text{clip}\left( \frac{\pi\theta(o_{i,t} \mid q, o_{i,< t})}{\pi_{\theta_{\text{old}}}(o_{i,t} \mid q, o_{i,< t})}, 1 - \epsilon, 1 + \epsilon \right) \hat{A}{i,t} \right) \textcolor{red}{M{i,t}} - \beta \mathbb{D}{\text{KL}}\left[\pi\theta | \pi_{\text{ref}}\right] \right]
$$

To enable this feature, use the off_policy_mask_threshold (corresponding to \( \textcolor{blue}{\delta} \)) in the [GRPOConfig]:

from trl import GRPOConfig

training_args = GRPOConfig(
    ...,
    off_policy_mask_threshold=0.5,
)

While the paper doesn't specify a \( \textcolor{blue}{\delta} \) value used, a good starting point could be \( \textcolor{blue}{\delta} = 0.5 \). If training seems too conservative or too many sequences are masked, you can increase the value.
For reference, \( \textcolor{blue}{\delta} = 1.0 \) corresponds to an average log-ratio divergence of 1 nat per token, i.e. on sequences where this threshold is exceeded, the old policy was on average \( e^1 \approx 2.7 \) times more likely to generate these tokens than the current policy.

• The Keep Sampling Mask strategy, which keeps top-k/top-p sampling unbiased. The sampler draws each token from a truncated support \( \textcolor{red}{\mathcal{S}_{i,t}} \), the tokens left after truncation, so the policy log-probability is renormalized over that same support:

$$
\log \pi_\theta(o_{i,t} \mid q, o_{i,<t}) = z_{o_{i,t}} - \log \sum_{v \in \textcolor{red}{\mathcal{S}_{i,t}}} e^{z_v}.
$$

Without it the importance ratio is biased low and the gradient keeps pushing probability mass onto tokens the sampler never draws, which collapses training. This is supported in the [experimental.async_grpo.AsyncGRPOTrainer]: with top_p < 1.0 or top_k > 0, the rollout worker records each token's support from vLLM (>= 0.28, started with --return-sampling-mask --enable-scale-out) and the trainer renormalizes over it:

from trl.experimental.async_grpo import AsyncGRPOConfig

training_args = AsyncGRPOConfig(
    ...,
    top_p=0.95,
    top_k=512,
)

GDPO: Group reward-Decoupled Normalization Policy Optimization for Multi-reward RL Optimization Copy Section

📜 Paper: https://huggingface.co/papers/2601.05242

GDPO is a reinforcement learning optimization method designed for multi-reward training. While existing approaches commonly apply Group Relative Policy Optimization (GRPO) in multi-reward settings, the authors show that this leads to reward advantages collapse, reducing training signal resolution and causing unstable or failed convergence. GDPO resolves this issue by decoupling reward normalization across individual rewards, preserving their relative differences and enabling more faithful preference optimization. To enable GDPO for multi-reward RL training, simply set:

For a group of \( N \) rewards and \( G \) samples per group, GDPO normalizes each reward independently:

$$
A_n^{(i,j)} = \frac{r_n^{(i,j)} - \text{mean}{r_n^{(i,1)}, \ldots, r_n^{(i,G)}}}{\text{std}{r_n^{(i,1)}, \ldots, r_n^{(i,G)}} + \epsilon}
$$

The normalized group advantage is then aggregated across rewards:

$$
A^{(i,j)} = \sum_{n=1}^{N} w_n A_n^{(i,j)}
$$

The final per-batch normalization produces:

$$
\hat{A}^{(i,j)} = \frac{A^{(i,j)} - \text{mean}{i',j'}{A^{(i',j')}}}{\text{std}{i',j'}{A^{(i',j')}} + \epsilon}
$$

Here, \( \text{mean}{i',j'}{A^{(i',j')}} \) and \( \text{std}{i',j'}{A^{(i',j')}} \) denote statistics over all groups in the batch.

from trl import GRPOConfig


training_args = GRPOConfig(
    ...,
    multi_objective_aggregation="normalize_then_sum",
)

Note that this method only has an effect when training involve more than one reward function.

The authors provide a easy-to-use, slurm-free training example that enable the community to quickly validate GDPO’s effectiveness over GRPO, see Experiment-"Aha" moment.

Length-Unbiased Sequence Policy Optimization: Revealing and Controlling Response Length Variation in RLVR Copy Section

📜 Paper: https://huggingface.co/papers/2602.05261

Length-Unbiased Sequence Policy Optimization (LUSPO) modifies GSPO by scaling each sequence's loss by its length. This corrects GSPO's gradient bias that penalizes longer responses. To reproduce the paper's setting, use this configuration:

from trl import GRPOConfig

training_args = GRPOConfig(
    loss_type="luspo",
    importance_sampling_level="sequence",
    epsilon=2e-3, # section 5.1 of the paper
    epsilon_high=2.5e-3, # section 5.1 of the paper
)

VESPO: Variational Sequence-Level Soft Policy Optimization for Stable Off-Policy LLM Training Copy Section

📜 Paper: https://huggingface.co/papers/2602.10693

VESPO addresses training instability in off-policy RL caused by policy staleness, asynchronous updates, and train-inference mismatches. Rather than relying on heuristic token-level clipping (GRPO) or sequence-length normalization (GSPO), VESPO derives a principled reshaping kernel from a variational framework. In practice, this yields a smooth, asymmetric Gamma weighting function that gracefully suppresses extreme sequence-level importance weights without introducing length bias.

$$
\mathcal{L}{\text{VESPO}}(\theta) = - \mathbb{E}{\tau \sim \mu} \left[ \underbrace{W(\tau)^{k} \cdot \exp\left(\lambda
(1 - W(\tau))\right)}{\phi(W) \text{ detached }} \cdot \mathcal{A}(\tau) \cdot \log \pi\theta(\tau) \right]
$$

with \( W(\tau) = \frac{\pi_\theta(\tau)}{\mu(\tau)} \) the sequence level importance ratio, and \( \phi(W) \) is detached from the computation graph to serve as a gradient scaling coefficient.

from trl import GRPOConfig

training_args = GRPOConfig(
    loss_type="vespo",
    use_vllm=True,  # or False if not using any token-level `vllm_importance_sampling_correction` methods
    vllm_importance_sampling_mode="token_truncate",  # default correction mode for VESPO, `token_mask` also supported
    vespo_k_pos=2.0,  # power exponent (c1 in paper Section 3.4) for positive advantages
    vespo_lambda_pos=3.0,  # decay factor (c2 in paper Section 3.4) for positive advantages
    vespo_k_neg=3.0,  # power exponent (c1 in paper Section 3.4) for negative advantages
    vespo_lambda_neg=2.0,  # decay factor (c2 in paper Section 3.4) for negative advantages
)

Optimal Advantage Regression Copy Chapter

Papers relating to the [experimental.a2po.A2POTrainer].

Accelerating RL for LLM Reasoning with Optimal Advantage Regression Copy Section

📜 Paper: https://huggingface.co/papers/2505.20686

A*-PO (Optimal Advantage Regression) is a two-stage RL method for LLM reasoning that avoids both an online critic and multi-sample group rollouts. Stage 1 estimates the optimal value V*(x) = β·log E_{y∼π_ref}[exp(r(x, y)/β)] offline from reference-policy samples; Stage 2 performs on-policy updates with a single generation per prompt using a squared-error regression loss (β·log(π(y|x)/π_ref(y|x)) − (r(x, y) − V*(x)))². This yields faster training and lower peak memory than PPO/GRPO/REBEL. The method assumes a binary verifiable reward and cannot exceed the reference policy's Pass@K. See A2PO. The official code can be found in ZhaolinGao/A-PO.

from trl.experimental.a2po import A2POConfig, A2POTrainer

# Hyperparameters from the paper (Qwen2.5, GSM8K/MATH)
training_args = A2POConfig(
    beta1=0.5,  # KL temperature for offline V* estimation
    beta2=1e-3,  # KL temperature for the on-policy regression target
    num_value_samples=8,  # N samples per prompt from the reference policy
)
trainer = A2POTrainer(
    model=...,
    reward_funcs=...,  # should return a binary reward in {0, 1}
    args=training_args,
    train_dataset=...,
)
trainer.train()

Direct Policy Optimization Copy Chapter

Papers relating to the [DPOTrainer]

Direct Preference Optimization: Your Language Model is Secretly a Reward Model Copy Section

📜 Paper: https://huggingface.co/papers/2305.18290

Direct Preference Optimization (DPO) fine-tunes language models more efficiently and with better performance compared to reinforcement learning from human feedback (RLHF), by directly optimizing policy training based on human preferences. To reproduce the paper's setting, use this configuration:

from trl import DPOConfig

training_args = DPOConfig(
    loss_type="sigmoid", # losses in Appendix B of the paper
    per_device_train_batch_size=64, #  batch size in Appendix B of the paper
    learning_rate=1e-6, # learning rate in Appendix B of the paper
    beta=0.1, # β in Appendix B of the paper
)

SLiC-HF: Sequence Likelihood Calibration with Human Feedback Copy Section

📜 Paper: https://huggingface.co/papers/2305.10425

Sequence Likelihood Calibration (SLiC) is shown to be an effective and simpler alternative to Reinforcement Learning from Human Feedback (RLHF) for learning from human preferences in language models. To reproduce the paper's setting, use this configuration:

from trl import DPOConfig

training_args = DPOConfig(
    loss_type="hinge", # Section 2 of the paper
    per_device_train_batch_size=512, #  batch size in Section 3.2 of the paper
    learning_rate=1e-4, # learning rate in Section 3.2 of the paper
)

These parameters only appear in the published version

Statistical Rejection Sampling Improves Preference Optimization Copy Section

📜 Paper: https://huggingface.co/papers/2309.06657

Proposes RSO, selecting stronger preference pairs via statistical rejection sampling to boost offline preference optimization; complements DPO/SLiC. They also introduce a new loss defined as:

$$
\mathcal{L}{\text{hinge-norm}}(\pi\theta)
= \mathbb{E}{(x, y_w, y_l) \sim \mathcal{D}}
\left[
\max\left(0,; 1 - \left[\gamma \log \frac{\pi\theta(y_w \mid x)}{\pi_\text{ref}(y_w \mid x)} - \gamma \log \frac{\pi_\theta(y_l \mid x)}{\pi_\text{ref}(y_l \mid x)}\right]\right)
\right]
$$

To train with RSO-filtered data and the hinge-norm loss, you can use the following code:

from trl import DPOConfig, DPOTrainer

dataset = ...

def rso_accept(example):  # replace with your actual filter/score logic
    return example["rso_keep"]

train_dataset = train_dataset.filter(rso_accept)

training_args = DPOConfig(
    loss_type="hinge",
    beta=0.05,  # correspond to γ in the paper
)

trainer = DPOTrainer(
    ...,
    args=training_args,
    train_dataset=train_dataset,
)
trainer.train()

Beyond Reverse KL: Generalizing Direct Preference Optimization with Diverse Divergence Constraints Copy Section

📜 Paper: https://huggingface.co/papers/2309.16240

Proposes (( f \)-DPO, extending DPO by replacing the usual reverse-KL regularizer with a general (( f \)-divergence, letting you trade off mode-seeking vs mass-covering behavior (e.g. forward KL, JS, (( \alpha \)-divergences). The only change is replacing the DPO log-ratio margin with an f′ score:

$$ \mathcal{L}_{f\text{-DPO}}(\pi_\theta) = \mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}} \left[ -\log \sigma\left( \beta \textcolor{red}{f'}\textcolor{red}{\Big(}\frac{\pi_\theta(y_w|x)}{\pi_{\text{ref}}(y_w|x)}\textcolor{red}{\Big)} Copy Chapter

\beta \textcolor{red}{f'}\textcolor{red}{\Big(}\frac{\pi_\theta(y_l|x)}{\pi_{\text{ref}}(y_l|x)}\textcolor{red}{\Big)}
\right)
\right]
$$

Where \( f' \) is the derivative of the convex function defining the chosen (( f \)-divergence.

To reproduce:

from trl import DPOConfig

training_args = DPOConfig(
    loss_type="sigmoid",
    beta=0.1,
    f_divergence_type="js_divergence",  # or "reverse_kl" (default), "forward_kl", "js_divergence", "alpha_divergence"
    f_alpha_divergence_coef=0.5,  # only used if f_divergence_type="alpha_divergence"
)

A General Theoretical Paradigm to Understand Learning from Human Preferences Copy Section

📜 Paper: https://huggingface.co/papers/2310.12036

Learning from human preferences can be written as a single KL-regularized objective over pairwise preference probabilities,

$$
\max_\pi ;\mathbb{E}\big[\Psi\left(p^*(y \succ y' \mid x)\right)\big] - \tau\mathrm{KL}(\pi||\pi_{\text{ref}}),
$$

which reveals RLHF and DPO as special cases corresponding to the logit choice of \( \Psi \).
The paper shows that this logit transform amplifies near-deterministic preferences and effectively weakens KL regularization, explaining overfitting.
Using the Identity transform (IPO) avoids this pathology by optimizing preferences directly, without assuming a Bradley–Terry reward model.
To reproduce the paper's setting, use this configuration:

from trl import DPOConfig

training_args = DPOConfig(
    loss_type="ipo",  # Section 5.1 of the paper
    per_device_train_batch_size=90,  #  mini-batch size in Section C.1 of the paper
    learning_rate=1e-2,  # learning rate in Section C.1 of the paper
)

These parameters only appear in the published version

Towards Efficient and Exact Optimization of Language Model Alignment Copy Section

📜 Paper: https://huggingface.co/papers/2402.00856

The paper shows that direct preference methods like DPO optimize the wrong KL direction, leading to blurred preference capture, and proposes EXO as an efficient way to exactly optimize the human‑preference alignment objective by leveraging reverse KL probability matching rather than forward KL approximations. To reproduce the paper's setting, use this configuration:

from trl import DPOConfig

training_args = DPOConfig(
    loss_type="exo_pair", # Section 3.2 of the paper
    # From Section B of the paper
    per_device_train_batch_size=64,
    learning_rate=1e-6,
    beta=0.1,
)

Noise Contrastive Alignment of Language Models with Explicit Rewards Copy Section

📜 Paper: https://huggingface.co/papers/2402.05369

The paper reframes language-model alignment as a noise-contrastive classification problem, proposing InfoNCA to learn a policy from explicit rewards (or preferences) by matching a reward-induced target distribution over responses, and showing DPO is a special binary case. It then introduces NCA, which adds an absolute likelihood term to prevent the likelihood collapse seen in purely relative (contrastive) objectives.

With pairwise preferences, treat the chosen/rejected \( K=2 \), define scores \( r=\beta(\log\pi_\theta-\log\pi_{\text{ref}}) \), and apply the NCA preference loss \( -\log\sigma(r_w)-\tfrac12\log\sigma(-r_w)-\tfrac12\log\sigma(-r_l) \).

To reproduce the paper's setting, use this configuration:

from trl import DPOConfig

training_args = DPOConfig(
    loss_type="nca_pair",
    # From Section C of the paper
    per_device_train_batch_size=32,
    learning_rate=5e-6,
    beta=0.01,
)

Provably Robust DPO: Aligning Language Models with Noisy Feedback Copy Section

📜 Paper: https://huggingface.co/papers/2403.00409

DPO breaks under noisy human preferences because label flips bias the objective. Robust DPO fixes this by analytically debiasing the DPO loss under a simple noise model, with provable guarantees.

$$
\mathcal{L}{\text{robust}}(\pi\theta) = \frac{(1-\varepsilon)\mathcal{L}{\text{DPO}}(y_w, y_l) - \varepsilon\mathcal{L}{\text{DPO}}(y_l, y_w)}
{1-2\varepsilon}
$$

Where \( \mathcal{L}_{\text{DPO}} \) is the DPO loss defined in Direct Preference Optimization: Your Language Model is Secretly a Reward Model and \( \varepsilon \) is the probability of a label flip.

This single correction turns noisy preference data into an unbiased estimator of the clean DPO objective.

from trl import DPOConfig

training_args = DPOConfig(
    loss_type="robust",
    per_device_train_batch_size=16,  # batch size in Section B of the paper
    learning_rate=1e-3,  # learning rate in Section B of the paper
    beta=0.1,  # β in Section B of the paper,
    max_length=512,  # max length in Section B of the paper
    label_smoothing=0.1  # label smoothing $\varepsilon$ in Section 6 of the paper
)

Binary Classifier Optimization for Large Language Model Alignment Copy Section

📜 Paper: https://huggingface.co/papers/2404.04656

Theoretical analysis and a new algorithm, Binary Classifier Optimization, explain and enhance the alignment of large language models using binary feedback signals. To reproduce the paper's setting, use this configuration:

BCO reframes language-model alignment as behavioral cloning from an optimal reward-weighted distribution, yielding simple supervised objectives that avoid RL while remaining theoretically grounded.
It supports both unpaired reward data and pairwise preference data, with a reward-shift–invariant formulation that reduces to a DPO-style loss in the preference setting.

For the pairwise preference setting, the BCO loss is defined as:

$$ \mathcal{L}_{\text{bco\_pair}}(\pi_\theta) = \mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}} \left[ -\log \sigma\Big( \beta[(\log\pi_\theta-\log\pi_{\text{ref}})(y_w) Copy Chapter

(\log\pi_\theta-\log\pi_{\text{ref}})(y_l)]
\Big)
\right]
$$

To reproduce the paper in this setting, use this configuration:

from trl import DPOConfig

training_args = DPOConfig(
    loss_type="bco_pair",
    # From Section C of the paper
    per_device_train_batch_size=128,
    learning_rate=5e-7,
    beta=0.01,
)

Learn Your Reference Model for Real Good Alignment Copy Section

📜 Paper: https://huggingface.co/papers/2404.09656

Trust Region DPO (TR-DPO) updates the reference policy during training, demonstrating effectiveness against DPO on the Anthropic HH and TLDR datasets, outperforming DPO by up to 19% measured by automatic evaluation with GPT-4, improving coherence, correctness, level of detail, helpfulness, and harmlessness. To reproduce the paper's setting, use this configuration:

from trl import DPOConfig

training_args = DPOConfig(
    sync_ref_model=True,  # enable TR-DPO (Section 3 of the paper)
    ref_model_mixup_alpha=0.6,  # α soft update weight (Table 1 of the paper)
    ref_model_sync_steps=512,  # τ update frequency in steps (Table 1 of the paper)
    beta=0.05,  # β temperature (Table 1 of the paper)
    learning_rate=1e-6,  # learning rate (Table 2 of the paper)
    num_train_epochs=1,  # Table 2 of the paper
    max_length=1024,  # max tokens length (Table 2 of the paper)
    max_grad_norm=2,  # max gradient norm (Table 2 of the paper)
    warmup_steps=100,  # warm-up steps (Table 2 of the paper)
)

Iterative Reasoning Preference Optimization Copy Section

📜 Paper: https://huggingface.co/papers/2404.19733

Iterative RPO improves reasoning by repeatedly generating chain-of-thought candidates, building preference pairs from correct vs. incorrect answers, and training with a DPO + NLL objective. The extra NLL term is key for learning to actually generate winning traces.

TRL can express the DPO + NLL objective by mixing "sigmoid" (DPO) with "sft" (NLL):

from trl import DPOConfig, DPOTrainer

training_args = DPOConfig(
    loss_type=["sigmoid", "sft"],
    loss_weights=[1.0, 1.0],  # alpha in the paper, recommended value is 1.0
)
trainer = DPOTrainer(
    ...,
    args=training_args,
)

Note that the paper uses an iterative loop: each iteration regenerates CoT candidates with the current model, then retrains on fresh preference pairs. TRL does not automate that loop for you.

Self-Play Preference Optimization for Language Model Alignment Copy Section

📜 Paper: https://huggingface.co/papers/2405.00675

A self-play method called SPPO for language model alignment achieves state-of-the-art performance by approximating Nash equilibrium policy in a constant-sum game setting, outperforming other approaches with limited data. To reproduce the paper's setting, use this configuration:

from trl import DPOConfig

training_args = DPOConfig(
    loss_type="sppo_hard",
    # From Section 5 of the paper
    beta=0.001,  # β = η^−1
    per_device_train_batch_size=64,
    learning_rate=5e-7,
)

Provably Mitigating Overoptimization in RLHF: Your SFT Loss is Implicitly an Adversarial Regularizer Copy Section

📜 Paper: https://huggingface.co/papers/2405.16436

Regularized Preference Optimization (RPO) mitigates overoptimization in RLHF by fusing the DPO loss with the SFT loss, provably preventing the policy from choosing actions with spurious high proxy rewards. To reproduce the paper's setting, use this configuration:

from trl import DPOConfig

training_args = DPOConfig(
    loss_type=["sigmoid", "sft"],  # RPO loss = DPO + SFT (Section 5 of the paper)
    loss_weights=[1.0, 0.005],  # η=0.005 SFT weight in Appendix E.1 of the paper
    beta=0.01,  # β in Appendix E.1 of the paper
    learning_rate=5e-7,  # learning rate in Appendix E.1 of the paper
    num_train_epochs=1,  # Appendix E.1 of the paper
)

Distributional Preference Alignment of LLMs via Optimal Transport Copy Section

📜 Paper: https://huggingface.co/papers/2406.05882

Alignment via Optimal Transport (AOT) aligns large language models distributionally by penalizing violations of stochastic dominance between positive and negative sample distributions, achieving state-of-the-art performance on alignment benchmarks. To reproduce the paper's setting, use this configuration:

from trl import DPOConfig

training_args = DPOConfig(
    loss_type="aot",
    beta=0.01,  # from the caption of Figure 2
)

or, for the unpaired version:

from trl import DPOConfig

training_args = DPOConfig(
    loss_type="aot_unpaired",
    beta=0.01,  # from the caption of Figure 2
)

There is no additional hyperparameter in the paper.

Discovering Preference Optimization Algorithms with and for Large Language Models Copy Section

📜 Paper: https://huggingface.co/papers/2406.08414

An LLM-driven method automatically discovers performant preference optimization algorithms, leading to a new algorithm called DiscoPOP that blends logistic and exponential losses. To reproduce the paper's setting, use this configuration:

from trl import DPOConfig

training_args = DPOConfig(
    loss_type="discopop",
    per_device_train_batch_size=64,  # batch size in Section B.1 of the paper
    learning_rate=5e-7,  # learning rate in Section B.1 of the paper
    beta=0.05,  # β in Section B.1 of the paper,
    discopop_tau=0.05  # τ in Section E of the paper
)

WPO: Enhancing RLHF with Weighted Preference Optimization Copy Section

📜 Paper: https://huggingface.co/papers/2406.11827

WPO reweights preference pairs by their policy probabilities to reduce the off-policy gap in DPO-style training. The loss is:

$$
\mathcal{L}{\text{WPO}} = -\mathbb{E}{(x, y_w, y_l) \sim \mathcal{D}} \left[ \textcolor{red}{w(x, y_w) w(x, y_l)} \log p(y_w \succ y_l \mid x) \right]
$$

where the weight \( w(x, y) \) is defined as:

$$
w(x, y) = \exp\left(\frac{1}{|y|}\sum_{t=1}^{|y|} \log \frac{\pi_\theta(y_t \mid x, y_{<t})}{\sum_{v \in \mathcal{V}} \pi_\theta(v \mid x, y_{<t})^2}\right)
$$

To reproduce the paper's setting, use this configuration:

from trl import DPOConfig

training_args = DPOConfig(
    loss_type="sigmoid",  # DPO loss used in the paper
    beta=0.01,  # Section 4 of the paper
    use_weighting=True,
)

Anchored Preference Optimization and Contrastive Revisions: Addressing Underspecification in Alignment Copy Section

📜 Paper: https://huggingface.co/papers/2408.06266

CLAIR and APO enhance LLM alignment through more contrastive preference pairs and controlled alignment objectives, improving model performance close to GPT4-turbo. To reproduce the paper's setting, use this configuration:

from trl import DPOConfig

training_args = DPOConfig(
    loss_type="apo_zero",  # Section 4 of the paper
    per_device_train_batch_size=64,  # batch size in Section B.1 of the paper
    learning_rate=2e-7,  # learning rate in Section 5.2 of the paper
    beta=0.1,  # β in Section 5.2 of the paper,
)

from trl import DPOConfig

training_args = DPOConfig(
    loss_type="apo_down",  # Section 4 of the paper
    per_device_train_batch_size=64,  # batch size in Section B.1 of the paper
    learning_rate=2e-7,  # learning rate in Section 5.2 of the paper
    beta=0.1,  # β in Section 5.2 of the paper,
)

These parameters only appear in the published version

Length Desensitization in Direct Preference Optimization Copy Section

📜 Paper: https://huggingface.co/papers/2409.06411

Shows that standard DPO is inherently length-sensitive, which often pushes models toward overly long or verbose generations. The paper proposes LD-DPO, which modifies the sequence log-prob aggregation by splitting the longer response into a shared prefix (up to the shorter response length) and an excess tail, then downweighting the tail with a factor \( \alpha \in [0,1] \):

$$
\log \pi_\theta(y_{\text{long}}|x) = \log \pi_\theta(y_{1:l_p}|x) + \alpha \cdot \log \pi_\theta(y_{l_p+1:l}|x, y_{1:l_p}),
\quad
l_p=\min(|y_w|,|y_l|).
$$

Setting \( \alpha=1 \) recovers standard DPO; smaller \( \alpha \) reduces verbosity while preserving preference quality.
The optimal \( \alpha \) depends on the model family and whether you’re training a base vs. instruct model, but the paper suggests \( \alpha=0.5 \) as a strong default starting point.

from trl import DPOConfig

training_args = DPOConfig(
    ld_alpha=0.5,
)

Enhancing the Reasoning Ability of Multimodal Large Language Models via Mixed Preference Optimization Copy Section

📜 Paper: https://huggingface.co/papers/2411.10442

Introduces Mixed Preference Optimization (MPO) to improve multimodal reasoning in MLLMs, addressing distribution shift and weak Chain-of-Thought (CoT) after standard pre-training and SFT. The paper contributes (1) MMPR, an automated pipeline for high-quality multimodal preference data, and (2) MPO, a combined preference objective (pairwise + BCO-style + SFT) that boosts CoT. InternVL2-8B-MPO reaches 67.0 on MathVista (+8.7 over InternVL2-8B), comparable to the 10× larger InternVL2-76B. Used in TRL via [DPOConfig] with composite loss. To reproduce the paper's setting, use this configuration:

from trl import DPOConfig

training_args = DPOConfig(
    loss_type=["sigmoid", "bco_pair", "sft"],  # ℒ = w_p·ℒ_p + w_q·ℒ_q + w_g·ℒ_g (Section 3.2 of the paper)
    loss_weights=[0.8, 0.2, 1.0],  # w_p, w_q, w_g loss weights (Section 7 of the paper)
    learning_rate=5e-6,  # learning rate (Section 7 of the paper)
)

TÜLU 3: Pushing Frontiers in Open Language Model Post-Training Copy Section

📜 Paper: https://huggingface.co/papers/2411.15124

The length-normalized sigmoid loss addresses length bias in DPO by dividing chosen and rejected log-ratio scores by their respective completion lengths before computing the Bradley-Terry loss. This per-token normalization was introduced in SimPO (Appendix I "DPO w/ LN" ablation) as an average log-probability reward for a reference-free setting, and was later adopted for standard reference-model-based DPO in post-training recipes such as Tülu 3 (Section 4.3). The loss is:

$$
\mathcal{L}{\text{sigmoid_norm}} = -\log\sigma!\left(\beta \left({\color{red}\frac{1}{|y_w|}}\log\frac{\pi\theta(y_w|x)}{\pi_{\text{ref}}(y_w|x)} - {\color{red}\frac{1}{|y_l|}}\log\frac{\pi_\theta(y_l|x)}{\pi_{\text{ref}}(y_l|x)}\right)\right),
$$
which can be set with:

from trl import DPOConfig

training_args = DPOConfig(
    loss_type="sigmoid_norm",
)

Kahneman–Tversky Optimization Copy Chapter

Papers relating to the [KTOTrainer]

KTO: Model Alignment as Prospect Theoretic Optimization Copy Section

📜 Paper: https://huggingface.co/papers/2402.01306

KTO derives an alignment objective from prospect theory and learns directly from binary human feedback (liked/disliked), matching or surpassing DPO-style methods while handling imbalanced/noisy signals well.
To reproduce the paper's setting, you can use the default configuration of [KTOTrainer]:

from trl import KTOConfig, KTOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)

trainer = KTOTrainer(
    model=model,
    processing_class=tokenizer,
    args=KTOConfig(),
    train_dataset=...,
)
trainer.train()

Supervised Fine-Tuning Copy Chapter

Papers relating to the [SFTTrainer]

TailSFT: Filtered Fine-Tuning Improves Post-Training Performance Copy Section

📜 Paper: https://huggingface.co/papers/2608.25756

TailSFT preserves response coverage for subsequent reinforcement learning by filtering the sequences whose length-normalized loss has improved the most relative to the initial policy. The tail_sft_gsm8k example implements the paper's distributed selection-batch filtering by subclassing [SFTTrainer]. It records each sequence's initial-policy loss before training and filters examples independently in every selection batch.

python examples/tail_sft_gsm8k/tail_sft.py

EMA Without the Lag: Bias-Corrected Iterate Averaging Schemes Copy Section

📜 Paper: https://huggingface.co/papers/2508.00180

Bias-Corrected Exponential Moving Average (BEMA) improves the stability and efficiency of language model fine-tuning by reducing stochasticity and eliminating bias. To use BEMA with SFT as described in the paper, you can use the [BEMACallback]:

from trl import BEMACallback, SFTTrainer

trainer = SFTTrainer(
    ...
    callbacks=[BEMACallback()],
)

On the Generalization of SFT: A Reinforcement Learning Perspective with Reward Rectification Copy Section

📜 Paper: https://huggingface.co/papers/2508.05629

Dynamic Fine-Tuning (DFT) improves the generalization of Large Language Models (LLMs) by dynamically rescaling gradients, outperforming standard Supervised Fine-Tuning (SFT) and showing competitive results in offline reinforcement learning.

$$
\mathcal{L}{\text{DFT}}(\theta) = \mathbb{E}{(x,y) \sim \mathcal{D}} \left[ - \sum_{t=1}^{|y|} \textcolor{red}{\text{sg}\big(\pi_\theta(y_t \mid y_{<t}, x)\big)} ; \log \pi_\theta(y_t \mid y_{<t}, x) \right]
$$

where \( \text{sg}(\cdot) \) is the stop-gradient operator. To use DFT with SFT as described in the paper, you can use the loss_type="dft" argument:

from trl import SFTConfig

training_args = SFTConfig(
    loss_type="dft",
    ...
)

To closely match the paper’s setup, you can use the following configuration (see Sec. 4.1). Authors also mention that the hyperparameters are not very sensitive (Sec. 4.3):

SFTConfig(
    loss_type="dft",
    learning_rate=5e-5,
    max_length=2048,
    # Target batch size 256; achieved via per-device batch 8 * grad accumulation 32
    per_device_train_batch_size=8,
    gradient_accumulation_steps=32,
)

Fewer Truncations Improve Language Modeling Copy Section

📜 Paper: https://huggingface.co/papers/2404.10830

The paper shows that the standard concatenate-then-split preprocessing (packing_strategy="wrapped") used for LLM training causes many documents to be arbitrarily truncated, which harms learning. It proposes packing document chunks into context windows using a Best-Fit Decreasing bin-packing algorithm, greatly reducing truncation while keeping high token utilization and improving model performance. TRL implements this as the "bfd_split" packing strategy in [SFTConfig]. For more details on packing, see the SFT documentation.

from trl import SFTConfig

training_args = SFTConfig(
    packing=True,
    packing_strategy="bfd_split",
    max_length=4096,
)

Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer Copy Section

📜 Paper: https://huggingface.co/papers/1910.10683

The T5 paper proposes a unified text-to-text framework for transfer learning and introduces sequence packing (Section 3.5.2): grouping multiple short sequences into fixed-length blocks to reduce padding and improve training efficiency. Packing is supported in TRL via [SFTConfig] with the [SFTTrainer]. To enable packing with TRL, use this configuration:

from trl import SFTConfig

training_args = SFTConfig(
    packing=True,  # enable sequence packing (Section 3.5.2 of the paper)
    max_length=512,  # packed sequence length (Section 3.5.2 of the paper)
)

Reparameterizing Uniform Diffusion Models: the Leave-One-Out Denoiser Copy Section

📜 Paper: https://huggingface.co/papers/2605.22765

For uniform discrete diffusion models, the usual plug-in cross-entropy does not learn the plain denoiser \( \mathbb{E}[X_0 \mid X_t] \) but the leave-one-out (LOO) posterior, which predicts each token without using its own corrupted value. The paper derives exact conversions between the denoiser, the LOO posterior, and the score, so a network can be trained as either parameterization with the same cross-entropy. Given logits \( f_\theta \) whose softmax parameterizes the LOO posterior, the denoiser used in the loss is

$$
d_\theta(x_t, t)^\ell = \mathrm{softmax}!\left( f_\theta(x_t, t)^\ell + \log!\left(1 + \frac{K,\alpha_t}{1 - \alpha_t}, x_t^\ell\right) \right),
$$

i.e. an additive correction on the logit of the observed token, where \( K \) is the vocabulary size and \( \alpha_t \) the probability the uniform forward keeps the clean token. The paper reports that the LOO parameterization consistently improves generation. The block-diffusion SFT example examples/sft_diffusion_gemma/sft_diffusion_gemma.py exposes both parameterizations via --model_prediction_type (mean for the plain denoiser, matching the released checkpoint, or mean_loo for the LOO posterior):

accelerate launch --config_file examples/accelerate_configs/deepspeed_zero3.yaml \
    examples/sft_diffusion_gemma/sft_diffusion_gemma.py \
    --use_peft \
    --gradient_checkpointing \
    --model_prediction_type mean_loo \
    --output_dir diffusiongemma-26B-A4B-it-gsm8k-lora

Reducing Activation Recomputation in Large Transformer Models Copy Section

📜 Paper: https://huggingface.co/papers/2205.05198

Full activation checkpointing recomputes every op in a checkpointed region during backward, including attention, even though only a few ops (attention among them) account for most of the recomputation cost relative to the memory they'd cost to keep. The paper's selective activation recomputation checkpoints only those expensive ops instead of the whole region. TRL implements this as selective activation checkpointing (SAC) for [SFTTrainer], saving the attention output during the forward pass so backward does not recompute it:

from trl import SFTConfig

training_args = SFTConfig(
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"selective": True},
)

Parameter-Efficient Fine-Tuning (PEFT) Copy Chapter

For general details on using PEFT with TRL, please refer to the PEFT Integration guide.

LoRA: Low-Rank Adaptation of Large Language Models Copy Section

📜 Paper: https://huggingface.co/papers/2106.09685

Low-Rank Adaptation (LoRA) reduces the number of trainable parameters and GPU memory usage in large-scale pre-trained models while maintaining or improving performance on downstream tasks. TRL integrates LoRA via the PEFT library and can be easily enabled in any TRL trainer by passing a [~peft.LoraConfig] to the peft_config argument. Here is an example of using LoRA with the [SFTTrainer]:

from trl import SFTTrainer
from peft import LoraConfig

trainer = SFTTrainer(
    ...,
    peft_config=LoraConfig(),
)

DoRA: Weight-Decomposed Low-Rank Adaptation Copy Section

📜 Paper: https://huggingface.co/papers/2402.09353

Weight-Decomposed Low-Rank Adaptation (DoRA) can improve the performance of LoRA, especially at low ranks. DoRA decomposes pre-trained weight into two component: magnitude and direction. Direction is handled by normal LoRA, and magnitude is learnable parameters. TRL integrate DoRA via the PEFT library and can be easily enable through setting use_dora=True to the [~peft.LoraConfig].

from peft import LoraConfig

config = LoraConfig(use_dora=True, ...)

Reinforce Leave-One-Out Copy Chapter

Papers relating to the [RLOOTrainer]

Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs Copy Section

📜 Paper: https://huggingface.co/papers/2402.14740

RLOO is a variant of REINFORCE that reduces variance by using leave-one-out baselines. It computes rewards by comparing each sample against the average of all other samples in the batch, providing more stable gradients than standard REINFORCE. To reproduce the paper's setting, use this configuration:

from trl import RLOOConfig

training_args = RLOOConfig(
    per_device_train_batch_size=512,  # section C Training Detail of the paper
    steps_per_generation=2,  # section C Training Detail of the paper
    beta=0.03,  # section C Training Detail of the paper
    num_generations=2,  # experiments of paper different num_generations={2,4}
    learning_rate=1e-6,  # section C Training Detail of the paper
)

REINFORCE++: A Simple and Efficient Approach for Aligning Large Language Models Copy Section

📜 Paper: https://huggingface.co/papers/2501.03262

REINFORCE++ is an enhanced variant of the classical REINFORCE algorithm that incorporates key optimization techniques from PPO while eliminating the need for a critic network. It achieves simplicity, enhanced training stability, and reduced computational overhead through global advantage normalization across the entire batch. To approximate the paper's setting with the [RLOOTrainer], use this configuration:

from trl import RLOOConfig

training_args = RLOOConfig(
    normalize_advantages=True,  # global advantage normalization, core of REINFORCE++
)

Odds Ratio Preference Optimization Copy Chapter

Papers relating to the [experimental.orpo.ORPOTrainer]

ORPO: Monolithic Preference Optimization without Reference Model Copy Section

📜 Paper: https://huggingface.co/papers/2403.07691

The introduction of a reference model-free monolithic odds ratio preference optimization algorithm (ORPO) enhances preference alignment during supervised fine-tuning, surpassing larger models in key evaluations. To reproduce the paper's setting, use this configuration:

from trl.experimental.orpo import ORPOConfig

training_args = ORPOConfig(
    beta=0.1,  # λ odds ratio loss weight (Table 7 of the paper, Mistral-ORPO-β)
    learning_rate=5e-6,  # learning rate (Appendix C of the paper)
    lr_scheduler_type="inverse_sqrt",  # scheduler (Appendix C of the paper)
    num_train_epochs=5,  # Appendix C of the paper
    warmup_steps=200,  # warm-up steps (Appendix C of the paper)
    per_device_train_batch_size=8,  # batch size (Appendix C of the paper)
)

Contrastive Preference Optimization Copy Chapter

Papers relating to the [experimental.cpo.CPOTrainer]

Contrastive Preference Optimization: Pushing the Boundaries of LLM Performance in Machine Translation Copy Section

📜 Paper: https://huggingface.co/papers/2401.08417

Introduces Contrastive Preference Optimization (CPO), a preference-based method for machine translation that trains models to avoid adequate-but-imperfect translations instead of mimicking references as in SFT. The paper analyzes limitations of SFT on MT (including reference quality issues) and shows that applying CPO to ALMA with only 22K parallel sentences yields ALMA-R, which matches or exceeds WMT competition winners and GPT-4 on WMT'21–WMT'23. Used in TRL via [experimental.cpo.CPOTrainer]. To reproduce the paper's setting, use this configuration:

from trl.experimental.cpo import CPOConfig

training_args = CPOConfig(
    loss_type="sigmoid",  # preference learning loss (Section 3 of the paper)
    cpo_alpha=1.0,  # NLL regularizer weight (Section 3 of the paper)
    beta=0.1,  # β temperature (Section 4.2 of the paper)
    learning_rate=1e-4,  # learning rate (official code)
    lr_scheduler_type="inverse_sqrt",  # scheduler (official code)
    num_train_epochs=1,  # Section 4.2 of the paper
    warmup_ratio=0.01,  # warm-up ratio (Section 4.2 of the paper)
    max_length=512,  # max sequence length (Section 4.2 of the paper)
)

SimPO: Simple Preference Optimization with a Reference-Free Reward Copy Section

📜 Paper: https://huggingface.co/papers/2405.14734

SimPO is a simpler yet more effective preference optimization approach that uses the average log probability of a sequence as the implicit reward, eliminating the need for a reference model. It introduces a target reward margin to the Bradley-Terry objective to encourage a larger margin between winning and losing responses. To reproduce the paper's setting, use this configuration:

from trl.experimental.cpo import CPOConfig

training_args = CPOConfig(
    loss_type="simpo",  # SimPO loss (Section 3 of the paper)
    cpo_alpha=0.0,  # no BC regularizer for SimPO
    beta=2.5,  # β in Appendix B of the paper
    simpo_gamma=1.375,  # γ target reward margin, from γ/β=0.55 in Appendix B of the paper
    learning_rate=1e-6,  # learning rate in Appendix B of the paper
    num_train_epochs=1,  # Appendix B of the paper
)

AlphaPO -- Reward shape matters for LLM alignment Copy Section

📜 Paper: https://huggingface.co/papers/2501.03884

AlphaPO is a new Direct Alignment Algorithms (DAAs) method that leverages an alpha-parameter to help change the shape of the reward function beyond the standard log reward. AlphaPO helps maintain fine-grained control over likelihood displacement and over-optimization. To reproduce the paper's setting, use this configuration:

from trl.experimental.cpo import CPOConfig

# Mistral-Instruct from Table 3 of the paper
training_args = CPOConfig(
    loss_type="alphapo",
    alpha=0.25,
    beta=2.5,
    simpo_gamma=0.1,
    learning_rate=7e-7,
    ...
)

Triple Preference Optimization Copy Chapter

Papers relating to the [experimental.tpo.TPOTrainer]

Triple Preference Optimization: Achieving Better Alignment using a Single Step Optimization Copy Section

📜 Paper: https://huggingface.co/papers/2405.16681

Introduces Triple Preference Optimization (TPO), a preference learning method that aligns an LLM with three responses per prompt — a gold (reference) completion, a preferred (chosen) completion and a dispreferred (rejected) completion — in a single optimization step. TPO combines a contrastive objective on the (chosen, rejected) pair with a supervised NLL term on the gold response, removing the need for a separate SFT stage and the reference model used in DPO. Used in TRL via [experimental.tpo.TPOTrainer]. To reproduce the paper's setting (Llama-3-Base, 5K), use this configuration:

from trl.experimental.tpo import TPOConfig

training_args = TPOConfig(
    loss_type="sigmoid",  # contrastive loss between chosen and rejected (Section 3 of the paper)
    tpo_alpha=1.0,  # weight of the NLL term on the gold response (Section 3 of the paper)
    beta=0.01,  # β temperature (Table 6 of the paper)
    learning_rate=5e-7,  # Table 6 of the paper
    num_train_epochs=1,
    max_length=1024,
)

To use the TPO-L variant (length-normalized log-probabilities with a target reward margin γ), set loss_type="tpo-l" and tpo_l_gamma:

from trl.experimental.tpo import TPOConfig

training_args = TPOConfig(
    loss_type="tpo-l",  # length-normalized variant (Section 3 of the paper)
    tpo_alpha=1.0,
    beta=0.01,
    tpo_l_gamma=0.5,  # γ target reward margin (Table 6 of the paper, Llama-3-Base 5K)
    learning_rate=5e-7,
    num_train_epochs=1,
)

Nash Learning from Human Feedback Copy Chapter

Nash Learning from Human Feedback Copy Section

📜 Paper: https://huggingface.co/papers/2312.00886

Introduces Nash-MD, an alternative to standard RLHF that learns a preference model conditioned on two inputs and finds a policy at the Nash equilibrium. Instead of optimizing against a reward model, Nash-MD produces policies that consistently generate responses preferred over those of any competing policy. The algorithm is based on mirror descent principles. TRL shipped an implementation as NashMDTrainer up to v1.13; it is no longer part of the library and remains available in the git history.

Reward Modeling Copy Chapter

Papers relating to the [RewardTrainer]

Helping or Herding? Reward Model Ensembles Mitigate but do not Eliminate Reward Hacking Copy Section

📜 Paper: https://huggingface.co/papers/2312.09244

This paper proposed an auxiliary loss function designed to directly learn a centered reward model. This auxiliary loss minimizes the squared sum of the rewards, encouraging the model to naturally produce mean-zero outputs and thereby resolving the issue of underdetermination.

$$
\mathcal{L}(\theta) = - \mathbb{E}{(x,y^+,y^-) \sim \mathcal{D}} \left[ \log \sigma(r\theta(x, y^+) - r_\theta(x, y^-)) \textcolor{red}{- \eta \cdot (r_\theta(x, y^+) + r_\theta(x, y^-))^2} \right].
$$

To use this auxiliary loss with [RewardTrainer], you can use the center_rewards_coefficient argument in [RewardConfig] as follows:

from trl import RewardConfig

training_args = RewardConfig(
    center_rewards_coefficient=0.01,  # η in the paper
    ...
)

Llama 2: Open Foundation and Fine-Tuned Chat Models Copy Section

📜 Paper: https://huggingface.co/papers/2307.09288

In this paper, the authors propose to leverage their preference ratings being decomposed as a scale of four points (e.g., significantly better) to provide more informative feedback to the reward model. This is done by adding a margin to the loss function, which encourages the reward model to assign larger gaps in scores for pairs with higher preference ratings.

$$
\mathcal{L}(\theta) = - \mathbb{E}{(x,y^+,y^-,\textcolor{red}{m}) \sim \mathcal{D}} \left[ \log \sigma(r\theta(x, y^+) - r_\theta(x, y^-) \textcolor{red}{- m}) \right].
$$

You can add a margin to the loss by adding a margin column to the dataset. The following example shows how to set up a the "Margin Small" setting of the paper.

def add_margin(example):
    preference_to_margin = {
        "significantly better": 1.0,
        "better": 2.0/3.0,
        "slightly better": 1.0/3.0,
        "negligibly better / unsure": 0.0,
    }
    return {"margin": preference_to_margin[example["preference_label"]]}

dataset = dataset.map(add_margin)

Online Direct Preference Optimization Copy Chapter

Papers relating to the [experimental.online_dpo.OnlineDPOTrainer]

Direct Language Model Alignment from Online AI Feedback Copy Section

📜 Paper: https://huggingface.co/papers/2402.04792

Online DPO improves direct alignment from preferences methods by providing real-time feedback from a model, outperforming both DPO and PPO methods.

To use Online DPO, you can use the [experimental.online_dpo.OnlineDPOTrainer].

Exploratory Preference Optimization: Harnessing Implicit Q*-Approximation for Sample-Efficient RLHF Copy Section

📜 Paper: https://huggingface.co/papers/2405.21046

XPO augments the DPO objective with a novel and principled exploration bonus, empowering the algorithm to explore outside the support of the initial model and human feedback data. It is a one-line change to online DPO that is provably sample-efficient and converges to a near-optimal language model policy. TRL shipped an implementation as XPOTrainer up to v1.13; it is no longer part of the library and remains available in the git history.

Distillation Copy Chapter

Papers relating to training a student model with the help of a teacher model.

On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes Copy Section

📜 Paper: https://huggingface.co/papers/2306.13649

Introduces Generalized Knowledge Distillation (GKD), which addresses distribution mismatch in KD for auto-regressive models by training the student on its own generated outputs with teacher feedback, instead of a fixed set of sequences. GKD supports flexible loss functions (e.g. beyond KL when the student cannot match the teacher) and integrates with RL fine-tuning (RLHF). The paper reports results on summarization, translation, arithmetic reasoning, and instruction-tuning. Used in TRL via [experimental.gkd.GKDTrainer], which exposes the paper's on/off-policy mixing (lmbda). [DistillationTrainer] implements the same generalized-JSD objective for the always-on-policy case. To reproduce the paper's setting, use this configuration:

from trl.experimental.gkd import GKDConfig

# XSum summarization task (Table A.1 of the paper)
training_args = GKDConfig(
    lmbda=0.5,  # λ student data fraction (Section 3 of the paper)
    beta=0.5,  # β Generalized JSD interpolation, 0=KL, 1=reverse KL (Section 3 of the paper)
    temperature=1.0,  # student training temperature (Appendix A of the paper)
    max_steps=40000,  # training steps (Table A.1 of the paper)
    learning_rate=3e-4,  # learning rate (Table A.1 of the paper)
    per_device_train_batch_size=32,  # batch size (Table A.1 of the paper)
    warmup_steps=2000,  # warm-up steps (Table A.1 of the paper)
    max_new_tokens=64,  # max output tokens (Table A.1 of the paper)
)

MOPD: Multi-Teacher On-Policy Distillation for Capability Integration in LLM Post-Training Copy Section

📜 Paper: https://huggingface.co/papers/2606.30406

Structures post-training as three stages: general SFT, independent per-domain RL training of one expert per domain (e.g. verifiable-answer RL for math, sandboxed agent RL for software engineering), and a final MOPD stage that fuses the frozen domain experts into a single unified student. In that final stage, the student generates a trajectory per prompt, each trajectory is dispatched to its corresponding domain teacher (never averaged or ensembled across teachers), and the student is updated by minimizing the per-token reverse KL against that one teacher's distribution along the trajectory. Used in TRL via [experimental.async_distillation.AsyncDistillationTrainer], which implements this third, fusion stage: passing more than one entry in teacher_server_urls enables MOPD, with each row's teacher_id column selecting which (already-trained) teacher scores it. Use beta=1.0 to match the paper's reverse-KL objective.

On the Position Bias of On-Policy Distillation Copy Section

📜 Paper: https://huggingface.co/papers/2606.22600

Introduces Importance-Weighted On-Policy Distillation (IW-OPD), which addresses the position bias in OPD by reweighting sampled-token distillation updates according to accumulated teacher-student prefix discrepancy. Early tokens keep larger weights, while later tokens after high drift are downweighted. Used in TRL via [experimental.iw_opd.IWOPDTrainer] with distillation_objective="iw_opd".

The paper optimizes IW-OPD with a clipped policy-gradient setup (verl) and vLLM rollouts. IWOPDTrainer exposes the matching distillation and rollout settings below; policy-optimization settings from the paper such as clipping range 0.2, dual-clip constant 3.0, inner PPO epochs, entropy coefficient, KL reward penalty, auxiliary KL, and rollout importance correction are not IWOPDConfig parameters.

from trl.experimental.iw_opd import IWOPDConfig

# Table 6 and Algorithm 1 of the paper, mapped to IWOPDConfig where available.
training_args = IWOPDConfig(
    distillation_objective="iw_opd",
    iw_opd_gamma=0.5,  # γ amplification, Algorithm 1 and Appendix C.3
    lmbda=1.0,  # fully on-policy rollouts
    learning_rate=1e-5,  # Table 6
    per_device_train_batch_size=1,  # Table 6 uses PPO micro-batch size 1 per GPU
    gradient_accumulation_steps=32,  # with 32 GPUs, this gives the paper's 1024-prompt batch
    num_generations=1,  # Table 6 rollout samples per prompt
    temperature=1.0,  # Table 6 training decoding temperature
    top_p=1.0,  # Table 6 training decoding top-p
    max_prompt_length=2048,  # Table 6
    max_completion_length=16384,  # Table 6
    warmup_ratio=0.0,  # Table 6
    use_vllm=True,  # Table 6 uses vLLM rollouts
    vllm_sync_frequency=1,  # refresh rollout policy after each update
    save_steps=10,  # Table 6 checkpoint frequency
    eval_steps=10,  # Table 6 validation frequency
)

On-Policy Distillation Copy Section

📰 Blog: https://thinkingmachines.ai/blog/on-policy-distillation/

On-Policy Distillation involves a student model generating rollouts for each batch of training data. We subsequently obtain the probability distributions for each token of the rollouts from both the student and teacher models. The student model is then optimized to minimize the negative Kullback-Leibler (KL) divergence between its own token distributions and those of the teacher model.

On-Policy Distillation has been shown to outperform SFT, GRPO and can be used to restore generalization capabilities lost during SFT.

Additionally on-policy distillation is more compute efficient and is less prone to overfitting when trained with limited data.

To train a model with on-policy distillation using TRL, you can use the following configuration, with the [DistillationTrainer] and [DistillationConfig]:

from trl import DistillationConfig

training_args = DistillationConfig(
    beta=1.0,  # to ensure reverse-kl as the loss function
    teacher_model_name_or_path="teacher-model",  # specify the teacher model
)

Alternatively, you can use the [experimental.gkd.GKDTrainer] and [experimental.gkd.GKDConfig]:

from trl.experimental.gkd import GKDConfig

training_args = GKDConfig(
    lmbda=1.0,  # student produces rollouts for all batches
    beta=1.0,  # to ensure reverse-kl as the loss function
    teacher_model_name_or_path="teacher-model",  # specify the teacher model
)

You can also use [experimental.gold.GOLDTrainer] for on-policy distillation; see Unlocking On-Policy Distillation for Any Model Family.

Unlocking On-Policy Distillation for Any Model Family Copy Section

🌐 Project: https://huggingface.co/spaces/HuggingFaceH4/general-on-policy-logit-distillation

Introduces General Online Logit Distillation (GOLD), which extends on-policy distillation to student/teacher pairs with different tokenizers and model families. GOLD aligns textual spans across tokenizers and can apply Universal Logit Distillation (ULD) so completion tokens are not dropped. Used in TRL via [experimental.gold.GOLDTrainer].

from trl.experimental.gold import GOLDConfig

training_args = GOLDConfig(
    lmbda=1.0,  # student produces rollouts for all batches
    beta=1.0,  # to ensure reverse-kl as the loss function
    use_uld_loss=True,  # Universal Logit Distillation for cross-tokenizer pairs
)

Knowledge Distillation of Large Language Models Copy Section

📜 Paper: https://huggingface.co/papers/2306.08543

MiniLLM is the first on-policy knowledge distillation method, which minimizes the sequence-level reverse KLD between the teacher and the student model and is optimized by reinforcement learning.

It is a generalized version of Think Machine Lab's On-Policy Distillation, with the option to add distribution-level single-step distillation signals (like GKD when beta=1) and long-context reverse KLD signals.

TRL shipped an implementation as MiniLLMTrainer up to v1.14; it is no longer part of the library and remains available in the git history.

Reinforcement Learning via Self-Distillation Copy Section

📜 Paper: https://huggingface.co/papers/2601.20802

Self-Distillation Policy Optimization (SDPO) enhances reinforcement learning with verifiable rewards by converting rich textual feedback (e.g., runtime errors, judge evaluations) into a dense learning signal without any external teacher or explicit reward model. SDPO treats the current model conditioned on feedback as a self-teacher and distills its feedback-informed next-token predictions back into the policy. Notably, SDPO also outperforms baselines in standard RLVR environments that only return scalar feedback by using successful rollouts as implicit feedback for failed attempts.

from trl.experimental.sdpo import SDPOConfig, SDPOTrainer

training_args = SDPOConfig(
    distillation_alpha=0.5,                # Jensen-Shannon divergence (recommended)
    distillation_mode="topk_logits",       # Explicitly select top-K logit distillation
    distillation_topk=100,                 # Required for top-K logit distillation
    distillation_is_clip=2.0,              # Importance sampling clipping
    distillation_weight=1.0,               # Convex weight: (1-w)*policy + w*distillation; 1.0 = pure distillation
    use_successful_as_teacher=True,        # Use successful rollouts as teacher
    teacher_model_kind="ema",              # Supported: "base", "live", "ema"
    teacher_update_rate=0.05,              # EMA update rate
    include_environment_feedback=True,     # required to use the dataset's privileged_context (defaults to False)
)

trainer = SDPOTrainer(
    model="Qwen/Qwen2.5-1.5B-Instruct",
    reward_funcs=...,
    args=training_args,
    train_dataset=...,
)
trainer.train()

Expected dataset columns:

• prompt
• privileged_context for optional environment feedback

For more details, see the SDPO Trainer documentation.

Self-Distillation Enables Continual Learning Copy Section

📜 Paper: https://huggingface.co/papers/2601.19897

Self-Distilled Fine-Tuning (SDFT) performs on-policy self-distillation by generating completions during training, then distilling an explicit teacher-conditioned view of those same completions back into the student. The teacher is selected by teacher_model_kind: "base" (the initial student), "live" (the current student), or "ema" (an exponentially averaged teacher). The teacher prompt is composed internally from the student prompt plus the dataset privileged_context.

from datasets import Dataset

from trl.experimental.sdft import SDFTConfig, SDFTTrainer

dataset = Dataset.from_dict(
    {
        "prompt": [[{"role": "user", "content": "Solve 2+2."}]],
        "privileged_context": ["Example answer: 4."],
    }
)

training_args = SDFTConfig(
    distillation_alpha=0.5,
    distillation_mode="topk_logits",
    distillation_topk=5,
    max_completion_length=64,
)

trainer = SDFTTrainer(
    model="Qwen/Qwen2.5-1.5B-Instruct",
    args=training_args,
    train_dataset=dataset,
)
trainer.train()

Expected dataset columns:

• prompt
• privileged_context containing only the extra teacher-only information

For more details, see the SDFT Trainer documentation.

Embarrassingly Simple Self-Distillation Improves Code Generation Copy Section

📜 Paper: https://huggingface.co/papers/2604.01193

Simple Self-Distillation (SSD) improves code generation by sampling completions from the model at a training-time temperature and truncation configuration, then fine-tuning on those raw, unverified samples with standard cross-entropy loss. No reward model, verifier, teacher model, or reinforcement learning is needed. SSD reshapes token distributions in a context-dependent way: suppressing distractor tails at "lock" positions (where syntax leaves little ambiguity) while preserving diversity at "fork" positions (where multiple valid continuations exist).

TRL shipped an implementation as SSDTrainer up to v1.15; it is no longer part of the library and remains available in the git history.

Distributed Training Copy Chapter

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models Copy Section

📜 Paper: https://huggingface.co/papers/1910.02054

ZeRO (Zero Redundancy Optimizer) eliminates memory redundancies in data- and model-parallel training by partitioning optimizer states, gradients, and parameters across devices while retaining low communication volume and high computational granularity. This allows for the efficient training of large models that would otherwise not fit in GPU memory.

TRL supports ZeRO via the DeepSpeed integration. To use it, provide a DeepSpeed configuration file with your desired settings,

# config.yaml
distributed_type: DEEPSPEED
num_processes: 2
deepspeed_config:
  zero_stage: 3

and launch the training script using accelerate launch --config_file config_file.

accelerate launch --config_file config.yaml train.py

Source Quickstart (`docs/source/quickstart.md`) Copy Section

Quickstart

TRL is a comprehensive library for post-training foundation models using techniques like Supervised Fine-Tuning (SFT), Group Relative Policy Optimization (GRPO), Direct Preference Optimization (DPO).

Quick Examples Copy Chapter

Get started instantly with TRL's most popular trainers. Each example uses compact models for quick experimentation.

Supervised Fine-Tuning Copy Section

from trl import SFTTrainer
from datasets import load_dataset

trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B",
    train_dataset=load_dataset("trl-lib/Capybara", split="train"),
)
trainer.train()

Group Relative Policy Optimization Copy Section

from trl import GRPOTrainer
from datasets import load_dataset
from trl.rewards import accuracy_reward

trainer = GRPOTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    train_dataset=load_dataset("trl-lib/DeepMath-103K", split="train"),
    reward_funcs=accuracy_reward,
)
trainer.train()

Direct Preference Optimization Copy Section

from trl import DPOTrainer
from datasets import load_dataset

trainer = DPOTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    train_dataset=load_dataset("trl-lib/ultrafeedback_binarized", split="train"),
)
trainer.train()

Reward Modeling Copy Section

from trl import RewardTrainer
from datasets import load_dataset

dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")

trainer = RewardTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    train_dataset=dataset,
)
trainer.train()

Command Line Interface Copy Chapter

Skip the code entirely - train directly from your terminal:

# SFT: Fine-tune on instructions
trl sft --model_name_or_path Qwen/Qwen2.5-0.5B \
    --dataset_name trl-lib/Capybara

# DPO: Align with preferences  
trl dpo --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
    --dataset_name trl-lib/ultrafeedback_binarized

# Reward: Train a reward model
trl reward --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
    --dataset_name trl-lib/ultrafeedback_binarized

What's Next? Copy Chapter

📚 Learn More Copy Section

• SFT Trainer - Complete SFT guide
• DPO Trainer - Preference alignment
• GRPO Trainer - Group relative policy optimization

🚀 Scale Up Copy Section

• Distributed Training - Multi-GPU setups
• Memory Optimization - Efficient training
• PEFT Integration - LoRA and QLoRA

💡 Examples Copy Section

• Example Scripts - Production-ready code
• Community Tutorials - External guides

Troubleshooting Copy Chapter

Out of Memory? Copy Section

Reduce batch size and enable optimizations:

training_args = SFTConfig(
    per_device_train_batch_size=1,  # Start small
    gradient_accumulation_steps=8,  # Maintain effective batch size
)

training_args = DPOConfig(
    per_device_train_batch_size=1,  # Start small
    gradient_accumulation_steps=8,  # Maintain effective batch size
)

training_args = GRPOConfig(
    per_device_train_batch_size=1,  # Start small
    gradient_accumulation_steps=8,  # Maintain effective batch size
    num_generations=4,              # Reduce from default 8 (GRPO generates num_generations completions per prompt)
    max_completion_length=512,      # Tune based on task; longer sequences cost more memory
)

Loss not decreasing? Copy Section

Try adjusting the learning rate:

training_args = SFTConfig(learning_rate=2e-5)  # Good starting point

For more help, open an issue on GitHub.

Source Toctree.Yml (`docs/source/_toctree.yml`) Copy Section

• sections:
  • local: index
    title: TRL
  • local: installation
    title: Installation
  • local: quickstart
    title: Quickstart
  • local: usage_stats
    title: Usage Stats Collection
    title: Getting started
• sections:
  • local: chat_templates
    title: Chat Templates
  • local: dataset_formats
    title: Dataset Formats
  • local: paper_index
    title: Paper Index
    title: Conceptual Guides
• sections: # Sorted alphabetically
  • local: distillation_trainer
    title: Distillation
  • local: dpo_trainer
    title: DPO
  • local: grpo_trainer
    title: GRPO
  • local: kto_trainer
    title: KTO
  • local: reward_trainer
    title: Reward
  • local: rloo_trainer
    title: RLOO
  • local: sft_trainer
    title: SFT
    title: Trainers
• sections:
  • local: clis
    title: Command Line Interface (CLI)
  • local: jobs_training
    title: Training using Jobs
  • local: customization
    title: Customizing the Training
  • local: reducing_memory_usage
    title: Reducing Memory Usage
  • local: speeding_up_training
    title: Speeding Up Training
  • local: distributing_training
    title: Distributing Training
  • local: aws_trainium
    title: Training on AWS Trainium
  • local: amd_gpus
    title: Training on AMD GPUs
  • local: long_context_training
    title: Training Beyond 1M Tokens
  • local: use_model
    title: Using Trained Models
    title: How-to guides
• sections:
  • local: deepspeed_integration
    title: DeepSpeed
  • local: harbor
    title: Harbor
  • local: kernels_hub
    title: Kernels Hub
  • local: openenv
    title: OpenEnv
  • local: openreward
    title: OpenReward
  • local: peft_integration
    title: PEFT
  • local: rapidfire_integration
    title: RapidFire AI
  • local: trackio_integration
    title: Trackio
  • local: unsloth_integration
    title: Unsloth
  • local: vllm_integration
    title: vLLM
    title: Integrations
• sections:
  • local: example_overview
    title: Example Overview
  • local: community_tutorials
    title: Community Tutorials
  • local: lora_without_regret
    title: LoRA Without Regret
    title: Examples
• sections:
  • sections:
    • local: chat_template_utils
      title: Chat Template Utilities
    • local: data_utils
      title: Data Utilities
    • local: script_utils
      title: Script Utilities
      title: Utilities
  • local: callbacks
    title: Callbacks
  • local: rewards
    title: Reward Functions
    title: API
• sections:
  • local: experimental_overview
    title: Experimental Overview
  • local: a2po_trainer # Sorted alphabetically
    title: A2PO
  • local: async_distillation_trainer
    title: Async Distillation
  • local: async_grpo_trainer
    title: Asynchronous GRPO
  • local: bema_for_reference_model
    title: BEMA for Reference Model
  • local: cpo_trainer
    title: CPO
  • local: gkd_trainer
    title: GKD
  • local: gmpo
    title: GMPO
  • local: gold_trainer
    title: GOLD
  • local: iw_opd_trainer
    title: IW-OPD
  • local: merge_model_callback
    title: MergeModelCallback
  • local: online_dpo_trainer
    title: Online DPO
  • local: orpo_trainer
    title: ORPO
  • local: sdft_trainer
    title: SDFT
  • local: sdpo_trainer
    title: SDPO
  • local: tpo_trainer
    title: TPO
    title: Experimental
    isExpanded: false

Source A2po Trainer (`docs/source/a2po_trainer.md`) Copy Section

A2PO

TRL supports A*-PO (Optimal Advantage Regression) as described in the paper Accelerating RL for LLM Reasoning with Optimal Advantage Regression by Kianté Brantley, Mingyu Chen, Zhaolin Gao, Jason D. Lee, Wen Sun, Wenhao Zhan, and Xuezhou Zhang.

The abstract from the paper is the following:

Reinforcement learning (RL) has emerged as a powerful tool for fine-tuning large language models (LLMs) to improve complex reasoning abilities. However, state-of-the-art policy optimization methods often suffer from high computational overhead and memory consumption, primarily due to the need for multiple generations per prompt and the reliance on critic networks or advantage estimates of the current policy. In this paper, we propose A*-PO, a novel two-stage policy optimization framework that directly approximates the optimal advantage function and enables efficient training of LLMs for reasoning tasks. In the first stage, we leverage offline sampling from a reference policy to estimate the optimal value function V*, eliminating the need for costly online value estimation. In the second stage, we perform on-policy updates using a simple least-squares regression loss with only a single generation per prompt. Theoretically, we establish performance guarantees and prove that the KL-regularized RL objective can be optimized without requiring complex exploration strategies. Empirically, A*-PO achieves competitive performance across a wide range of mathematical reasoning benchmarks, while reducing training time by up to 2× and peak memory usage by over 30% compared to PPO, GRPO, and REBEL.

Usage Copy Chapter

A*-PO assumes a binary, verifiable reward (r ∈ {0, 1}) and runs in two stages:

1. Offline value estimation. Before training, num_value_samples completions are sampled from the reference policy for every prompt and scored with reward_funcs. The optimal value V*(x) = β₁·log(mean_i exp(r(x, yᵢ)/β₁)) is estimated and cached per prompt.
2. On-policy regression. During training, a single completion is generated per prompt from the current policy. The loss is the squared error between the implicit reward β₂·log(π(y|x)/π_ref(y|x)) and the optimal advantage r(x, y) − V*(x).

from trl.experimental.a2po import A2POConfig, A2POTrainer

# A*-PO assumes a binary, verifiable reward in {0, 1}.
def reward_correct(completions, ground_truth, **kwargs):
    return [float(completion.strip() == truth) for completion, truth in zip(completions, ground_truth)]

training_args = A2POConfig(
    output_dir="Qwen2.5-0.5B-A2PO",
    num_value_samples=8,  # Stage 1: samples per prompt from the reference policy to estimate V*
    beta1=0.5,  # Stage 1: KL temperature for the V* estimate
    beta2=1e-3,  # Stage 2: KL temperature for the regression target
)
trainer = A2POTrainer(
    model="Qwen/Qwen2.5-0.5B",
    reward_funcs=reward_correct,
    args=training_args,
    train_dataset=...,
)
trainer.train()

Because V* is estimated entirely from reference-policy samples, A*-PO cannot exceed the reference policy's Pass@K. The official implementation can be found at ZhaolinGao/A-PO.

A2POTrainer Copy Chapter

[[autodoc]] experimental.a2po.A2POTrainer
- train
- save_model
- push_to_hub

A2POConfig Copy Chapter

[[autodoc]] experimental.a2po.A2POConfig

Source Amd Gpus (`docs/source/amd_gpus.md`) Copy Section

Training on AMD GPUs

TRL trainers run on AMD Instinct GPUs without code changes. PyTorch's ROCm build exposes them as cuda devices, so the commands in the rest of the documentation work as they are.

Installation Copy Chapter

Install PyTorch's ROCm build, then TRL in the same environment:

pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm7.2
pip install trl

Attention kernels Copy Chapter

The FlashAttention kernels used in the examples of this documentation (kernels-community/flash-attn2, flash-attn3 and flash-attn4) have no ROCm build. On AMD GPUs, use kernels-community/aiter-flash-attn, downloaded from the Hub by the kernels library (pip install trl[kernels]):

from trl import SFTConfig

training_args = SFTConfig(model_init_kwargs={"attn_implementation": "kernels-community/aiter-flash-attn@v2"})

The kernel has no attention-logit soft-capping, and Transformers skips attn_logit_softcapping without a warning, so models that use it, such as Gemma 2, train with uncapped attention.

Generation with vLLM Copy Chapter

GRPO and RLOO with use_vllm=True need vLLM's ROCm build, see the vLLM installation guide.

Source Async Distillation Trainer (`docs/source/async_distillation_trainer.md`) Copy Section

Async Distillation

[!IMPORTANT]
This trainer requires vllm and transformers>=5.2.0. For distributed training, only FSDP2 is supported (DeepSpeed ZeRO is not).

Currently, vllm and transformers have conflicting dependency constraints. To work around this, install vLLM first and then force-install transformers:

bash Copy

pip install vllm
pip install 'transformers>=5.2.0' --no-deps

Overview Copy Chapter

[experimental.async_distillation.AsyncDistillationTrainer] is the async counterpart to [~trl.experimental.distillation.DistillationTrainer], architected like [experimental.async_grpo.AsyncGRPOTrainer]: a background rollout worker generates the student's own on-policy completions and scores them against a teacher, while training proceeds concurrently instead of alternating between generation and gradient updates. Unlike the synchronous trainer, the teacher is never loaded locally — only a vLLM server URL is needed, so the teacher can run on entirely separate hardware from the student and trainer, or even be a much larger model than would otherwise fit alongside the student.

compute_loss minimizes a generalized Jensen-Shannon Divergence between the student's and teacher's per-position token distributions (beta=0.0 is forward KL, beta=1.0 is reverse KL, values in between interpolate), the same objective [~trl.experimental.distillation.DistillationTrainer] and [~trl.experimental.server_distillation.ServerDistillationTrainer] use. Training is always on-policy: the student generates every completion it trains on.

Multi-teacher on-policy distillation (MOPD) Copy Section

MOPD is not part of the 2306.13649 paper this trainer's core objective is based on (which describes a synchronous, single-teacher setting); it's a separate method, described in MOPD: Multi-Teacher On-Policy Distillation for Capability Integration in LLM Post-Training. There, MOPD is the third of three stages — general SFT, then independent per-domain RL training of one expert per domain, then MOPD fuses those frozen experts into a single student. AsyncDistillationTrainer implements that third, fusion stage only: the per-domain expert teachers must already exist (e.g. trained separately with [GRPOTrainer]/[RLOOTrainer]) and be served over HTTP before you point teacher_server_urls at them. The paper's own Stage 3 uses reverse KL (beta=1.0), not this trainer's default beta=0.0 (forward KL).

teacher_server_urls accepts more than one entry. With a single entry, every sample is scored by that one teacher (plain on-policy distillation). With multiple entries, each training sample's teacher_id column selects which teacher scores it — for example, routing math prompts to a math-specialist teacher and code prompts to a code-specialist teacher, each served independently. Each sample is dispatched to its one matching teacher, never averaged or ensembled across teachers. A sample with a missing or unmapped teacher_id raises rather than silently falling back to the wrong teacher. See examples/async_distillation_math/async_distillation_mopd.py for a runnable two-teacher example.

[!WARNING]
Every teacher must share the student's tokenizer. Completions travel to the teachers as raw token ids, and the candidate ids a teacher reports back index the student's own vocabulary directly in compute_loss. Teachers from the same model family as the student (as in the MOPD example, where a Qwen2.5 student is fused from Qwen2.5 and Qwen2.5-Coder experts) satisfy this; a teacher with a different vocabulary trains the student against the wrong tokens, silently unless its vocabulary is larger than the student's.

Checkpoint and resume: ignore_data_skip defaults to True; the base Trainer's skip-and-replay loop does not apply to a live rollout queue. Instead, the index of the first prompt not yet trained on is saved to rollout_state.json alongside each checkpoint and restored on resume, so the worker fast-forwards to that prompt without replaying samples. It is the trained position, not the generator's: the worker runs ahead of training by the rollout queue depth, and those buffered samples are lost when the run ends, so resuming from the generator's position would skip prompts that were generated but never trained on. Streaming datasets (IterableDataset) cannot be repositioned; their worker restarts from prompt 0 on resume.

How it differs from [`~trl.experimental.distillation.DistillationTrainer`] Copy Chapter

In [~trl.experimental.distillation.DistillationTrainer], the teacher is a locally loaded model: generation, teacher forward pass, and the gradient update all happen sequentially in the same process. [experimental.async_distillation.AsyncDistillationTrainer] separates these concerns the same way [experimental.async_grpo.AsyncGRPOTrainer] separates GRPO's rollout from its update:

• Rollout worker (background process) — generates completions from the student's vLLM server, sends the full sequence back to the resolved teacher's /v1/completions with prompt_logprobs for teacher-forced scoring (no new tokens generated by the teacher), and pushes ready-to-train samples into a queue.
• Training loop (main process) — pulls samples from the queue, computes the generalized-JSD loss, and updates the student's weights.

Because the teacher is scored over HTTP rather than a local forward pass, only a sparse, top-k slice of its distribution is ever transmitted (teacher_top_k), not the full vocabulary — see [experimental.async_distillation.AsyncDistillationConfig]'s beta and teacher_top_k documentation for exactly which candidates the wire protocol guarantees a teacher logprob for at each beta regime.

After every weight_sync_steps training steps, the updated student weights are transferred to its vLLM server via NCCL. As with [experimental.async_grpo.AsyncGRPOTrainer], generation runs ahead of training, so samples may reflect a slightly stale policy; max_staleness controls how many weight updates a sample can lag behind before being discarded.

Quick start Copy Chapter

# train_async_distillation.py
from datasets import load_dataset
from trl.experimental.async_distillation import AsyncDistillationTrainer

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = AsyncDistillationTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    train_dataset=dataset,
)
trainer.train()

The teacher server, the student's vLLM server, and the trainer must run on separate GPUs. Both are plain vllm serve instances, but they need different flags: the teacher is only ever scored (--logprobs-mode processed_logprobs so teacher_temperature reaches its returned logprobs, --max-logprobs -1 so teacher_top_k can exceed vLLM's default cap of 20), while the student is only ever generated from and needs NCCL weight transfer enabled so the trainer can push updated weights into it:

# Terminal 1: teacher server on GPU 0 (static, never updated)
CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen2.5-1.5B-Instruct \
    --port 8001 \
    --logprobs-mode processed_logprobs \
    --max-logprobs -1

# Terminal 2: student's vLLM server on GPU 1 (dev mode + NCCL weight transfer are required)
CUDA_VISIBLE_DEVICES=1 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen2.5-0.5B-Instruct \
    --port 8000 \
    --weight-transfer-config '{"backend":"nccl"}'

# Terminal 3: training on GPU 2
CUDA_VISIBLE_DEVICES=2 accelerate launch train_async_distillation.py

LoRA Copy Chapter

Pass a peft_config to train a LoRA adapter on the student instead of the full model:

from peft import LoraConfig

trainer = AsyncDistillationTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    train_dataset=dataset,
    peft_config=LoraConfig(r=32, lora_alpha=64, target_modules="all-linear"),
)

Every weight sync has to get the updated student to its server. By default the trainer merges the adapter into the base model and sends the full merged weights, which is around 15 GB for a 7B model on every sync. Start the student's vLLM server with --enable-lora and it sends only the adapter instead, about 160 MB at r=32:

CUDA_VISIBLE_DEVICES=1 VLLM_SERVER_DEV_MODE=1 VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 vllm serve Qwen/Qwen2.5-0.5B-Instruct \
    --port 8000 \
    --weight-transfer-config '{"backend":"nccl"}' \
    --enable-lora \
    --max-lora-rank 32 \
    --max-loras 6

VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 exposes the endpoint the trainer posts each new adapter to. Add VLLM_WORKER_MULTIPROC_METHOD=spawn if you serve tensor-parallel, and keep --weight-transfer-config either way: the trainer only chooses a sync mode when it starts, by which point the server is already up, and merged sync is the fallback. The teacher server is unaffected: it is never updated.

--max-lora-rank must be one of 1, 8, 16, 32, 64, 128, 256, 320, 512. It sets the highest rank the server can serve rather than the rank it will serve, so an r=4 adapter works fine under 8. --max-loras must be at least max_staleness + 2: the trainer keeps max_staleness + 1 adapter versions registered so a rollout that started under an older policy can finish under it instead of switching policies mid-generation, and each sync loads the next version before it unloads the oldest. The trainer checks both values when it starts and tells you what to restart the server with.

vLLM loads adapters from disk rather than over the network, so the trainer writes each version under <output_dir>/.vllm_lora/ and gives the server the path. If the server runs on a different machine, that directory has to be on a filesystem both of them can see.

[!WARNING]
<output_dir>/.vllm_lora/ is a serving cache, not where your trained adapter is kept. It only holds the most recent versions, and each one is deleted shortly after it falls out of the staleness window. Checkpoints are what preserve the adapter you trained, so do not turn them off: with save_strategy="no" a multi-hour run finishes and leaves nothing behind. output_dir also has to be readable by the server and has to outlive the job, so node-local scratch will not work, even though it would keep syncs off a slow shared filesystem.

[!TIP]
Adapter sync is not always faster. vLLM generates roughly 1.33–1.39x slower per token when it serves a LoRA adapter than when it serves merged weights, and generation is usually what dominates wall-clock time. What you get in return is a much shorter pause at each sync and about 1% of the sync bandwidth. If generation throughput is your bottleneck, merged sync is often still the better choice.

Logged metrics Copy Chapter

A rollout passes through several stages before it becomes a gradient. Let's first take a look at the different stages, from a prompt to our final batch composition:

Dataset row = PROMPT (message list + optional `teacher_id`)
 └─ ROLLOUT            (1 prompt -> 1 student completion -> 1 teacher scoring call)
     ├─ generate       (student's vLLM /v1/completions, sampled)
     └─ score          (the routed teacher's /v1/completions with `prompt_logprobs`, teacher-forced)
         └─ SAMPLE     (prompt + completion + the teacher's sparse per-position distribution)

════════════ Trainer process boundary: `rollout_buffer` (mp.Queue) ════════════

         └─ SAMPLE (Pulled 1 at a time; dropped if staleness > `max_staleness`)
             └─ ROW            (Planner assigns it to one of `dp` rows, Σ Lᵢ²-balanced)
                 └─ MICRO-BATCH (`dp` rows, one per rank)
                     └─ PACKED ROW (1 concat sequence, `position_ids` reset per sample)
                         └─ FORWARD (`compute_loss`, bs=1, inter-rank padding stripped)
             └─ OPTIMIZER STEP (from `grad_accum` micro-batches)

Reading it top to bottom:

• a rollout is one prompt generated once and scored once. There is no group: distillation has no advantage baseline to compute across generations, so a prompt is not repeated and one rollout yields exactly one training sample.
• a sample is that completion plus the teacher's top-teacher_top_k candidates per completion position. This is what crosses the process boundary.
• past the boundary the trainer plans samples into rows (one per DP rank), packs each row into one concatenated sequence, and accumulates grad_accum micro-batches into a single optimizer step.

So we split our metrics based on which stage (or entity) they are about:

Everything the loss itself measures is logged unprefixed (jsd, entropy, teacher_entropy), with a teacher_*/<teacher_id> breakdown under MOPD — see the divergence below.

Let's also define some terms to make token metrics unambiguous:

• generated tokens are what the student produced (completions/*).
• forwarded tokens are every token the forward pass processes: the prompt and the generated tokens.
• trained tokens are the subset the loss is taken over where completion_mask == 1. Trained ≠ generated: a completion position the teacher scored no candidate for is masked out of the divergence but still forwarded.

How the batch numbers compose Copy Section

A row is what one DP rank forwards in one micro-batch: several samples concatenated into a single sequence, with position_ids restarting at each sample boundary. The planner decides which samples land in which row (balancing Σ Lᵢ² so no rank straggles), and a data collator does the concatenating.

In all our metrics a "step" always means one FULL optimizer step, never a micro-batch. gradient_accumulation_steps micro-batches make one step, and every metric named _per_step is per optimizer step, matching global_step, logging_steps and the rest of the [Trainer] vocabulary. Where a metric is per micro-batch it says so (batch/microbatches_per_step) or it is a per-row quantity (batch/row_*, batch/samples_per_row).

One step therefore holds a fixed number of row-slots:

row-slots per step = gradient_accumulation_steps x world_size

which makes the batch metrics compose into a ladder you can check against each other:

sample                                     sample/forwarded_tokens_mean
  └─ packed into a ROW                     batch/samples_per_row, batch/row_tokens_mean
      └─ one row per DP rank               = one micro-batch
          └─ gradient_accumulation_steps micro-batches = one STEP
                                           batch/samples_per_step, batch/forwarded_tokens_per_step

So batch/samples_per_step ≈ row-slots x batch/samples_per_row, and likewise for tokens. The two sides agree only up to the variation between micro-batches inside the step — the per-step metrics are sums, the per-row ones are means — so expect a fraction of a percent, not an exact match.

batch/row_fill_frac is the one to watch when samples are long: a 10k-token sample tiles a 32k budget badly (three fit, four never do, so the packer often gets two and the row runs ~77% full), while 1k-token samples tile it almost perfectly. That is quantization, not a bug, and token_budget is the lever — bearing in mind that attention is O(L²) per sequence, so a fuller row of long sequences does not cost linearly more memory.

The rollout queue Copy Section

The worker pushes scored samples into a queue and the trainer pulls from it. Four metrics describe that one queue, and they answer different questions:

The last two metrics (perf/rollout_wait_s and rollout/backpressure_s) are mirror images and are never both large! Reading them together with the queue size tells you which side is the bottleneck:

• queue near empty, perf/rollout_wait_s high → generation-bound. The trainer is starving; look at rollout/generated_tok_s, rollout/inflight and rollout/score_s.
• queue near full, rollout/backpressure_s high → trainer-bound. Generation is throttled and its output is aging in the queue, so watch sample/staleness_mean climb.
• both near zero → balanced !

Completions Copy Section

A completion is what the student's vLLM server generated for one prompt.

Rollouts Copy Section

A rollout is one prompt taken all the way through: generated by the student, then teacher-forced through its teacher. Because the teacher call is on the critical path of every rollout (unlike GRPO, where scoring is a separate loop over completed groups), a slow teacher shows up directly as rollout/duration_s.

Samples arriving from the queue Copy Section

Each sample is one RolloutSample: a prompt, the student's completion, and the teacher's sparse scoring of it. The worker pushes them into rollout_buffer; everything below is measured on the trainer side, as it pulls them for training.

Batches Copy Section

One micro-batch is world_size rows (remember packing flattens into 1 sequence), so one per DP rank, so that rank i forwards row i. gradient_accumulation_steps of those make one optimizer step. A step then covers gradient_accumulation_steps × world_size rows in total. Every _per_step metric below is a sum over the whole step and across every rank; the batch/row_* ones are means over the rows.

The divergence Copy Section

What the objective itself measures, averaged over the trained tokens of the window. These are the numbers to watch for learning, as opposed to for throughput.

There is no per-teacher entropy: the student's entropy is a property of its own policy, not of which teacher scored the sample, so the blended metric already covers it.

Performance Copy Section

Throughput measures how many tokens the trainer processes per second. Model FLOPs Utilization (MFU) estimates the percentage of the training devices' theoretical compute capacity used for model forward and backward passes. The capacity used for this estimate depends on the hardware and training precision.

Throughput and MFU are reported on two time bases:

• _fwd_bwd measures performance during the forward and backward passes.
• _wall_clock measures performance over the complete optimizer step, including optimizer updates, weight synchronization, and time spent waiting for rollouts.

The difference helps distinguish compute performance from the cost of waiting and coordination. Both MFU metrics describe the training devices only: they do not include the computation or hardware capacity of external rollout or teacher servers. Wall-clock MFU reflects the time the trainer spends waiting for those services, not the services' own utilization.

MFU is reported only when peak compute capacity is known for the local training device at the selected precision. Otherwise, MFU is omitted; training continues and throughput and timing metrics remain available.

Design philosophy Copy Chapter

This trainer is intentionally kept minimal and is not meant to grow into a general-purpose solution. If you need a feature that is not supported, we recommend cloning the repository and adapting the trainer to your needs directly. New features will only be considered when there is significant community demand.

AsyncDistillationConfig Copy Chapter

[[autodoc]] trl.experimental.async_distillation.AsyncDistillationConfig

AsyncDistillationTrainer Copy Chapter

[[autodoc]] trl.experimental.async_distillation.AsyncDistillationTrainer

Source Async Grpo Trainer (`docs/source/async_grpo_trainer.md`) Copy Section

Asynchronous GRPO

[!IMPORTANT]
This trainer requires vllm and transformers>=5.2.0. For distributed training, only FSDP2 is supported (DeepSpeed ZeRO is not).

Currently, vllm and transformers have conflicting dependency constraints. To work around this, install vLLM first and then force-install transformers:

bash Copy

pip install vllm
pip install 'transformers>=5.2.0' --no-deps

Overview Copy Chapter

[experimental.async_grpo.AsyncGRPOTrainer] implements the same GRPO algorithm but decouples rollout generation from training. A background worker continuously streams completions from a vLLM server while the training loop consumes them, so generation and gradient updates overlap instead of alternating. The API mirrors [GRPOTrainer] — for full details on the GRPO method itself (advantage computation, KL estimation, loss formulation, reward functions, etc.), see the GRPO Trainer documentation. Not all features from [GRPOTrainer] are available; refer to [experimental.async_grpo.AsyncGRPOConfig] for the supported parameters.

This trainer was contributed by Quentin Gallouédec and Amine Dirhoussi.

How it differs from [`GRPOTrainer`] Copy Chapter

In the standard [GRPOTrainer], generation and training are sequential: generate a batch, compute the loss, update weights, repeat. Even in vLLM colocate mode, where generation runs on the same GPUs, one phase must finish before the other begins.

[experimental.async_grpo.AsyncGRPOTrainer] separates these two concerns:

• Rollout worker (background process) — sends prompts to a vLLM server, scores completions with reward functions, computes advantages, and pushes ready-to-train samples into a queue.
• Training loop (main process) — pulls samples from the queue, computes the clipped surrogate loss, and updates the model weights.

The rollout worker runs in a separate process spawned from the trainer, so reward computation never contends with the training loop for the GIL. This has two consequences for what you can pass as reward_funcs, tools, and environment_factory (for the latter, see the OpenEnv guide, which covers the contract and the available integrations):

[!WARNING]
Because we run the rollout worker in a separate process, everything passed to it is pickled. Each reward function, tool, and environment_factory (and anything they close over) must therefore be picklable: use a module-level function, functools.partial, or a callable class instance. Tools are the exception: the model calls a tool by its name, so each tool is registered under its __name__ and must be a module-level function. Lambdas and closures will raise a TypeError at trainer.train(). This is a difference from [GRPOTrainer], where reward functions are called in-process and closures work.

The rollout process also runs with CUDA_VISIBLE_DEVICES="", so it cannot use the GPU. A GPU-backed reward model (e.g. an AutoModelForSequenceClassification scorer) still loads without error but silently falls back to CPU (note that in [GRPOTrainer], such a reward model shares the trainer's GPUs). Keep reward functions CPU-side and lightweight (verifiers like accuracy_reward, format/length checks).

If you do need a GPU reward model, the recommended approach is to serve it behind its own inference engine (vLLM, TGI, …) on separate GPUs and have a lightweight, picklable reward function call it over HTTP. This keeps the reward model on its own device while the rollout process stays CPU-only, and it scales independently of the trainer.

Tools and environment methods can be synchronous or asynchronous. Synchronous tools run on a thread pool sized to max_inflight_tasks, so a slow tool never blocks the other in-flight vLLM requests. Within one assistant turn, tool calls run in the order the model emitted them. Up to max_inflight_tasks synchronous tools therefore run at once, so a tool that reads or writes shared state must be thread-safe.

After every weight_sync_steps training steps, the updated weights are transferred to the vLLM server via NCCL so that subsequent generations reflect the latest policy.

Because generation and training run concurrently, the training samples may have been generated by a slightly older version of the model. The max_staleness parameter controls how many weight updates a sample can lag behind before being discarded.

Concurrency Copy Section

The number of concurrent requests sent to the vLLM server is capped by max_inflight_tasks. By default it is set automatically to max_staleness × per_device_train_batch_size × gradient_accumulation_steps × num_processes — the maximum number of samples the trainer can consume before they become stale. Generating more than this is wasteful since the excess samples will be discarded.

max_inflight_tasks is a ceiling, not the number actually in flight. The worker polls the server's /metrics endpoint every few seconds and keeps a live cap underneath it (rollout/inflight_cap): past a certain concurrency, more requests make vLLM slower, because the KV cache fills, sequences are preempted and recomputed, and cached prefixes are evicted and re-prefilled. The cap starts at KV cache tokens / max_model_len, read from the server at startup: as many requests as fit at full length, so the first step cannot thrash. While usage stays under 0.6 with nothing queued, the cap grows by 25% each time a cap's worth of rollouts has completed, up to the ceiling. When KV cache usage passes 0.8, the cap is cut by 30%; no running request is cancelled, admissions just stop until the pool has drained below the new cap, and no further cut is taken until then. If the server does not expose /metrics, the cap stays at the ceiling.

Checkpoint and resume: ignore_data_skip defaults to True; the base Trainer's skip-and-replay loop does not apply to a live rollout queue. Instead, the index of the first prompt not yet trained on is saved to rollout_state.json alongside each checkpoint and restored on resume, so the worker fast-forwards to that prompt without replaying samples. It is the trained position, not the generator's: the worker runs ahead of training by the rollout queue depth, and those buffered samples are lost when the run ends, so resuming from the generator's position would skip prompts that were generated but never trained on. Streaming datasets (IterableDataset) cannot be repositioned; their worker restarts from prompt 0 on resume.

Quick start Copy Chapter

# train_async_grpo.py
from datasets import load_dataset
from trl.experimental.async_grpo import AsyncGRPOTrainer
from trl.rewards import accuracy_reward

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = AsyncGRPOTrainer(
    model="Qwen/Qwen3-4B",
    reward_funcs=accuracy_reward,
    train_dataset=dataset,
)
trainer.train()

The vLLM server and the trainer must run on separate GPUs. Use CUDA_VISIBLE_DEVICES to partition your GPUs. For example, with 2 GPUs, you can run the vLLM server on GPU 0 and the trainer on GPU 1 as follows:

# Terminal 1: vLLM server on GPU 0 (dev mode + NCCL weight transfer are required)
CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3-4B \
    --max-model-len 4096 \
    --logprobs-mode processed_logprobs \
    --weight-transfer-config '{"backend":"nccl"}'

[!TIP]
Set --max-model-len to the maximum total sequence length (prompt + completion) you expect. A lower value reduces GPU memory usage on the server, freeing more memory for the KV cache and increasing throughput. A good starting point is the prompt length plus max_completion_length from your config.

# Terminal 2: training on GPU 1
CUDA_VISIBLE_DEVICES=1 accelerate launch train_async_grpo.py

AMD GPUs Copy Chapter

The trainer runs on AMD GPUs (ROCm) with no extra setup: the commands above work as they are. Training uses the kernels-community/aiter-flash-attn attention kernel on AMD GPUs, kernels-community/flash-attn4 on NVIDIA Blackwell GPUs and kernels-community/flash-attn3 on other NVIDIA GPUs, all downloaded from the Hub by the kernels library. Models that cap their attention logits (attn_logit_softcapping), such as Gemma 2, are not supported on AMD GPUs or on NVIDIA Blackwell GPUs.

Vision-language models Copy Chapter

Vision-language models (Qwen3.5, Qwen3.6, Qwen3-VL, …) can be trained on text-only datasets. Pass the model id as usual; nothing else changes:

CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3.5-2B \
    --max-model-len 4096 \
    --logprobs-mode processed_logprobs \
    --weight-transfer-config '{"backend":"nccl"}'

The vision tower is frozen. A text-only dataset never produces image tokens, so the tower is never exercised by the forward pass. Everything outside the text tower (vision tower, multimodal projector) has requires_grad=False: it gets no gradients and no optimizer state, and weight sync skips it entirely — the server keeps the values it loaded from the checkpoint. The tower is still loaded, so it costs GPU memory for its weights.

[!WARNING]
Images are not supported yet. Prompts containing images are not passed to the vLLM server or to the training forward pass. Multimodal training also needs the padding-free packing path to build 3D M-RoPE positions from the image grid, which is not implemented here.

[!WARNING]
Hybrid models (Qwen3.5, Qwen3.6) need flash-linear-attention installed, or their gated-DeltaNet layers silently fall back to a pure-PyTorch scan that costs ~20x (measured on Qwen3.5-2B: 1046 vs 52 µs/token). Those layers also carry recurrent state across a padding-free packed row, which the trainer does not reset at sample boundaries, so their training log-probs drift from what the server generated as more sequences are packed per row.

Top-k and top-p sampling Copy Chapter

Top-k and top-p renormalize the sampling distribution over a restricted candidate set rather than the whole vocabulary. The server's log-probabilities reflect that (--logprobs-mode processed_logprobs), so the trainer has to normalize over the same set, otherwise every importance ratio is biased: a token's probability under the whole vocabulary is lower than under the candidate set it was actually drawn from (DeepSeek-V3.2, section 3.1). With top_p < 1 or top_k > 0, the worker asks the server for the candidate set of each sampled token and the trainer renormalizes its log-probabilities over it. This needs two more server flags, and top_k has to be set, since vLLM bounds the size of the candidate set with it:

CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3-4B \
    --max-model-len 4096 \
    --logprobs-mode processed_logprobs \
    --weight-transfer-config '{"backend":"nccl"}' \
    --return-sampling-mask \
    --enable-scale-out

config = AsyncGRPOConfig(top_p=0.95, top_k=512)

--return-sampling-mask makes vLLM record the candidate set (vLLM 0.28 or later, on its V2 model runner: VLLM_USE_V2_MODEL_RUNNER=1 if the server refuses the flag), --enable-scale-out exposes the token-in-token-out endpoint that returns it. The worker checks both on the server when it starts. An older vLLM cannot return the set, so the trainer keeps normalizing over the whole vocabulary and warns about the biased ratio.

LoRA Copy Chapter

Pass a peft_config to train a LoRA adapter instead of the full model:

from peft import LoraConfig

trainer = AsyncGRPOTrainer(
    model="Qwen/Qwen3-4B",
    reward_funcs=reward_len,
    train_dataset=dataset,
    peft_config=LoraConfig(r=32, lora_alpha=64, target_modules="all-linear"),
)

Every weight sync has to get the updated policy to the server. By default the trainer merges the adapter into the base model and sends the full merged weights, which is around 15 GB for a 7B model on every sync. Start the vLLM server with --enable-lora and it sends only the adapter instead, about 160 MB at r=32:

CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 vllm serve Qwen/Qwen3-4B \
    --max-model-len 4096 \
    --logprobs-mode processed_logprobs \
    --weight-transfer-config '{"backend":"nccl"}' \
    --enable-lora \
    --max-lora-rank 32 \
    --max-loras 6

VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 exposes the endpoint the trainer posts each new adapter to. Keep --weight-transfer-config even with --enable-lora: the trainer only chooses a sync mode when it starts, by which point the server is already up, and merged sync is the fallback.

--max-lora-rank must be one of 1, 8, 16, 32, 64, 128, 256, 320, 512. It sets the highest rank the server can serve rather than the rank it will serve, so an r=4 adapter works fine under 8. --max-loras must be at least max_staleness + 2: the trainer keeps max_staleness + 1 adapter versions registered so a rollout that started under an older policy can finish under it instead of switching policies mid-generation, and each sync loads the next version before it unloads the oldest. The trainer checks both values when it starts and tells you what to restart the server with.

vLLM loads adapters from disk rather than over the network, so the trainer writes each version under <output_dir>/.vllm_lora/ and gives the server the path. If the server runs on a different machine, that directory has to be on a filesystem both of them can see.

[!WARNING]
<output_dir>/.vllm_lora/ is a serving cache, not where your trained adapter is kept. It only holds the most recent versions, and each one is deleted shortly after it falls out of the staleness window. Checkpoints are what preserve the adapter you trained, so do not turn them off: with save_strategy="no" a multi-hour run finishes and leaves nothing behind. output_dir also has to be readable by the server and has to outlive the job, so node-local scratch will not work, even though it would keep syncs off a slow shared filesystem.

[!TIP]
Adapter sync is not always faster. vLLM generates roughly 1.33–1.39x slower per token when it serves a LoRA adapter than when it serves merged weights, and generation is usually what dominates wall-clock time in RL. What you get in return is a much shorter pause at each sync and about 1% of the sync bandwidth. If generation throughput is your bottleneck, merged sync is often still the better choice.

Design philosophy Copy Chapter

This trainer is intentionally kept minimal and is not meant to grow into a general-purpose solution. If you need a feature that is not supported, we recommend cloning the repository and adapting the trainer to your needs directly. New features will only be considered when there is significant community demand.

Logged metrics Copy Chapter

A rollout passes through several stages before it becomes a gradient. Let's first take a look at the different stages, from a prompt to our final batch composition:

Dataset row = PROMPT (message list + reward kwargs)
 └─ GROUP                  (1 prompt, `num_generations` rollouts, 1 advantage baseline)
     └─ ROLLOUT            (1 conversation, keyed by `rollout_id`)
         ├─ TURN           (1 vLLM /v1/completions call + tool messages fed back)
         │   └─ TurnRecord (prompt_ids, output_ids, output_log_probs)
         └─ reconcile      (`_chain_to_sequences` classifies drift per turn: CLEAN/REALIGN/FORK)
             └─ SEQUENCE   (≥1 per rollout; new one per fork; dropped if no trained token)
                 └─ SAMPLE (Sequence + group advantage + reward + metrics)

════════════ Trainer process boundary: `rollout_buffer` (mp.Queue) ════════════

                 └─ SAMPLE (Pulled 1 at a time; dropped if staleness > `max_staleness`)
                     └─ ROW            (Planner inserts it into one of `dp` prefix tries, attention-cost-balanced)
                         └─ MICRO-BATCH (`dp` rows, one per rank)
                             └─ PACKED ROW (Collator lays each trie out: unique tokens, depth as `position_ids`)
                                 └─ FORWARD (`compute_loss`, bs=1, padding stripped, FA3 per trie segment)
                     └─ OPTIMIZER STEP (from `grad_accum` micro-batches)

Reading it top to bottom:

• a group is one prompt generated num_generations times. The advantage baseline is computed inside it, which is why the group is also the unit that gets scored.
• a rollout is one conversation. Every turn re-tokenizes the whole message list, so drift between turns has to be reconciled at the end.
• reconciling can fork one rollout into several sequences, which is why a single conversation can produce more than one training sample.
• a sample is a sequence plus its advantage, reward and metrics. This is what crosses the process boundary.
• past the boundary the trainer plans samples into rows (one per DP rank), folds each row into a prefix trie so shared tokens are forwarded once, and accumulates grad_accum micro-batches into a single optimizer step.

So we split our metrics based on which stage (or entity) they are about:

Let's also define some terms to make token metrics unambiguous:

• generated tokens are what the model produced (completions/*).
• forwarded tokens are every token the forward pass processes: the prompt, any tool results, and the generated tokens.
• trained tokens are the subset the loss is taken over where completion_mask == 1. For example, a fork or a realign can demote already-generated tokens to context, so trained ≠ generated.

How the batch numbers compose Copy Section

A row is what one DP rank forwards in one micro-batch: several samples packed into a single prefix forest, so tokens they share are forwarded once — see Tree packing. The planner decides which samples land in which row (balancing attention cost so no rank straggles), and a data collator does the packing.

In all our metrics a "step" always means one FULL optimizer step, never a micro-batch. gradient_accumulation_steps micro-batches make one step, and every metric named _per_step is per optimizer step, matching global_step, logging_steps and the rest of the [Trainer] vocabulary. Where a metric is per micro-batch it says so (batch/microbatches_per_step) or it is a per-row quantity (batch/row_*, batch/samples_per_row).

One step therefore holds a fixed number of row-slots:

row-slots per step = gradient_accumulation_steps x world_size

which makes the batch metrics compose into a ladder you can check against each other:

sample                                     sample/forwarded_tokens_mean
  └─ packed into a ROW                     batch/samples_per_row, batch/row_tokens_mean
      └─ one row per DP rank               = one micro-batch
          └─ gradient_accumulation_steps micro-batches = one STEP
                                           batch/samples_per_step, batch/forwarded_tokens_per_step

So batch/samples_per_step ≈ row-slots x batch/samples_per_row, and likewise for tokens. The two sides agree only up to the variation between micro-batches inside the step — the per-step metrics are sums, the per-row ones are means — so expect a fraction of a percent, not an exact match.

batch/row_fill_frac is the one to watch when samples are long: a 10k-token sample tiles a 32k budget badly (three fit, four never do, so the packer often gets two and the row runs ~77% full), while 1k-token samples tile it almost perfectly. That is quantization, not a bug, and token_budget is the lever — bearing in mind that attention is O(L²) per sequence, so a fuller row of long sequences does not cost linearly more memory.

Tree packing Copy Section

A row is a prefix trie, not a concatenation: a token shared by several samples is forwarded once, and attention still shows each sample exactly its own history. Rollouts of one prompt all repeat it, and every turn of a multi-turn rollout re-sends the conversation, so the decoder's work drops by the packing ratio (raw tokens ÷ forwarded tokens). It is a layout change only: every trained token keeps its own target, advantage and old log-probability, and its log-probability matches forwarding the sample on its own.

Where it happens:

planner    TreeRow.add            the trie grows sample by sample; a sample is priced by its novel tokens
   │                              (TreeRow.fits), so `token_budget` bounds unique tokens, not raw ones
collator   TreeRow.pack           DFS layout: unique tokens, depth as position_ids, one loss term per
   │                              original trained token (token_index), and per segment its queries and
   │                              the K/V indices of its root-to-segment path, all built on CPU
attention  tree_attention         gather K/V along the paths, one FA3 varlen call, causal

How it works, on two rollouts of one prompt:

/* Detailed source-code truncated for AI context efficiency. */

A few consequences:

• Any common prefix is shared: the rollouts of a prompt, and across prompts a common system prompt or tool schema. The planner sends a sample to the row already holding its group_id, where its prompt already is. When that row is full, the sample goes to another row if a rank is still empty, or if that does not raise the most loaded row, which sets the micro-batch's time; otherwise the micro-batch closes.
• Nothing shared is plain packing: every sample is then one segment whose path is itself, which is ordinary padding-free FA3 varlen packing.
• Only K/V is duplicated: a shared prefix is gathered once per descendant segment, which is about what unpacked training would hold, and small with GQA.

What it buys: with nothing shared, it runs at about the speed of plain padding-free packing, and the speed-up grows roughly with the packing ratio (about 2.8–3x at a ratio of 3.4 on Qwen3-0.6B/1.7B, H100). The ratio is shared prefix × samples per row, so batch/samples_per_row is the leading indicator: near 1, nothing can be shared, and the lever is token_budget.

Judge it on throughput, batch/trained_tokens_per_step / perf/fwd_bwd_s. A packed micro-batch holds more samples, so perf/fwd_bwd_s rises and perf/mfu_*, which divides by forwarded tokens, falls — neither is a regression. More samples per row also means more rollouts drawn per step, so a generation-bound run becomes more so.

The rollout queue Copy Section

The worker pushes scored samples into a queue and the trainer pulls from it. Four metrics describe that one queue, and they answer different questions:

The last two metrics (perf/rollout_wait_s and rollout/backpressure_s) are mirror images and are never both large! Reading them together with the queue size tells you which side is the bottleneck:

• queue near empty, perf/rollout_wait_s high → generation-bound. The trainer is starving; look at rollout/generated_tok_s, rollout/inflight and rollout/score_queue_size.
• queue near full, rollout/backpressure_s high → trainer-bound. Generation is throttled and its output is aging in the queue, so watch sample/staleness_mean climb.
• both near zero → balanced !

Completions Copy Section

A completion is what vLLM generated for one prompt, counted across every turn of the rollout.

Rollouts Copy Section

A rollout is one full conversation: a prompt generated to completion, including every tool round-trip it took to get there. num_generations rollouts share a prompt and form a group.

Tools Copy Section

Logged only when the model executes tools. tools/<name>/* repeats per tool name, the same way rewards/<func> repeats per reward function, because one failing tool is invisible in an average over all of them. The built-in tool loop and the OpenEnv harness worker (HarnessRolloutWorker, see training on harnesses) log the same names, so one dashboard reads both; the path column says which produce each metric.

Environments Copy Section

Logged only with an environment_factory. Both calls run on the worker's tool thread pool, so a slow environment delays its own rollout but not the others in flight.

Samples arriving from the queue Copy Section

Each sample is one RolloutSample: a reconciled sequence plus the group advantage, the reward and its per-sample metrics, assembled by the worker in _score_group once every generation in a group has finished and been scored. The worker pushes them into rollout_buffer; everything below is measured on the trainer side, as it pulls them for training.

Batches Copy Section

One micro-batch is world_size rows (remember packing flattens into 1 sequence), so one per DP rank, so that rank i forwards row i. gradient_accumulation_steps of those make one optimizer step. A step then covers gradient_accumulation_steps × world_size rows in total. Every _per_step metric below is a sum over the whole step and across every rank; the batch/row_* ones are means over the rows.

Performance Copy Section

Throughput measures how many tokens the trainer processes per second. Model FLOPs Utilization (MFU) estimates the percentage of the training devices' theoretical compute capacity used for model forward and backward passes. The capacity used for this estimate depends on the hardware and training precision.

Throughput and MFU are reported on two time bases:

• _fwd_bwd measures performance during the forward and backward passes.
• _wall_clock measures performance over the complete optimizer step, including optimizer updates, weight synchronization, and time spent waiting for rollouts.

The difference helps distinguish compute performance from the cost of waiting and coordination. Both MFU metrics describe the training devices only: they do not include the computation or hardware capacity of external rollout or teacher servers. Wall-clock MFU reflects the time the trainer spends waiting for those services, not the services' own utilization.

MFU is reported only when peak compute capacity is known for the local training device at the selected precision. Otherwise, MFU is omitted; training continues and throughput and timing metrics remain available.

AsyncGRPOConfig Copy Chapter

[[autodoc]] trl.experimental.async_grpo.AsyncGRPOConfig

AsyncGRPOTrainer Copy Chapter

[[autodoc]] trl.experimental.async_grpo.AsyncGRPOTrainer

RolloutWorkerProtocol Copy Chapter

[[autodoc]] trl.experimental.async_grpo.async_grpo_trainer.RolloutWorkerProtocol

Source Aws Trainium (`docs/source/aws_trainium.md`) Copy Section

Training on AWS Trainium

AWS Trainium runs PyTorch natively through the torch-neuronx package, which registers the neuron device. TRL trainers run on it without code changes. Neuron compiles a graph for each input shape it sees, and each compilation takes about a minute, so most of this page is about keeping shapes static. For installation and the Transformers side, see the Transformers AWS Trainium guide.

Installation Copy Chapter

Install torch-neuronx by following the AWS Neuron setup guide, then install TRL in the same environment.

pip install trl

Supervised fine-tuning Copy Chapter

The [SFTTrainer] quick start runs as is:

from datasets import load_dataset
from trl import SFTTrainer

trainer = SFTTrainer(
    model="Qwen/Qwen3-0.6B",
    train_dataset=load_dataset("trl-lib/Capybara", split="train"),
)
trainer.train()

On Neuron, pad_to_multiple_of defaults to 128 (its "auto" value), so the sequences are padded to a multiple of 128. This bounds the number of shapes, and so of compilations, to max_length / 128. Without it, every new sequence length triggers a compilation. On one NeuronCore, with per_device_train_batch_size=4, 20 steps of the example above take 130 s with the default and 260 s with pad_to_multiple_of=None.

The same applies to every trainer that has a pad_to_multiple_of option: DPO, KTO, Reward, GRPO, RLOO and Distillation.

Compiling the model Copy Chapter

Set torch_compile=True to compile the model with the neuron backend. Dynamo's automatic dynamic shapes produce graphs the Neuron backend can't compile yet (#7660), so turn them off:

import torch
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer

torch._dynamo.config.automatic_dynamic_shapes = False

trainer = SFTTrainer(
    model="Qwen/Qwen3-0.6B",
    args=SFTConfig(torch_compile=True),
    train_dataset=load_dataset("trl-lib/Capybara", split="train"),
)
trainer.train()

The first step compiles for about 2 minutes. After that, with per_device_train_batch_size=4, Qwen3-0.6B trains at 1.15 s/step on one NeuronCore, against 1.71 s/step without compilation.

Online methods Copy Chapter

vLLM doesn't run on Neuron, so GRPO and RLOO generate with Transformers. With the default dynamic cache, every generated token has a new shape, and a step takes several minutes. Use the static cache, which keeps the decoding shapes fixed:

from datasets import load_dataset
from trl import GRPOConfig, GRPOTrainer


# Reward completions close to 200 characters
def reward_len(completions, **kwargs):
    return [-abs(200 - len(completion)) / 100 for completion in completions]


trainer = GRPOTrainer(
    model="Qwen/Qwen3-0.6B",
    reward_funcs=reward_len,
    args=GRPOConfig(
        per_device_train_batch_size=4,
        num_generations=4,
        max_completion_length=128,
        cache_implementation="static",
        model_init_kwargs={"dtype": "bfloat16"},
    ),
    train_dataset=load_dataset("trl-lib/tldr", split="train"),
)
trainer.train()

The prompts are padded to pad_to_multiple_of before generation too, so the static cache only takes a few sizes. Each new size still compiles for about 12 minutes on one NeuronCore; once a size is compiled, a step of this example takes about 25 s.

[!NOTE]
The static cache needs a bf16 model for now: with the default fp32 model, the keys and values end up in different dtypes (#7664).

Data parallelism Copy Chapter

On a trn2.3xlarge instance, which contains one Trainium2 chip, the default Logical NeuronCore configuration (LNC=2) pairs the chip's 8 physical NeuronCores into 4 logical NeuronCores.

To train on all four, use Accelerate:

accelerate launch --num_processes 4 train.py

In this setup, SFT runs 3.6x faster on 4 NeuronCores than on 1.

To restrict a run to a subset of NeuronCores, set NEURON_RT_VISIBLE_CORES to 0,1, 2,3, or 0-3. Other combinations such as 1,2 are not supported: the Neuron runtime fails to set up collectives across them.

Source Bema For Reference Model (`docs/source/bema_for_reference_model.md`) Copy Section

BEMA for Reference Model

This feature implements the BEMA algorithm to update the reference model during DPO training.

Usage Copy Chapter

from trl.experimental.bema_for_ref_model import BEMACallback, DPOTrainer
from datasets import load_dataset

dataset = load_dataset("trl-internal-testing/zen", "standard_preference", split="train")

bema_callback = BEMACallback(update_ref_model=True)

trainer = DPOTrainer(
    model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
    train_dataset=dataset,
    callbacks=[bema_callback],
)
trainer.train()

DPOTrainer Copy Chapter

[[autodoc]] experimental.bema_for_ref_model.DPOTrainer
- train
- save_model
- push_to_hub

BEMACallback Copy Chapter

[[autodoc]] experimental.bema_for_ref_model.BEMACallback

Source Callbacks (`docs/source/callbacks.md`) Copy Section

Callbacks

RichProgressCallback Copy Chapter

[[autodoc]] RichProgressCallback

LogCompletionsCallback Copy Chapter

[[autodoc]] LogCompletionsCallback

BEMACallback Copy Chapter

[[autodoc]] BEMACallback

WeaveCallback Copy Chapter

[[autodoc]] WeaveCallback

Source Chat Template Utils (`docs/source/chat_template_utils.md`) Copy Section

Chat template utilities

For an overview of the chat templates bundled with TRL and the rationale behind the training patches, see Chat Templates.

clone_chat_template Copy Chapter

[[autodoc]] clone_chat_template

is_chat_template_prefix_preserving Copy Chapter

[[autodoc]] chat_template_utils.is_chat_template_prefix_preserving

get_training_chat_template Copy Chapter

[[autodoc]] chat_template_utils.get_training_chat_template

Source Chat Templates (`docs/source/chat_templates.md`) Copy Section

Chat Templates

A chat template is a Jinja2 snippet that formats messages into the string a model was trained on. For example:

>>> from transformers import AutoTokenizer
>>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
>>> tokenizer.chat_template
"{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}{% endif %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
>>> tokenizer.apply_chat_template([{"role": "user", "content": "Hi!"}], tokenize=False)
'<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nHi!<|im_end|>\n'

In most cases you don't need to worry about chat templates: models ship their template along with the tokenizer, and TRL applies it for you. The whole thing is transparent. But some TRL recipes rely on features that most shipped templates don't include:

• SFT with assistant_only_loss=True needs {% generation %} / {% endgeneration %} markers around assistant output, so the loss mask can target only assistant tokens.
• GRPO with tool calls needs the template to be prefix-preserving: appending a tool message must not change how earlier messages are rendered.

TRL ships patched templates under trl/chat_templates/ for common families (Qwen, Llama, DeepSeek-V3, GPT-OSS, ...) and swaps them in automatically for supported models. For any other model, you'll need to patch its template yourself. The rest of this page catalogs what's bundled.

Supported model families Copy Chapter

TRL stores reference copies of the original templates so it can identify supported models at init and swap in a training template when needed. The following families are recognized: Cohere, Cohere2, DeepSeek-V3, DeepSeek-R1-Distill, DiffusionGemma, Gemma, Gemma3, Gemma4, GLM-4-MoE, GPT-OSS, Idefics3, LFM2, LFM2.5 (including LFM2.5-2.6B and LFM2.5-24B-A2B), LFM2.5-VL, Llama 3 / 3.1 / 3.2, Llava-Next, Muse Glimmer, Nemotron 3 (Nano, Super, Ultra), Nemotron 3.5 Lightning, Phi-3, Phi-3.5, Qwen2-VL, Qwen2.5, Qwen2.5-VL, Qwen3 (including the Instruct-2507 variant), Qwen3-VL, Qwen3.5, Qwen3.6, Qwen3.8, SmolVLM (including SmolVLM2).

Training templates Copy Chapter

Patched templates that fix training-specific issues. Swapped in at init when tools are enabled (GRPO) or when assistant_only_loss=True (SFT).

`cohere_training.jinja` Copy Section

Patched Cohere template. Diff vs cohere.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`cohere2_training.jinja` Copy Section

Patched Cohere2 template. Diff vs cohere2.jinja:

Move the trailing <|END_OF_TURN_TOKEN|> from after the role-dispatch {% endif %} into each role branch, and wrap the assistant branch (<|START_RESPONSE|>...<|END_RESPONSE|><|END_OF_TURN_TOKEN|>) with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`deepseek_r1_distill_training.jinja` Copy Section

Patched DeepSeek-R1-Distill template. Diff vs deepseek_r1_distill.jinja:

• Removes the '' split on assistant content. The original drops everything up to and including ``, which would silently remove the reasoning from the training target.
• Wraps assistant message output with {% generation %} / {% endgeneration %} markers for SFT assistant-only loss.

`deepseekv3_training.jinja` Copy Section

Patched DeepSeek-V3 template. Diff vs deepseekv3.jinja:

• Uses | tojson on tool['function']['arguments'] so that arguments can be passed as a dict (the documented format per transformers docs). The original template uses raw string concatenation, which crashes on dict inputs.
• Wraps assistant message output with {% generation %} / {% endgeneration %} markers for SFT assistant-only loss.

`diffusion_gemma_training.jinja` Copy Section

Patched DiffusionGemma template. Diff vs diffusion_gemma.jinja:

Wrap the parts of a model turn that the model emits (thinking channel, tool calls, content, and the closing <turn|>) with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss. Tool responses embedded in the turn come from the environment and stay outside the generation block.

`gemma_training.jinja` Copy Section

Patched Gemma template (shared by Gemma and Gemma2, which ship identical chat templates). Diff vs gemma.jinja:

Split the unified assistant output so that the <start_of_turn>model\n header (a prompt cue, not generated by the model) sits outside the generation block, and wrap the assistant content with {% generation %} / {% endgeneration %} markers for SFT assistant-only loss.

`gemma3_training.jinja` Copy Section

Patched Gemma 3 template. Same diff as gemma_training.jinja (split the unified output line into role-specific branches so the <start_of_turn>model\n prompt cue sits outside the generation block, and wrap the assistant content with {% generation %} / {% endgeneration %}), applied to the Gemma 3 base template that supports system messages and multimodal content blocks.

`glm4moe_training.jinja` Copy Section

Patched GLM-4-MoE template. Diff vs glm4moe.jinja:

Require both and to be present before parsing, to avoid incorrect splitting when the model generates only one tag:

- {%- if '</think>' in content %}
+ {%- if '<think>' in content and '</think>' in content %}

Wrap assistant message output (including the thinking block and tool calls) with {% generation %} / {% endgeneration %} markers for SFT assistant-only loss.

`qwen3_training.jinja` Copy Section

Patched Qwen3 template. Diff vs qwen3.jinja:

Require both and to be present before parsing, to avoid incorrect splitting when the model generates only one tag:

- {%- if '</think>' in content %}
+ {%- if '<think>' in content and '</think>' in content %}

Always include the thinking block regardless of message position. The original conditionally omits it based on loop.last, which changes the assistant rendering when a tool message is appended, breaking prefix-preservation:

- {%- if loop.index0 > ns.last_query_index %}
-     {%- if loop.last or (not loop.last and reasoning_content) %}
-         {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
-     {%- else %}
-         {{- '<|im_start|>' + message.role + '\n' + content }}
-     {%- endif %}
- {%- else %}
-     {{- '<|im_start|>' + message.role + '\n' + content }}
- {%- endif %}
+ {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`qwen3_vl_training.jinja` Copy Section

Patched Qwen3-VL template. Diff vs qwen3_vl.jinja:

Wrap assistant message output (both content and tool_calls) with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`gptoss_training.jinja` Copy Section

Patched GPT-OSS template. Diff vs gptoss.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`idefics3_training.jinja` Copy Section

Patched Idefics3 template. Diff vs idefics3.jinja:

Split the assistant message into its own branch so the {% generation %} / {% endgeneration %} markers wrap the assistant content. This enables return_assistant_tokens_mask=True to produce correct masks for SFT assistant-only loss.

`lfm2_training.jinja` Copy Section

Patched LFM2 template. Diff vs lfm2.jinja:

Split the unified message output line into role-specific branches, so the <|im_start|>assistant\n prompt cue sits outside the generation block (it is not generated by the model), while the assistant's content and <|im_end|>\n (which the model must learn to produce and to stop on) sit inside. Wrap the assistant content with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`lfm2_v2_training.jinja` Copy Section

Patched LFM2 template, later revision. Diff vs lfm2_v2.jinja:

Wrap the assistant output with {% generation %} / {% endgeneration %}, in both the branch that renders tool_calls and the plain-content branch, so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss. The <|im_start|>assistant\n prompt cue stays outside the generation block, while the content, the tool calls and <|im_end|>\n sit inside.

`llama3_training.jinja` Copy Section

Patched Llama 3 template. Diff vs llama3.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`llava_next_training.jinja` Copy Section

Patched Llava-Next template. Diff vs llava_next.jinja:

• Wraps assistant message output with {% generation %} / {% endgeneration %} markers for SFT assistant-only loss.
• Keeps the space that follows </s> outside the generation block. The tokenizer merges it into the next turn's first token (▁[), which would otherwise be masked as assistant output and hide the </s> stop token from the end of the masked span.

`muse_glimmer_training.jinja` Copy Section

Patched Muse Glimmer template. Diff vs muse_glimmer.jinja: wrap the whole assistant branch — every ATEM channel of the turn — with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`nemotron_3_nano_training.jinja` Copy Section

Patched Nemotron Nano template. Diff vs nemotron_3_nano.jinja: the original is already prefix-preserving, so the only change is wrapping assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`nemotron_3_super_training.jinja` Copy Section

Patched Nemotron Super template. Diff vs nemotron_3_super.jinja: same as nemotron_3_nano_training.jinja — the original is already prefix-preserving, so the only change is wrapping assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`nemotron_3_ultra_training.jinja` Copy Section

Patched Nemotron Ultra template. Diff vs nemotron_3_ultra.jinja: same as nemotron_3_nano_training.jinja — the original is already prefix-preserving, so the only change is wrapping assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`nemotron_3_5_lightning_training.jinja` Copy Section

Patched Nemotron 3.5 Lightning template. Diff vs nemotron_3_5_lightning.jinja: same as nemotron_3_nano_training.jinja — the original is already prefix-preserving, so the only change is wrapping assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`phi3_training.jinja` Copy Section

Patched Phi-3 template. Diff vs phi3.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`phi3_5_training.jinja` Copy Section

Patched Phi-3.5 template. Diff vs phi3_5.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`qwen2_5_training.jinja` Copy Section

Patched Qwen2.5 template. Diff vs qwen2_5.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`qwen2_5_vl_training.jinja` Copy Section

Patched Qwen2.5-VL template (also used for Qwen2-VL, which ships a byte-identical template). Diff vs qwen2_5_vl.jinja:

Split the assistant message into its own branch so the {% generation %} / {% endgeneration %} markers wrap the assistant content. This enables return_assistant_tokens_mask=True to produce correct masks for SFT assistant-only loss.

`qwen3_instruct_2507_training.jinja` Copy Section

Patched Qwen3-Instruct-2507 template (used by models like Qwen3-4B-Instruct-2507, which ship a simpler Qwen3 variant without reasoning_content / `` parsing, multi_step_tool tracking, or the enable_thinking flag). Diff vs qwen3_instruct_2507.jinja:

Wrap assistant message output with {% generation %} / {% endgeneration %} so that return_assistant_tokens_mask=True produces correct masks for SFT assistant-only loss.

`qwen3_5_think_training.jinja` / `qwen3_5_nothink_training.jinja` Copy Section

Patched Qwen3.5 templates, shared logic across both flavors (they differ only in the default value of the enable_thinking flag — qwen3_5_think_training.jinja defaults to thinking enabled, used by Qwen3.5-4B and larger; qwen3_5_nothink_training.jinja defaults to thinking disabled, used by Qwen3.5-2B and smaller). Diff vs qwen3_5_think.jinja / qwen3_5_nothink.jinja: same set of changes as qwen3_training.jinja — require both and to be present before parsing, drop the loop.index0 > ns.last_query_index conditional so the thinking block is always emitted (prefix-preservation), and wrap assistant output with {% generation %} / {% endgeneration %} markers for SFT assistant-only loss.

`qwen3_6_training.jinja` Copy Section

Patched Qwen3.6 template. Diff vs qwen3_6.jinja: same set of changes as qwen3_training.jinja — require both and to be present before parsing, drop the loop.index0 > ns.last_query_index conditional so the thinking block is always emitted (prefix-preservation), and wrap assistant output with {% generation %} / {% endgeneration %} markers for SFT assistant-only loss.

`qwen3_8_training.jinja` Copy Section

Patched Qwen3.8 template. Diff vs qwen3_8.jinja: drop the preserve_thinking / loop.index0 > ns.last_query_index conditional so the thinking block is always emitted — prefix-preservation then holds even when the caller passes preserve_thinking=False — and wrap assistant output with {% generation %} / {% endgeneration %} markers for SFT assistant-only loss.

`smolvlm_training.jinja` Copy Section

Patched SmolVLM template (also used for SmolVLM2, which ships a byte-identical template). Diff vs smolvlm.jinja:

Split the assistant message into its own branch so the {% generation %} / {% endgeneration %} markers wrap the assistant content. This enables return_assistant_tokens_mask=True to produce correct masks for SFT assistant-only loss.

Related utilities Copy Chapter

See Chat Template Utilities for the helper functions ([clone_chat_template], [is_chat_template_prefix_preserving], [get_training_chat_template]) that operate on these templates.

Source Clis (`docs/source/clis.md`) Copy Section

Command Line Interfaces (CLIs)

TRL provides a powerful command-line interface (CLI) to fine-tune large language models (LLMs) using methods like Supervised Fine-Tuning (SFT), Direct Preference Optimization (DPO), and more. The CLI abstracts away much of the boilerplate, letting you launch training jobs quickly and reproducibly.

Commands Copy Chapter

Currently supported commands are:

Training Commands Copy Section

• trl distillation: distill a teacher model into a student with on-policy knowledge distillation
• trl dpo: fine-tune a LLM with DPO
• trl grpo: fine-tune a LLM with GRPO
• trl kto: fine-tune a LLM with KTO
• trl reward: train a Reward Model
• trl rloo: fine-tune a LLM with RLOO
• trl sft: fine-tune a LLM with SFT

Other Commands Copy Section

• trl env: get the system information
• trl vllm-serve: serve a model with vLLM (deprecated, use vllm serve directly, see vLLM integration)

Fine-Tuning with the TRL CLI Copy Chapter

Basic Usage Copy Section

You can launch training directly from the CLI by specifying required arguments like the model and dataset:

trl sft \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name stanfordnlp/imdb

trl dpo \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name anthropic/hh-rlhf

trl reward \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name trl-lib/ultrafeedback_binarized

trl grpo \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name HuggingFaceH4/Polaris-Dataset-53K \
  --reward_funcs accuracy_reward

trl rloo \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name HuggingFaceH4/Polaris-Dataset-53K \
  --reward_funcs accuracy_reward

trl kto \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name trl-lib/kto-mix-14k

Using Configuration Files Copy Section

To keep your CLI commands clean and reproducible, you can define all training arguments in a YAML configuration file:

# sft_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: stanfordnlp/imdb

trl sft --config sft_config.yaml

# dpo_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: anthropic/hh-rlhf

trl dpo --config dpo_config.yaml

# reward_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: trl-lib/ultrafeedback_binarized

trl reward --config reward_config.yaml

# grpo_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: HuggingFaceH4/Polaris-Dataset-53K
reward_funcs:
  - accuracy_reward

trl grpo --config grpo_config.yaml

# rloo_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: HuggingFaceH4/Polaris-Dataset-53K
reward_funcs:
  - accuracy_reward

trl rloo --config rloo_config.yaml

# kto_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: trl-lib/kto-mix-14k

trl kto --config kto_config.yaml

Scaling Up with Accelerate Copy Section

TRL CLI natively supports 🤗 Accelerate, making it easy to scale training across multiple GPUs, machines, or use advanced setups like DeepSpeed — all from the same CLI.

You can pass any accelerate launch arguments directly to trl, such as --num_processes. For more information see Using accelerate launch.

trl sft \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name stanfordnlp/imdb \
  --num_processes 4

# sft_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: stanfordnlp/imdb
num_processes: 4

trl sft --config sft_config.yaml

trl dpo \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name anthropic/hh-rlhf \
  --num_processes 4

# dpo_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: anthropic/hh-rlhf
num_processes: 4

trl dpo --config dpo_config.yaml

trl reward \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name trl-lib/ultrafeedback_binarized \
  --num_processes 4

# reward_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: trl-lib/ultrafeedback_binarized
num_processes: 4

trl reward --config reward_config.yaml

trl grpo \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name HuggingFaceH4/Polaris-Dataset-53K \
  --reward_funcs accuracy_reward \
  --num_processes 4

# grpo_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: HuggingFaceH4/Polaris-Dataset-53K
reward_funcs:
  - accuracy_reward
num_processes: 4

trl grpo --config grpo_config.yaml

trl rloo \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name HuggingFaceH4/Polaris-Dataset-53K \
  --reward_funcs accuracy_reward \
  --num_processes 4

# rloo_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: HuggingFaceH4/Polaris-Dataset-53K
reward_funcs:
  - accuracy_reward
num_processes: 4

trl rloo --config rloo_config.yaml

trl kto \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name trl-lib/kto-mix-14k \
  --num_processes 4

# kto_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: trl-lib/kto-mix-14k
num_processes: 4

trl kto --config kto_config.yaml

Using `--accelerate_config` for Accelerate Configuration Copy Section

The --accelerate_config flag lets you easily configure distributed training with 🤗 Accelerate. This flag accepts either:

• the name of a predefined config profile (built into TRL), or
• a path to a custom Accelerate YAML config file.

Predefined Config Profiles

TRL provides several ready-to-use Accelerate configs to simplify common training setups:

To use one of these, just pass the name to --accelerate_config. TRL will automatically load the corresponding config file from trl/accelerate_config/.

Example Usage

trl sft \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name stanfordnlp/imdb \
  --accelerate_config zero2  # or path/to/my/accelerate/config.yaml

# sft_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: stanfordnlp/imdb
accelerate_config: zero2  # or path/to/my/accelerate/config.yaml

trl sft --config sft_config.yaml

trl dpo \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name anthropic/hh-rlhf \
  --accelerate_config zero2  # or path/to/my/accelerate/config.yaml

# dpo_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: anthropic/hh-rlhf
accelerate_config: zero2  # or path/to/my/accelerate/config.yaml

trl dpo --config dpo_config.yaml

trl reward \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name trl-lib/ultrafeedback_binarized \
  --accelerate_config zero2  # or path/to/my/accelerate/config.yaml

# reward_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: trl-lib/ultrafeedback_binarized
accelerate_config: zero2  # or path/to/my/accelerate/config.yaml

trl reward --config reward_config.yaml

trl grpo \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name HuggingFaceH4/Polaris-Dataset-53K \
  --reward_funcs accuracy_reward \
  --accelerate_config zero2  # or path/to/my/accelerate/config.yaml

# grpo_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: HuggingFaceH4/Polaris-Dataset-53K
reward_funcs:
  - accuracy_reward
accelerate_config: zero2  # or path/to/my/accelerate/config.yaml

trl grpo --config grpo_config.yaml

trl rloo \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name HuggingFaceH4/Polaris-Dataset-53K \
  --reward_funcs accuracy_reward \
  --accelerate_config zero2  # or path/to/my/accelerate/config.yaml

# rloo_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: HuggingFaceH4/Polaris-Dataset-53K
reward_funcs:
  - accuracy_reward
accelerate_config: zero2  # or path/to/my/accelerate/config.yaml

trl rloo --config rloo_config.yaml

trl kto \
  --model_name_or_path Qwen/Qwen2.5-0.5B \
  --dataset_name trl-lib/kto-mix-14k \
  --accelerate_config zero2  # or path/to/my/accelerate/config.yaml

# kto_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: trl-lib/kto-mix-14k
accelerate_config: zero2  # or path/to/my/accelerate/config.yaml

trl kto --config kto_config.yaml

Using dataset mixtures Copy Section

You can use dataset mixtures to combine multiple datasets into a single training dataset. This is useful for training on diverse data sources or when you want to mix different types of data.

# sft_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
datasets:
  - path: stanfordnlp/imdb
  - path: roneneldan/TinyStories

trl sft --config sft_config.yaml

# dpo_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
datasets:
  - path: BAAI/Infinity-Preference
  - path: argilla/Capybara-Preferences

trl dpo --config dpo_config.yaml

# reward_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
datasets:
  - path: trl-lib/tldr-preference
  - path: trl-lib/lm-human-preferences-sentiment

trl reward --config reward_config.yaml

# grpo_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
datasets:
  - path: HuggingFaceH4/Polaris-Dataset-53K
  - path: trl-lib/DeepMath-103K
reward_funcs:
  - accuracy_reward

trl grpo --config grpo_config.yaml

# rloo_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
datasets:
  - path: HuggingFaceH4/Polaris-Dataset-53K
  - path: trl-lib/DeepMath-103K
reward_funcs:
  - accuracy_reward

trl rloo --config rloo_config.yaml

# kto_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
datasets:
  - path: trl-lib/kto-mix-14k
  - path: argilla/ultrafeedback-binarized-preferences-cleaned

trl kto --config kto_config.yaml

To see all the available keywords for defining dataset mixtures, refer to the [scripts.utils.DatasetConfig] and [DatasetMixtureConfig] classes.

Getting the System Information Copy Chapter

You can get the system information by running the following command:

trl env

This will print out the system information, including the GPU information, the CUDA version, the PyTorch version, the transformers version, the TRL version, and any optional dependencies that are installed.

Copy-paste the following information when reporting an issue:

- Platform: Linux-5.15.0-1048-aws-x86_64-with-glibc2.31
- Python version: 3.11.9
- PyTorch version: 2.4.1
- accelerator(s): NVIDIA H100 80GB HBM3
- Transformers version: 4.45.0.dev0
- Accelerate version: 0.34.2
- Accelerate config: 
  - compute_environment: LOCAL_MACHINE
  - distributed_type: DEEPSPEED
  - mixed_precision: no
  - use_cpu: False
  - debug: False
  - num_processes: 4
  - machine_rank: 0
  - num_machines: 1
  - rdzv_backend: static
  - same_network: True
  - main_training_function: main
  - enable_cpu_affinity: False
  - deepspeed_config: {'gradient_accumulation_steps': 4, 'offload_optimizer_device': 'none', 'offload_param_device': 'none', 'zero3_init_flag': False, 'zero_stage': 2}
  - downcast_bf16: no
  - tpu_use_cluster: False
  - tpu_use_sudo: False
  - tpu_env: []
- Datasets version: 3.0.0
- HF Hub version: 0.24.7
- TRL version: 0.12.0.dev0+acb4d70
- bitsandbytes version: 0.41.1
- DeepSpeed version: 0.15.1
- Diffusers version: 0.30.3
- Liger-Kernel version: 0.3.0
- LLM-Blender version: 0.0.2
- OpenAI version: 1.46.0
- PEFT version: 0.12.0
- vLLM version: not installed

This information is required when reporting an issue.

Source Community Tutorials (`docs/source/community_tutorials.md`) Copy Section

Community Tutorials

Community tutorials are made by active members of the Hugging Face community who want to share their knowledge and expertise with others. They are a great way to learn about the library and its features, and to get started with core classes and modalities.

Language Models Copy Chapter

Tutorials Copy Section

Videos Copy Section

⚠️ Deprecated features notice for "How to fine-tune a smol-LM with Hugging Face, TRL, and the smoltalk Dataset" (click to expand)

[!WARNING]
The tutorial uses two deprecated features:

• SFTTrainer(..., tokenizer=tokenizer): Use SFTTrainer(..., processing_class=tokenizer) instead, or simply omit it (it will be inferred from the model).
• setup_chat_format(model, tokenizer): Use SFTConfig(..., chat_template_path="Qwen/Qwen3-0.6B"), where chat_template_path specifies the model whose chat template you want to copy.

Vision Language Models Copy Chapter

Tutorials Copy Section

Speech Language Models Copy Chapter

Tutorials Copy Section

Contributing Copy Chapter

If you have a tutorial that you would like to add to this list, please open a PR to add it. We will review it and merge it if it is relevant to the community.

Source Cpo Trainer (`docs/source/cpo_trainer.md`) Copy Section

CPO Trainer

Overview Copy Chapter

Contrastive Preference Optimization (CPO) as introduced in the paper Contrastive Preference Optimization: Pushing the Boundaries of LLM Performance in Machine Translation by Haoran Xu, Amr Sharaf, Yunmo Chen, Weiting Tan, Lingfeng Shen, Benjamin Van Durme, Kenton Murray, and Young Jin Kim. At a high level, CPO trains models to avoid generating adequate, but not perfect, translations in Machine Translation (MT) tasks. However, CPO is a general approximation of the DPO loss and can be applied to other domains, such as chat.

CPO aims to mitigate two fundamental shortcomings of SFT. First, SFT’s methodology of minimizing the discrepancy between predicted outputs and gold-standard references inherently caps model performance at the quality level of the training data. Secondly, SFT lacks a mechanism to prevent the model from rejecting mistakes in translations. The CPO objective is derived from the DPO objective.

Quick start Copy Chapter

This example demonstrates how to train a model using the CPO method. We use the Qwen 0.5B model as the base model. We use the preference data from the UltraFeedback dataset. You can view the data in the dataset here:

Below is the script to train the model:

# train_cpo.py
from datasets import load_dataset
from trl.experimental.cpo import CPOConfig, CPOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
train_dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")

training_args = CPOConfig(output_dir="Qwen2-0.5B-CPO")
trainer = CPOTrainer(model=model, args=training_args, processing_class=tokenizer, train_dataset=train_dataset)
trainer.train()

Execute the script using the following command:

accelerate launch train_cpo.py

Expected dataset type Copy Chapter

CPO requires a preference dataset. The [experimental.cpo.CPOTrainer] supports both conversational and standard dataset formats. When provided with a conversational dataset, the trainer will automatically apply the chat template to the dataset.

Logged metrics Copy Chapter

While training and evaluating, we record the following metrics:

• rewards/chosen: the mean log probabilities of the policy model for the chosen responses scaled by beta
• rewards/rejected: the mean log probabilities of the policy model for the rejected responses scaled by beta
• rewards/accuracies: mean of how often the chosen rewards are > than the corresponding rejected rewards
• rewards/margins: the mean difference between the chosen and corresponding rejected rewards
• nll_loss: the mean negative log likelihood loss of the policy model for the chosen responses

CPO variants Copy Chapter

Simple Preference Optimization (SimPO) Copy Section

Simple Preference Optimization (SimPO) by Yu Meng, Mengzhou Xia, and Danqi Chen proposes a simpler and more effective preference optimization algorithm than DPO without using a reference model. The key designs in SimPO are (1) using length-normalized log likelihood as the implicit reward, and (2) incorporating a target reward margin in the Bradley-Terry ranking objective. The official code can be found at princeton-nlp/SimPO.

The abstract from the paper is the following:

Direct Preference Optimization (DPO) is a widely used offline preference optimization algorithm that reparameterizes reward functions in reinforcement learning from human feedback (RLHF) to enhance simplicity and training stability. In this work, we propose SimPO, a simpler yet more effective approach. The effectiveness of SimPO is attributed to a key design: using the average log probability of a sequence as the implicit reward. This reward formulation better aligns with model generation and eliminates the need for a reference model, making it more compute and memory efficient. Additionally, we introduce a target reward margin to the Bradley-Terry objective to encourage a larger margin between the winning and losing responses, further enhancing the algorithm's performance. We compare SimPO to DPO and its latest variants across various state-of-the-art training setups, including both base and instruction-tuned models like Mistral and Llama3. We evaluated on extensive instruction-following benchmarks, including AlpacaEval 2, MT-Bench, and the recent challenging Arena-Hard benchmark. Our results demonstrate that SimPO consistently and significantly outperforms existing approaches without substantially increasing response length. Specifically, SimPO outperforms DPO by up to 6.4 points on AlpacaEval 2 and by up to 7.5 points on Arena-Hard. Our top-performing model, built on Llama3-8B-Instruct, achieves a remarkable 44.7 length-controlled win rate on AlpacaEval 2 -- surpassing Claude 3 Opus on the leaderboard, and a 33.8 win rate on Arena-Hard -- making it the strongest 8B open-source model.

The SimPO loss is integrated in the [experimental.cpo.CPOTrainer], as it's an alternative loss that adds a reward margin, allows for length normalization, and does not use BC regularization. To use this loss, just turn on loss_type="simpo" and cpo_alpha=0.0 in the [experimental.cpo.CPOConfig] and set the simpo_gamma to a recommended value.

CPO-SimPO Copy Section

We also offer the combined use of CPO and SimPO, which enables more stable training and improved performance. Learn more details at CPO-SimPO GitHub. To use this method, simply enable SimPO by setting loss_type="simpo" and a non-zero cpo_alpha in the [experimental.cpo.CPOConfig].

AlphaPO Copy Section

The AlphaPO -- Reward shape matters for LLM alignment (AlphaPO) method by Aman Gupta, Shao Tang, Qingquan Song, Sirou Zhu, Jiwoo Hong, Ankan Saha, Viral Gupta, Noah Lee, Eunki Kim, Jason Zhu, Natesh Pillai, and S. Sathiya Keerthi is also implemented in the [experimental.cpo.CPOTrainer]. AlphaPO is an alternative method that applies a transformation to the reward function shape in the context of SimPO loss. The abstract from the paper is the following:

Reinforcement Learning with Human Feedback (RLHF) and its variants have made huge strides toward the effective alignment of large language models (LLMs) to follow instructions and reflect human values. More recently, Direct Alignment Algorithms (DAAs) have emerged in which the reward modeling stage of RLHF is skipped by characterizing the reward directly as a function of the policy being learned. Some popular examples of DAAs include Direct Preference Optimization (DPO) and Simple Preference Optimization (SimPO). These methods often suffer from likelihood displacement, a phenomenon by which the probabilities of preferred responses are often reduced undesirably. In this paper, we argue that, for DAAs the reward (function) shape matters. We introduce AlphaPO, a new DAA method that leverages an α-parameter to help change the shape of the reward function beyond the standard log reward. AlphaPO helps maintain fine-grained control over likelihood displacement and overoptimization. Compared to SimPO, one of the best performing DAAs, AlphaPO leads to about 7% to 10% relative improvement in alignment performance for the instruct versions of Mistral-7B and Llama3-8B while achieving 15% to 50% relative improvement over DPO on the same models. The analysis and results presented highlight the importance of the reward shape and how one can systematically change it to affect training dynamics, as well as improve alignment performance.

To use this loss as described in the paper, we can set the loss_type="alphapo" which automatically sets loss_type="simpo" and cpo_alpha=0.0, together with alpha and simpo_gamma to recommended values in the [experimental.cpo.CPOConfig]. Alternatively, you can manually set loss_type="simpo", cpo_alpha=0.0, together with alpha and simpo_gamma to recommended values. Other variants of this method are also possible, such as setting loss_type="ipo" and alpha to any non-zero value.

Loss functions Copy Chapter

The CPO algorithm supports several loss functions. The loss function can be set using the loss_type parameter in the [experimental.cpo.CPOConfig]. The following loss functions are supported:

For Mixture of Experts Models: Enabling the auxiliary loss Copy Section

MOEs are the most efficient if the load is about equally distributed between experts.
To ensure that we train MOEs similarly during preference-tuning, it is beneficial to add the auxiliary loss from the load balancer to the final loss.

This option is enabled by setting output_router_logits=True in the model config (e.g., [~transformers.MixtralConfig]).
To scale how much the auxiliary loss contributes to the total loss, use the hyperparameter router_aux_loss_coef=... (default: 0.001) in the model config.

CPOTrainer Copy Chapter

[[autodoc]] experimental.cpo.CPOTrainer
- train
- save_model
- push_to_hub

CPOConfig Copy Chapter

[[autodoc]] experimental.cpo.CPOConfig

Source Customization (`docs/source/customization.md`) Copy Section

Training customization

TRL is designed with modularity in mind so that users are able to efficiently customize the training loop for their needs. Below are examples on how you can apply and test different techniques.

[!NOTE]
Although these examples use the [DPOTrainer], these customization methods apply to most (if not all) trainers in TRL.

Use different optimizers and schedulers Copy Chapter

By default, the [DPOTrainer] creates a torch.optim.AdamW optimizer. You can create and define a different optimizer and pass it to [DPOTrainer] as follows:

from datasets import load_dataset
from torch import optim
from transformers import AutoModelForCausalLM
from trl import DPOTrainer

dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
optimizer = optim.SGD(model.parameters(), lr=1e-6)

trainer = DPOTrainer(
    model=model,
    train_dataset=dataset,
    optimizers=(optimizer, None),
)
trainer.train()

Add a learning rate scheduler Copy Section

You can also add learning rate schedulers by passing both optimizer and scheduler:

from torch import optim

optimizer = optim.AdamW(model.parameters(), lr=1e-6)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

trainer = DPOTrainer(..., optimizers=(optimizer, lr_scheduler))

Pass 8-bit reference models Copy Chapter

Since trl supports all keyword arguments when loading a model from transformers using from_pretrained, you can also leverage load_in_8bit from transformers for more memory efficient fine-tuning.

Read more about 8-bit model loading in transformers Load in 8bit or 4bit.

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(load_in_8bit=True)
ref_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", quantization_config=quantization_config)

trainer = DPOTrainer(..., ref_model=ref_model)

Add custom callbacks Copy Chapter

You can customize the training loop by adding callbacks for logging, monitoring, or early stopping. Callbacks allow you to execute custom code at specific points during training.

from transformers import TrainerCallback


class CustomLoggingCallback(TrainerCallback):
    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is not None:
            print(f"Step {state.global_step}: {logs}")


trainer = DPOTrainer(..., callbacks=[CustomLoggingCallback()])

Add custom evaluation metrics Copy Chapter

You can define custom evaluation metrics to track during training. This is useful for monitoring model performance on specific tasks.

def compute_metrics(eval_preds):
    logits, labels = eval_preds
    # Add your metric computation here
    return {"custom_metric": 0.0}


training_args = DPOConfig(..., eval_strategy="steps", eval_steps=100)

trainer = DPOTrainer(..., eval_dataset=eval_dataset, compute_metrics=compute_metrics)

Use mixed precision training Copy Chapter

Mixed precision training can significantly speed up training and reduce memory usage. You can enable it by setting bf16=True or fp16=True in the training config.

# Use bfloat16 precision (recommended for modern GPUs)
training_args = DPOConfig(..., bf16=True)

Note: Use bf16=True for Ampere GPUs (A100, RTX 30xx) or newer, and fp16=True for older GPUs.

Use gradient accumulation Copy Chapter

When training with limited GPU memory, gradient accumulation allows you to simulate larger batch sizes by accumulating gradients over multiple steps before updating weights.

# Simulate a batch size of 32 with per_device_train_batch_size=4 and gradient_accumulation_steps=8
training_args = DPOConfig(
    ...,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
)

Change the training objective Copy Chapter

Subclassing a trainer is the simplest way to customize training. The loss is defined in the compute_loss method, so to train with a different objective, subclass the trainer and override it. Data preparation, generation, logging, and checkpointing are inherited, so the subclass holds only what actually changes.

import torch

from trl import DPOTrainer
from trl.trainer.utils import selective_log_softmax


class MyDPOTrainer(DPOTrainer):
    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        # In this example: the hinge loss from https://huggingface.co/papers/2309.06657
        input_ids, attention_mask = inputs["input_ids"], inputs["attention_mask"]
        shift_labels, shift_completion_mask = input_ids[:, 1:], inputs["completion_mask"][:, 1:]
        logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
        with torch.no_grad():
            ref_logits = self.ref_model(input_ids=input_ids, attention_mask=attention_mask).logits

        # Sum the log-probs over the completion tokens, for the policy and for the reference model
        logps = (selective_log_softmax(logits[:, :-1], shift_labels) * shift_completion_mask).sum(dim=1)
        ref_logps = (selective_log_softmax(ref_logits[:, :-1], shift_labels) * shift_completion_mask).sum(dim=1)

        chosen_logps, rejected_logps = logps.chunk(2, dim=0)  # batch is [chosen, rejected]
        ref_chosen_logps, ref_rejected_logps = ref_logps.chunk(2, dim=0)
        delta = (chosen_logps - ref_chosen_logps) - (rejected_logps - ref_rejected_logps)
        return torch.relu(1 - self.beta * delta).mean()

Add parameters to the config Copy Section

Subclass the config to declare the parameters your loss needs:

from dataclasses import dataclass, field

from trl import DPOConfig


@dataclass
class MyDPOConfig(DPOConfig):
    my_coef: float = field(default=0.1, metadata={"help": "Coefficient of the custom term."})

Change the batch format Copy Section

When the loss needs inputs that the default collator doesn't produce, pass your own collator, and override _prepare_dataset to leave the dataset untouched when it is already in the expected format:

class MyDPOTrainer(DPOTrainer):
    def _prepare_dataset(self, dataset, *args, **kwargs):
        return dataset


trainer = MyDPOTrainer(..., data_collator=my_collator)

Complete example Copy Section

The block-diffusion SFT example examples/sft_diffusion_gemma/sft_diffusion_gemma.py combines the three: DiffusionGemmaSFTConfig extends [SFTConfig] with the canvas and corruption parameters, and DiffusionGemmaSFTTrainer extends [SFTTrainer] and replaces the autoregressive cross-entropy with a block-diffusion denoising objective. Neither requires a change to the library.

Antidoom, an open-source tool from Liquid AI, extends [DPOTrainer] with Final Token Preference Optimization (Antislop), a preference loss over a single token position that reduces repetition loops in reasoning models.

Source Data Utils (`docs/source/data_utils.md`) Copy Section

Data Utilities

is_conversational Copy Chapter

[[autodoc]] is_conversational

maybe_convert_to_chatml Copy Chapter

[[autodoc]] maybe_convert_to_chatml

extract_prompt Copy Chapter

[[autodoc]] extract_prompt

unpair_preference_dataset Copy Chapter

[[autodoc]] unpair_preference_dataset

Source Dataset Formats (`docs/source/dataset_formats.md`) Copy Section

Dataset formats and types

This guide provides an overview of the dataset formats and types supported by each trainer in TRL.

Overview of the dataset formats and types Copy Chapter

• The format of a dataset refers to how the data is structured, typically categorized as either standard or conversational.
• The type is associated with the specific task the dataset is designed for, such as prompt-only or preference. Each type is characterized by its columns, which vary according to the task, as shown in the table.

{"role": "assistant", "content": "It is blue."}]} |

| Prompt-only | {"prompt": "The sky is"} | {"prompt": [{"role": "user", "content": "What color is the sky?"}]} |
| Prompt-completion | {"prompt": "The sky is",
"completion": " blue."} | {"prompt": [{"role": "user", "content": "What color is the sky?"}],
"completion": [{"role": "assistant", "content": "It is blue."}]} |
| Preference | {"prompt": "The sky is",
"chosen": " blue.",
"rejected": " green."}
or, with implicit prompt:
{"chosen": "The sky is blue.",
"rejected": "The sky is green."} | {"prompt": [{"role": "user", "content": "What color is the sky?"}],
"chosen": [{"role": "assistant", "content": "It is blue."}],
"rejected": [{"role": "assistant", "content": "It is green."}]}
or, with implicit prompt:
{"chosen": [{"role": "user", "content": "What color is the sky?"},
{"role": "assistant", "content": "It is blue."}],
"rejected": [{"role": "user", "content": "What color is the sky?"},
{"role": "assistant", "content": "It is green."}]} |

Formats Copy Section

Standard

The standard dataset format typically consists of plain text strings. The columns in the dataset vary depending on the task. This is the format expected by TRL trainers. Below are examples of standard dataset formats for different tasks:

# Language modeling
language_modeling_example = {"text": "The sky is blue."}
# Preference
preference_example = {"prompt": "The sky is", "chosen": " blue.", "rejected": " green."}
# Unpaired preference
unpaired_preference_example = {"prompt": "The sky is", "completion": " blue.", "label": True}

Conversational

Conversational datasets are used for tasks involving dialogues or chat interactions between users and assistants. Unlike standard dataset formats, these contain sequences of messages where each message has a role (e.g., "user" or "assistant") and content (the message text).

messages = [
    {"role": "user", "content": "Hello, how are you?"},
    {"role": "assistant", "content": "I'm doing great. How can I help you today?"},
    {"role": "user", "content": "I'd like to show off how chat templating works!"},
]

Just like standard datasets, the columns in conversational datasets vary depending on the task. Below are examples of conversational dataset formats for different tasks:

# Prompt-completion
prompt_completion_example = {"prompt": [{"role": "user", "content": "What color is the sky?"}],
                             "completion": [{"role": "assistant", "content": "It is blue."}]}
# Preference
preference_example = {
    "prompt": [{"role": "user", "content": "What color is the sky?"}],
    "chosen": [{"role": "assistant", "content": "It is blue."}],
    "rejected": [{"role": "assistant", "content": "It is green."}],
}

Tool Calling

Some chat templates support tool calling, which allows the model to interact with external functions—referred to as tools—during generation. This extends the conversational capabilities of the model by enabling it to output a "tool_calls" field instead of a standard "content" message whenever it decides to invoke a tool.

After the assistant initiates a tool call, the tool executes and returns its output. The assistant can then process this output and continue the conversation accordingly.

Here’s a simple example of a tool-calling interaction:

messages = [
    {"role": "user", "content": "Turn on the living room lights."},
    {"role": "assistant", "tool_calls": [
        {"type": "function", "function": {
            "name": "control_light",
            "arguments": {"room": "living room", "state": "on"}
        }}]
    },
    {"role": "tool", "name": "control_light", "content": "The lights in the living room are now on."},
    {"role": "assistant", "content": "Done!"}
]

When preparing datasets for Supervised Fine-Tuning (SFT) with tool calling, it is important that your dataset includes an additional column named tools. This column contains the list of available tools for the model, which is usually used by the chat template to construct the system prompt.

The tools must be specified in a codified JSON schema format. You can automatically generate this schema from Python function signatures using the [~transformers.utils.get_json_schema] utility:

import json
from transformers.utils import get_json_schema

def control_light(room: str, state: str) -> str:
    """
    Controls the lights in a room.

    Args:
        room: The name of the room.
        state: The desired state of the light ("on" or "off").

    Returns:
        str: A message indicating the new state of the lights.
    """
    return f"The lights in {room} are now {state}."

# Generate JSON schema
json_schema = get_json_schema(control_light)

The generated schema would look like:

{"type": "function", "function": {"name": "control_light", "description": "Controls the lights in a room.", "parameters": {"type": "object", "properties": {"room": {"type": "string", "description": "The name of the room."}, "state": {"type": "string", "description": "The desired state of the light (\"on\" or \"off\")."}}, "required": ["room", "state"]}, "return": {"type": "string", "description": "str: A message indicating the new state of the lights."}}}

A complete dataset entry for SFT might look like:

{"messages": messages, "tools": [json_schema]}

To get a Dataset you need to use the Json() type for tool arguments since they are arbitrary JSON objects, and not dictionaries with fixed fields and types: