From 316aaefdf2d9a2001535ff82d3cfbae80eed3f17 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 8 Jul 2024 10:44:19 -0700 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index e41a45b266..de7f7bc49b 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -897,9 +897,15 @@ def CausalLM_fast_forward(fast_forward_inference): logit_softcapping = logit_softcapping, ) elif logit_softcapping != 0: - logits *= (1.0 / logit_softcapping) - logits = torch.tanh(logits, out = logits if not logits.requires_grad else None) - logits *= logit_softcapping + if logits.requires_grad: + logits = (1.0 / logit_softcapping) * logits + logits = torch.tanh(logits) + logits = logit_softcapping * logits + else: + logits *= (1.0 / logit_softcapping) + torch.tanh(logits, out = logits) + logits *= logit_softcapping + pass pass if not return_dict: