From 1aeca6f00cd35d4f5f7b06f478debdb543748acb Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 26 Feb 2024 14:35:44 +1100 Subject: [PATCH] RoPE --- unsloth/models/gemma.py | 12 ++++++++---- unsloth/models/llama.py | 7 +++++++ 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 67fda48127..cf27bacf7f 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -108,9 +108,6 @@ class FastGemmaRotaryEmbedding(torch.nn.Module): print(freqs) print(self.cos_cached2) - print(freqs.to(torch.int64)) - print(self.cos_cached2.to(torch.int64)) - raise 1 # return self.cos_cached[:,:length], self.sin_cached[:,:length] return self.cos_cached, self.sin_cached @@ -581,7 +578,7 @@ class FastGemmaModel(FastLlamaModel): # https://github.com/huggingface/transformers/pull/27931 # https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py import transformers.models.gemma.modeling_gemma - transformers.models.gemma.modeling_gemma.GemmaRotaryEmbedding = FastGemmaRotaryEmbedding + transformers.models.gemma.modeling_gemma.GemmaRotaryEmbedding = LlamaRotaryEmbedding return pass @@ -631,6 +628,13 @@ class FastGemmaModel(FastLlamaModel): quant_state.dtype = correct_dtype pass pass + # Downcast RoPE embedding to correct data type + if (name.endswith("rotary_emb") or hasattr(module, "cos_cached")) \ + and (module.cos_cached.dtype != expected_dtype): + module.cos_cached = module.cos_cached.to(expected_dtype) + module.sin_cached = module.sin_cached.to(expected_dtype) + pass + pass pass # Add 1 to weight diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index f552cb473f..047eafc56e 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1174,6 +1174,13 @@ class FastLlamaModel: quant_state.dtype = correct_dtype pass pass + # Downcast RoPE embedding to correct data type + if (name.endswith("rotary_emb") or hasattr(module, "cos_cached")) \ + and (module.cos_cached.dtype != expected_dtype): + module.cos_cached = module.cos_cached.to(expected_dtype) + module.sin_cached = module.sin_cached.to(expected_dtype) + pass + pass pass # Clear deleted GPU items