From 549c5be61c2b0223d3c64eb94447b8ea73c509cf Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 Nov 2024 21:39:19 -0800 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index a17af2dd8a..36fd151f4e 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -980,7 +980,8 @@ def CausalLM_fast_forward(fast_forward_inference): elif num_logits_to_keep != 0: logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :].to(lm_head.dtype)) else: - if "return_logits" not in kwargs and HAS_CUT_CROSS_ENTROPY and labels is not None: + RETURN_LOGITS = os.environ.get("UNSLOTH_RETURN_LOGITS", "0") == "1" + if not RETURN_LOGITS and HAS_CUT_CROSS_ENTROPY and labels is not None: n_items = kwargs.get("num_items_in_batch", None) or kwargs.get("n_items", None) loss = fused_linear_cross_entropy( hidden_states = hidden_states, @@ -995,7 +996,7 @@ def CausalLM_fast_forward(fast_forward_inference): return CausalLMOutputWithPast( loss=loss, - logits=EMPTY_LOGITS, + logits=None, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions,