trl

Features

🔥 Mistral Large 4: fine-tune a 1T-parameter model on one node

TRL now supports mistralai/Mistral-Large-4-1T-A52B: 1T parameters, 52B active. The new LoRA example fine-tunes it with SFTTrainer on a single 8×B300 node, on top of its FP8 weights: the experts are split 8 ways with token dispatch, everything else is sharded across the 8 GPUs at load time by transformers' DistributedConfig, and LoRA trains attention and the dense and shared-expert MLPs.

Measured: 32.7k tokens per step in ~8 s, 142 GiB per GPU, with packing and assistant-only loss on reasoning conversations.

https://github.com/user-attachments/assets/5aa86ec7-6444-4bed-bfdc-0445e99bea0d

training_args = SFTConfig(
    model_init_kwargs={
        "dtype": "auto",  # keep the FP8 weights
        "distributed_config": DistributedConfig(fsdp_size=8, ep_size=8),
    },
    packing=True,
    assistant_only_loss=True,
    ...
)
trainer = SFTTrainer(model="mistralai/Mistral-Large-4-1T-A52B", args=training_args, peft_config=LoraConfig(...), ...)

assistant_only_loss=True works out of the box: TRL ships a training chat template that matches the released checkpoint's, including its reasoning_effort levels (low, medium, high). Needs transformers>=5.22.0, peft>=0.21.0 and accelerate>=1.15.0.

• Add Mistral Large 4 by @qgallouedec in https://github.com/huggingface/trl/pull/8017
• Sync the Mistral Large 4 chat template with the released checkpoint by @albertvillanova in https://github.com/huggingface/trl/pull/8018

🎯 Unbiased importance sampling under top-k, top-p and min-p

vLLM normalizes its log-probabilities over the tokens left after top_k, top_p and min_p filtering, while the trainer normalized over the whole vocabulary. Every sampled token's log-ratio was off by the log of the probability mass the filter kept, so the vLLM importance sampling ratio was biased towards 0 (around 1e-4 instead of 1 with top_p=0.9). The trainers now normalize over the candidate set each token was sampled from, following Keep Sampling Mask from DeepSeek-V3.2.

• AsyncGRPO: renormalize log-probs over the top-k/top-p candidate set by @AmineDiro in https://github.com/huggingface/trl/pull/7627
• GRPO: renormalize log-probs over the top-k/top-p candidate set (vLLM colocate mode, vLLM >= 0.28; server mode will follow) by @albertvillanova in https://github.com/huggingface/trl/pull/7866
• Require top_k with min_p in AsyncGRPO so the importance sampling ratio stays unbiased by @albertvillanova in https://github.com/huggingface/trl/pull/7937
• Track the support maximum in the chunked log-prob kernel, so a far-below-max candidate set no longer gives a log-probability of 0 and non-finite gradients, by @kashif in https://github.com/huggingface/trl/pull/7983

💥 What this changes for you. vLLM bounds the candidate set it returns with top_k, so with vLLM and the importance sampling correction on (the default), GRPOConfig now raises when top_p < 1.0 or min_p > 0.0 is set without top_k > 0. Set top_k (e.g. top_k=512), or vllm_importance_sampling_correction=False.

⚖️ Score centering is the default AsyncGRPO objective

AsyncGRPO's old log-probs are the vLLM sampler's, so its ratio mixes staleness with the numerical mismatch between the server and the trainer. With an FP8 sampler and a bf16 trainer, the clipped objective collapses: reward goes from 0.60 to 0.34 (FP8 weights) or 0.10 (FP8 weights + KV cache) by step 300, against 0.71 with a bf16 sampler.

[Score centering](htt

Features

🛠️ TRL now trains on AWS Trainium

The fused LM head from v1.15 is what made Neuron tractable; the remaining gaps are closed here. There is a guide to go with it.

Neuron recompiles once per input shape, so an unpadded run stalls about a minute on every new sequence length. pad_to_multiple_of now defaults to "auto": 128 on Neuron, None elsewhere, so nothing changes on GPU. Trainium2, Qwen3-0.6B SFT, 20 steps: 130 s with the default against 260 s without.

• Add a guide for training on AWS Trainium by @qgallouedec in https://github.com/huggingface/trl/pull/7729
• Default pad_to_multiple_of to "auto" by @qgallouedec in https://github.com/huggingface/trl/pull/7724
• Run the fused LM head and distillation kernels on devices without Triton by @qgallouedec in https://github.com/huggingface/trl/pull/7669
• Report the Neuron device in trl env and telemetry by @qgallouedec in https://github.com/huggingface/trl/pull/7672
• Support any accelerator in activation offloading by @qgallouedec in https://github.com/huggingface/trl/pull/7670

🌳 Tree attention for AsyncGRPO

AsyncGRPO now packs every micro-batch row as a prefix tree instead of a flat sequence. Multi-turn and multi-generation rollouts repeat a lot of prefix: on a shared prompt, those tokens used to be forwarded once per row. Tree packing forwards them once.

rollout 0:  [ prompt ] [ completion A ]              prompt ──┬── completion A ── later turn
rollout 1:  [ prompt ] [ completion B ]        =>             └── completion B
rollout 2:  [ prompt ] [ completion A ] [ … ]

Chains without a fork collapse into segments, and each segment's queries attend to the K/V gathered along its root-to-segment path in a single FA3 varlen call. No mask tensor: every layer except attention runs once per unique token, and only K/V is duplicated along shared prefixes.

Qwen3-1.7B, 32K tokens per rollout, forward + backward:

There is no packing option to choose: with nothing shared a tree row is exactly the old sequence packing, at a ~1% cost.

by @AmineDiro in https://github.com/huggingface/trl/pull/7328

Muse Spark Open

Support for meta-models/Muse-Spark-Open-816B-A42B: the tiny-model script, coverage in the VLM tests (SFT, DPO, KTO, chat-template utils, fused LM head), and a LoRA SFT example on 8 nodes of 8×H100 with DistributedConfig(fsdp_size=64, ep_size=8), token dispatch and packing. Measured at 14.2 s/step, 131k tokens/step, 40.2 GB per GPU.

Needs transformers>=5.20.0; the adapter save additionally needs transformers#49500 with peft 0.21.2.

by @qgallouedec in https://github.com/huggingface/trl/pull/7536

AMD GPUs in the async trainers

• Support AMD GPUs in the async GRPO and async distillation trainers by @Abdennacer-Badaoui in https://github.com/huggingface/trl/pull/7689
• Hide the GPU from the async rollout worker on ROCm by @Abdennacer-Badaoui in https://github.com/huggingface/trl/pull/7685

Tool-call arguments typed by the tool's JSON schema

parse_response now receives the tools, so tool-call arguments come back typed from the tool's own schema instead of whatever the raw parse pr

Features

⚡ Fused LM head: up to 6.9× longer sequences on the same GPU

SFT, DPO, KTO, GRPO, RLOO and Distillation now score tokens with a fused LM head: a Triton kernel projects the hidden states through the LM head in tiles and reduces to per-token log-probs and entropy directly, so the [batch, seq, vocab] logits tensor is never built. It is on by default, there is nothing to turn on.

Max trainable sequence length (gemma-3-1b, 262k vocabulary, 79 GiB of GPU memory, bf16, batch size 1, gradient checkpointing, sdpa):

At 8,192 tokens, peak memory drops 52% to 82% and steps are 2.3% to 10.9% faster:

The gain is largest for small models with large vocabularies. SFT's previous default (chunked_nll) already avoided full logits, so it is unchanged. Scoring needs Triton on a GPU (Linux with CUDA, ROCm or XPU).

💥 What this changes for you. The full-logits scoring path, the use_liger_kernel chunked path and _forward_redirection are gone from GRPO, RLOO, DPO and KTO:

• A PEFT adapter on lm_head now raises. Use modules_to_save=["lm_head"] instead.
• use_liger_kernel=True is deprecated in these four trainers and will be removed in v2.0.0. Liger's layer kernels still apply, but fused_linear_cross_entropy is forced to False (it would replace the fused head) and setting it explicitly raises. Use model_init_kwargs={"use_kernels": True}.
• In DPO and KTO, compute_metrics and compute_loss(..., return_outputs=True) still receive the full logits, from an extra forward pass taken only when they are used. Both warn that this is deprecated and will be removed in v2.0.0.

Nothing else the trainers logged is lost: the kernel gained log_sum_sq_probs (WPO weighting), mean_logits (logits/chosen, logits/rejected) and is_top1 (DPO's mean_token_accuracy) as opt-in outputs, each checked against the full logits in tests.

1× B300 with the PyTorch allocator capped at 79 GiB (about an H100 80GB's usable memory), google/gemma-3-1b-pt built from its config with random bf16 weights, lr 1e-6, AdamW, synthetic token ids at exact length. Max length: one fresh process per attempt, 2 training steps, doubling from 8k then bisection to 1024 tokens, arms interleaved round-robin in the same job. Throughput: 12 steps, first 2 discarded, median of 3 A/B/A/B repeats in the same job. Noise floor (A vs A) is 0.1% to 0.8%, below every gap reported. GRPO/RLOO replace generation with fixed random completions, so only scoring is measured.

Not measured: H100 hardware itself, flash-attention (sdpa only), real generation, multi-GPU.

• Compute the chunked log-probabilities with a Triton kernel by @qgallouedec in https://github.com/huggingface/trl/pull/7386
• Add an opt-in fused LM head to the patched forward by @qgallouedec in https://github.com/huggingface/trl/pull/7387
• 💥 Score GRPO and RLOO tokens with the fused LM head by @qgallouedec in https://github.com/huggingface/trl/pull/7389
• 💥 Score DPO and KTO tokens with the fused LM head by @qgall

Patch release fixing two cases of silently wrong training and three crashes.

Worth calling out: #7505 affects GRPO, RLOO and Distillation on the default generation path (use_vllm=False). Models that end a turn with an id only their generation config declares (Gemma 3/4, Phi-3.5, ...) did not stop at the end of the turn, and everything generated after it was trained on, up to max_completion_length. The only visible sign was completions/clipped_ratio close to 1. With vLLM, generation stopped correctly but the metrics were still wrong, and mask_truncated_completions=True dropped every finished completion. Runs on models whose tokenizer eos is their only end-of-turn id are unchanged.

What's Changed

• Keep the DFT loss finite on a batch with no trainable tokens by @albertvillanova in https://github.com/huggingface/trl/pull/7471
• Fix precompute_ref_log_probs under FSDP with no reference model by @qgallouedec in https://github.com/huggingface/trl/pull/7507
• Fix: Start the completion where the tokenized prompt and prompt+completion diverge by @qgallouedec in https://github.com/huggingface/trl/pull/7463
• End completions on every eos id the model declares in GRPO, RLOO and Distillation by @albertvillanova in https://github.com/huggingface/trl/pull/7505
• Guard the requests and urllib3 imports in the vLLM client by @qgallouedec in https://github.com/huggingface/trl/pull/7545

Full Changelog: https://github.com/huggingface/trl/compare/v1.14.1...v1.14.2