TRL v1.15 ships fused LM head cutting peak memory up to 82% and enabling 7x longer sequences

Hugging Face's TRL library releases version 1.15 with what maintainers call the largest optimization in its history: a fused language-model head now enabled by default across every RL and distillation trainer. Instead of materializing a full [batch, sequence, vocab] logits tensor, a Triton kernel tiles the projection of hidden states through the LM head and directly computes log-probabilities and per-token entropy. The practical result: up to 82% peak VRAM reduction and the ability to train sequences seven times longer on the same hardware.

The change applies to SFT, DPO, KTO, GRPO, RLOO, and Distillation. The new kernel replaces three prior compute paths — the full-logits path, the Liger Kernel chunked path, and the _forward_redirection mechanism — with a single pass that emits only what each trainer needs: log-probs, entropy, and optionally log_sum_sq_probs for WPO weighting, mean_logits for chosen/rejected pairs, and is_top1 for average token accuracy in DPO. All auxiliary outputs are validated against full logits in automated tests. The sole requirement is Triton on GPU with CUDA, ROCm, or XPU under Linux.

Benchmarks ran on a single B300 with the PyTorch allocator capped at 79 GiB (matching usable H100 80GB memory), using google/gemma-3-1b-pt with random bfloat16 weights, a 262K-token vocabulary, batch size 1, gradient checkpointing, and SDPA. At 8,192 tokens, peak memory fell 52–82% and step time improved 2.3–10.9%. Maximum trainable length jumped in DPO from 10K to 59K tokens and in GRPO from 29K to 115K tokens. Gains are largest on small models with large vocabularies; SFT was unchanged because it already used chunked_nll to avoid full logits.

Breaking changes and migration: the four RL trainers (GRPO, RLOO, DPO, KTO) drop the legacy paths. A PEFT adapter on lm_head will now error; the fix is modules_to_save=["lm_head"]. The flag use_liger_kernel=True is deprecated and will be removed in v2.0.0; Liger kernels still apply, but fused_linear_cross_entropy is forced to False and an explicit setting will error. Migrate to model_init_kwargs={"use_kernels": True}. The compute_metrics and compute_loss callbacks with return_outputs=True still receive full logits via an extra forward pass, but both now warn that this will be removed in v2.0.0.

A separate addition in this release: selective checkpointing in SFT. When gradient checkpointing is enabled, the attention output is now saved during the forward pass instead of being recomputed in the backward pass. This recovers most of the checkpointing slowdown in long-context training at the cost of one additional hidden tensor per layer. The approach mirrors the Eager SAC used by torchtitan under FSDP2 and does not require torch.compile.