From 681c16e031c21761b005bfa8fa489d7335d93e60 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 12 Jun 2026 09:38:58 +0000 Subject: [PATCH] Restore UNSLOTH_RETURN_LOGITS after prediction_step instead of forcing 0 unsloth_prediction_step forces UNSLOTH_RETURN_LOGITS=1 to materialize logits during evaluation, then hardcoded it back to 0 afterwards. That silently turned off an explicit user setting: a user who sets UNSLOTH_RETURN_LOGITS=1 before training (the documented way to keep real logits) had it reset to 0 after the first eval step. Save the prior value before forcing 1 and restore it after. Verified: with an explicit '1' the value is preserved across trainer.evaluate(); with the default it is restored to '0' as before. --- unsloth/models/rl.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 5fb8cf6ddc..a3808fe660 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -276,6 +276,9 @@ def PatchRL(FastLanguageModel): else: labels = None + # Force logits during eval, but restore the user's prior setting after + # so an explicit UNSLOTH_RETURN_LOGITS="1" is not silently turned off. + _old_return_logits = os.environ.get("UNSLOTH_RETURN_LOGITS", "0") os.environ["UNSLOTH_RETURN_LOGITS"] = "1" with torch.no_grad(): if has_labels or loss_without_labels: @@ -315,7 +318,7 @@ def PatchRL(FastLanguageModel): # TODO: this needs to be fixed and made cleaner later. if self.args.past_index >= 0: self._past = outputs[self.args.past_index - 1] - os.environ["UNSLOTH_RETURN_LOGITS"] = "0" + os.environ["UNSLOTH_RETURN_LOGITS"] = _old_return_logits if prediction_loss_only: return (loss, None, None)