logits

unsloth/models/_utils.py

@@ -54,6 +54,7 @@ __all__ = [
     "unpatch_gradient_checkpointing",
 
     "HAS_CUT_CROSS_ENTROPY",
+    "EMPTY_LOGITS",
     "fused_linear_cross_entropy",
     "patch_unsloth_smart_gradient_checkpointing",
     "unpatch_unsloth_smart_gradient_checkpointing",
@@ -1164,3 +1165,27 @@ def unsloth_compile_transformers(
     pass
     return model_types
 pass
+
+# We need an empty logits flag to warn people logits will not be returned anymore unless asked ie
+# os.environ['UNSLOTH_RETURN_LOGITS'] = '1'
+LOGITS_ERROR_STRING = \
+    "Unsloth: Logits are empty from 2024.11 onwards. To get raw logits again, please "\
+    'set the environment variable `UNSLOTH_RETURN_LOGITS` to `"1" BEFORE starting to train ie before `trainer.train()`. For example:\n\n'\
+    "import os\n"\
+    "os.environ['UNSLOTH_RETURN_LOGITS'] = '1'\n"\
+    "... trainer.train() ..."
+
+def raise_logits_error(*args, **kwargs): raise NotImplementedError(LOGITS_ERROR_STRING)
+class EmptyLogits(torch.Tensor):
+    def __init__(self): return
+    __getitem__ = raise_logits_error
+    __getattr__ = raise_logits_error
+    def __repr__(self): return LOGITS_ERROR_STRING
+    def __str__ (self): return LOGITS_ERROR_STRING
+pass
+EMPTY_LOGITS = EmptyLogits()
+functions = dir(torch.Tensor)
+for function in functions:
+    try: exec(f"EMPTY_LOGITS.{function} = raise_logits_error", globals(), locals())
+    except: continue
+pass

unsloth/models/llama.py

@@ -974,8 +974,6 @@ def CausalLM_fast_forward(fast_forward_inference):
         logit_softcapping = getattr(self.config, "final_logit_softcapping", 0)
         logit_scaling     = getattr(self.config, "logit_scale", 0)
 
-        print(kwargs)
-
         if bsz == 1 and q_len == 1:
             logits = torch.mv(lm_head, hidden_states.ravel().to(lm_head.dtype))
             logits = logits.unsqueeze(0).unsqueeze(0)
@@ -997,7 +995,7 @@ def CausalLM_fast_forward(fast_forward_inference):
 
                 return CausalLMOutputWithPast(
                     loss=loss,
-                    logits=None,
+                    logits=EMPTY_LOGITS,
                     past_key_values=outputs.past_key_values,
                     hidden_states=outputs.hidden_states,
                     attentions=outputs.attentions,