From fa58ed8a08ce0b6cc7eb26c76001782b549a9ad4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 23 Jul 2024 10:53:31 -0700 Subject: [PATCH] refix rope --- unsloth/models/gemma.py | 5 ++--- unsloth/models/llama.py | 9 +++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 3dccf63ae4..e3f1e615db 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -208,8 +208,8 @@ class GemmaFixedRotaryEmbedding(torch.nn.Module): def __init__(self, dim = None, max_position_embeddings=2048, base=10000, device=None, config = None, # [TODO] Hack to pass in config - need to remove later ): - if config is not None: return # [TODO] Hack to pass in config - need to remove later super().__init__() + if config is not None: return # [TODO] Hack to pass in config - need to remove later self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base @@ -270,9 +270,8 @@ class GemmaFixedLinearScalingRotaryEmbedding(GemmaFixedRotaryEmbedding): def __init__(self, dim = None, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0, config = None, # [TODO] Hack to pass in config - need to remove later ): - if config is not None: return # [TODO] Hack to pass in config - need to remove later self.scaling_factor = scaling_factor - super().__init__(dim, max_position_embeddings, base, device) + super().__init__(dim = dim, max_position_embeddings = max_position_embeddings, base = base, device = device, config = config) pass def _set_cos_sin_cache(self, seq_len, device, dtype): diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index d043f03d1f..a4a6527ff2 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -980,8 +980,9 @@ class LlamaRotaryEmbedding(torch.nn.Module): def __init__(self, dim = None, max_position_embeddings=2048, base=10000, device=None, config = None, # [TODO] Hack to pass in config - need to remove later ): - if config is not None: return # [TODO] Hack to pass in config - need to remove later super().__init__() + if config is not None: return # [TODO] Hack to pass in config - need to remove later + self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base @@ -1036,9 +1037,8 @@ class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding): def __init__(self, dim = None, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0, config = None, # [TODO] Hack to pass in config - need to remove later ): - if config is not None: return # [TODO] Hack to pass in config - need to remove later self.scaling_factor = scaling_factor - super().__init__(dim, max_position_embeddings, base, device) + super().__init__(dim = dim, max_position_embeddings = max_position_embeddings, base = base, device = device, config = config) pass def _set_cos_sin_cache(self, seq_len, device, dtype): @@ -1064,8 +1064,9 @@ class LlamaExtendedRotaryEmbedding(torch.nn.Module): def __init__(self, dim = None, max_position_embeddings=2048, base=10000, device=None, config = None, # [TODO] Hack to pass in config - need to remove later ): - if config is not None: return # [TODO] Hack to pass in config - need to remove later super().__init__() + if config is not None: return # [TODO] Hack to pass in config - need to remove later + self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base