#7063·trl

[Tracking] Move off liger-kernel and make the memory-efficient loss the default

Author: qgallouedecCreated Sep 4, 2026Updated Sep 17, 2026

Context

use_liger_kernel=True does two unrelated things in TRL today:

  1. transformers' Trainer patches the model with Liger's Triton kernels (RMSNorm, SwiGLU, RoPE, fused linear cross-entropy). This is the SFT speedup.
  2. DPO, KTO, GRPO and the distillation trainers swap their loss for Liger's chunked_loss fused losses, which compute the loss from the hidden states without materializing the (batch, seq_len, vocab) logits.

#7059 vendors the second part into trl/losses/ as-is. This issue tracks what comes after.

Why

  • chunked_loss is pure PyTorch (no Triton) and mostly TRL-authored. Upstream stopped reviewing our PRs (6 open since June, several without a single comment), so we now own the code anyway.
  • The vendored losses are a second copy of each trainer's loss math. That is where the bugs live: Liger#1082 (GRPO grad_norm off vs the default path), KTO ignoring the reference under PEFT, KTO class weights and DPO label_smoothing silently dropped, chunk_size > 1 broadcast. Every one of them is a drift between the two copies.
  • The fused path carries 13 restrictions (single loss_type, no f-divergence, no entropy bonus, no off-policy masking, ...) that exist only because the loss math was reimplemented.
  • All DPO/KTO/GRPO actually need is per-token log-probs of the selected tokens. TRL already has a primitive for that: _ChunkedLogProbFunction in trl/trainer/utils.py (vocab-streamed online logsumexp, recompute in backward, used by async and zero-sync GRPO). SFT already made the equivalent (loss_type="chunked_nll") its default with neutral wall time (#5575).
  • The Triton half is already covered by kernels-community/liger-kernels + use_kernels=True in transformers.

Destination

  • In every trainer, the per-token log-probs come from backbone -> _ChunkedLogProbFunction, then the trainer's existing loss code runs unchanged. One implementation of each loss.
  • That path is the default. No trl/losses/, no fused loss objects, no flag to select a fused path.
  • Distillation keeps its token-chunked JSD (it needs full distributions) and drops the fused JSD, which duplicates it.
  • use_liger_kernel only ever meant "model kernels". It is deprecated in transformers in favour of use_kernels=True, and TRL drops the liger extra.

Plan

Each item is one PR, independently reviewable, ordered smallest first.

  • #7059: vendor liger_kernel.chunked_loss into trl/losses/, import swap only.
  • #7062 (kashif): KTO reference under PEFT, KTO class-weight guard, DPO label_smoothing / discopop_tau, chunk_size > 1 broadcast. Small fixes on the vendored code while it exists.
  • #7064 (kashif): DistillationTrainer: delete the fused JSD branch. The chunked path is already the default and does more. Pure deletion.
  • #7065 (kashif): Harden _ChunkedLogProbFunction: lm_head bias, re-gather the weight in backward under ZeRO-3 (today it reads the saved weight, which is partitioned once the gather context exits; tests only pass because tiny models stay below stage3_param_persistence_threshold), unit tests against full logits.
    • #7076 (kashif): Skip LM-head gradient buffers and GEMMs when its parameters are frozen.
  • #7077 (kashif): GRPO: fused path on _ChunkedLogProbFunction in _get_per_token_logps_and_entropies, _compute_loss untouched. Deletes grpo_loss.py / fused_linear_ppo.py, lifts 4 restrictions, adds a loss + gradient equivalence test.
  • #7243 (kashif): DPO: same, on _compute_loss. Deletes dpo_loss.py / fused_linear_preference.py.
  • #7075 (kashif): KTO: same, including the KL log-probs path. Deletes kto_loss.py / fused_linear_unpaired_preference.py.
  • Experimental GKD, GOLD, SDPO, SDFT, IW-OPD: drop the fused JSD like DistillationTrainer. Deletes jsd_loss.py / fused_linear_distillation.py, i.e. all of trl/losses/.
  • Benchmark the streamed path vs full logits, then make it the default in DPO/KTO/GRPO. Keep the full-logits path behind a flag only if the benchmark says so (and for DPO compute_metrics, which needs logits).
  • transformers: deprecate use_liger_kernel in favour of use_kernels=True. Check the layer coverage of the Hub mapping against liger's per-model patches first.
  • TRL: drop the liger extra, is_liger_kernel_available, liger_kernel_integration.md; point SFT users at model_init_kwargs={"use_kernels": True}.

Notes

Performance:

  • GRPO loses nothing (Liger's GRPO base already streams the vocab and recomputes in backward).
  • DPO/KTO pay one extra lm_head GEMM pass in backward, roughly 1/3 of the head's cost, a few percent of the step on an 8B/150k-vocab model, more on small models with big vocabularies. In exchange the loss-side peak drops from a full (T, V) fp32 logits chunk to one (N, 8192) tile. The benchmark step decides whether that trade needs a torch.compile on the epilogue.