From ae336e96631b467504034974686fad0fe385bdcb Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 26 Feb 2024 04:36:10 +1100 Subject: [PATCH] Update gemma.py --- unsloth/models/gemma.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index f00552b0e5..4840d09aa9 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -92,7 +92,8 @@ class FastGemmaRotaryEmbedding(torch.nn.Module): position_ids = torch.arange(self.max_position_embeddings, device=x.device, dtype=torch.int64).unsqueeze(0) inv_freq_expanded = self.inv_freq[None, :, None].float().expand(1, -1, 1) position_ids_expanded = position_ids[:, None, :].float() - freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) + with torch.autocast(enabled = False): + freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) self.cos_cached = emb.cos().to(dtype=x.dtype) self.sin_cached = emb.sin().to(dtype=x.dtype) @@ -106,13 +107,6 @@ class FastGemmaRotaryEmbedding(torch.nn.Module): cos_cached = emb.cos().to(dtype=x.dtype) sin_cached = emb.sin().to(dtype=x.dtype) - print(position_ids) - print(inv_freq_expanded) - print(position_ids_expanded) - print(freqs) - print(self.cos_cached2) - raise 1 - # return self.cos_cached[:,:length], self.sin_cached[:,:length] return self.cos_cached, self.sin_cached pass