From 85c052d8fbd7636a673023289a41c8d22322bf33 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Wed, 6 Mar 2024 04:48:26 +1100 Subject: [PATCH] sqrt --- unsloth/models/gemma.py | 5 +++-- unsloth/models/llama.py | 3 ++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index c23f9171f8..ba7aac7d31 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -156,8 +156,9 @@ def GemmaModel_fast_forward_inference( hidden_states = self.embed_tokens(input_ids) # 3072**0.5 = 55.5000 in bfloat16, whilst 55.4256 in float32 # 2048**0.5 = 45.2500 in bfloat16, whilst 45.2548 in float32 - inputs_embeds *= torch.tensor(math_sqrt(self.config.hidden_size), dtype = inputs_embeds.dtype) - + # inputs_embeds *= torch.tensor(math_sqrt(self.config.hidden_size), dtype = inputs_embeds.dtype) + inputs_embeds *= math_sqrt(self.config.hidden_size) + next_decoder_cache = [] for idx, decoder_layer in enumerate(self.layers): # Self Attention diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index d7743216fb..4b9161a1a5 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -523,7 +523,8 @@ def LlamaModel_fast_forward( # inputs_embeds *= math_sqrt(self.config.hidden_size) # Ie 3072**0.5 = 55.5000 in bfloat16, whilst 55.4256 in float32 # & 2048**0.5 = 45.2500 in bfloat16, whilst 45.2548 in float32 - inputs_embeds *= torch.tensor(math_sqrt(self.config.hidden_size), dtype = inputs_embeds.dtype) + # inputs_embeds *= torch.tensor(math_sqrt(self.config.hidden_size), dtype = inputs_embeds.dtype) + inputs_embeds *= math_sqrt(self.config.hidden_size) if inputs_requires_grad: inputs_embeds.requires_grad_(True) pass