[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:
- transformers'
Trainerpatches the model with Liger's Triton kernels (RMSNorm, SwiGLU, RoPE, fused linear cross-entropy). This is the SFT speedup. - DPO, KTO, GRPO and the distillation trainers swap their loss for Liger's
chunked_lossfused 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_lossis 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_smoothingsilently dropped,chunk_size > 1broadcast. 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:
_ChunkedLogProbFunctionintrl/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=Truein 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_kernelonly ever meant "model kernels". It is deprecated in transformers in favour ofuse_kernels=True, and TRL drops theligerextra.
Plan
Each item is one PR, independently reviewable, ordered smallest first.
- #7059: vendor
liger_kernel.chunked_lossintotrl/losses/, import swap only. - #7062 (kashif): KTO reference under PEFT, KTO class-weight guard, DPO
label_smoothing/discopop_tau,chunk_size > 1broadcast. 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_headbias, 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 belowstage3_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
_ChunkedLogProbFunctionin_get_per_token_logps_and_entropies,_compute_lossuntouched. Deletesgrpo_loss.py/fused_linear_ppo.py, lifts 4 restrictions, adds a loss + gradient equivalence test. - #7243 (kashif): DPO: same, on
_compute_loss. Deletesdpo_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 oftrl/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_kernelin favour ofuse_kernels=True. Check the layer coverage of the Hub mapping against liger's per-model patches first. - TRL: drop the
ligerextra,is_liger_kernel_available,liger_kernel_integration.md; point SFT users atmodel_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_headGEMM 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 atorch.compileon the epilogue.
Source: huggingface/trl