ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Restore UNSLOTH_RETURN_LOGITS after prediction_step instead of forcing 0

Browse source 
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.
This commit is contained in:
Daniel Han 2026-06-12 09:38:58 +00:00
commit 681c16e031
1 changed files with 4 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

5
unsloth/models/rl.py
View file

@ -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)