From bc428f4e7b8f1399b613e545e302d5a1b877e944 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sun, 25 Feb 2024 04:29:01 +1100 Subject: [PATCH] Update gemma.py --- unsloth/models/gemma.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index b6b6152d4e..913569b6bd 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -78,6 +78,8 @@ class FastGemmaRotaryEmbedding(torch.nn.Module): emb = torch.cat((freqs, freqs), dim=-1) self.cos_cached = emb.cos().to(dtype=torch.bfloat16) self.sin_cached = emb.sin().to(dtype=torch.bfloat16) + print(self.cos_cached) + print(self.sin_cached) def forward(self, x, position_ids, seq_len=None): # x: [bs, num_attention_heads, seq_len, head_size] @@ -87,15 +89,17 @@ class FastGemmaRotaryEmbedding(torch.nn.Module): ) length = position_ids.shape[1] - if self.cos_cached is None: - 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) - 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) - pass + #if self.cos_cached is None: + 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) + 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) + # pass + print(self.cos_cached) + print(self.sin_cached) + raise return self.cos_cached[:,:length], self.sin_cached[:,:length] pass