From 480240ce2429c0f5e6181418e334538feea11c41 Mon Sep 17 00:00:00 2001 From: Daniel Shen Date: Mon, 6 Apr 2026 14:10:26 -0700 Subject: [PATCH] fix: liger unnecessarily materializes logits in VRAM during eval, causing OOM --- src/transformers/trainer.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 235189fe8320..f18dfc4d46ae 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -2951,6 +2951,8 @@ def prediction_step( if has_labels or loss_without_labels: with self.compute_loss_context_manager(): num_items_in_batch = self._get_num_items_in_batch([inputs], self.args.device) + if self.args.use_liger_kernel and prediction_loss_only: + inputs = {**inputs, "skip_logits": True} loss, outputs = self.compute_loss( model, inputs, return_outputs=True, num_items_in_batch=num_items_in_batch )