From e279c5ab26090cc73dde0d626fa3d8d5e554f768 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 10 Dec 2024 02:46:14 -0800 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 38ad6062db..cfeeab0800 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -993,6 +993,9 @@ def CausalLM_fast_forward(fast_forward_inference): logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :].to(lm_head.dtype)) else: RETURN_LOGITS = os.environ.get("UNSLOTH_RETURN_LOGITS", "0") == "1" + # < 1024 Normal Unsloth uses less VRAM! + if bsz*q_len <= 1024: RETURN_LOGITS = True + 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(