ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

patch RoPE

Browse source
This commit is contained in:
Daniel Han 2024-07-23 10:48:45 -07:00
commit 4a46220131
3 changed files with 16 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

12
unsloth/models/_utils.py
View file

@ -752,7 +752,7 @@ def patch_linear_scaling(
fix_rope_function = """
if getattr(self.config, "rope_scaling", None) is None:
self.rotary_emb = {rope_function}(
self.head_dim,
dim = self.head_dim,
max_position_embeddings=self.max_position_embeddings,
base=self.rope_theta,
)
@ -761,7 +761,7 @@ def patch_linear_scaling(
scaling_factor = self.config.rope_scaling["factor"]
if scaling_type == "linear":
self.rotary_emb = {scaled_rope_function}(
self.head_dim,
dim = self.head_dim,
max_position_embeddings=self.max_position_embeddings,
scaling_factor=scaling_factor,
base=self.rope_theta,
@ -827,7 +827,7 @@ def patch_llama_rope_scaling(
fix_rope_function = """
if getattr(self.config, "rope_scaling", None) is None:
self.rotary_emb = {rope_function}(
self.head_dim,
dim = self.head_dim,
max_position_embeddings=self.max_position_embeddings,
base=self.rope_theta,
)
@ -836,17 +836,17 @@ def patch_llama_rope_scaling(
scaling_type2 = self.config.rope_scaling.get("rope_type", None)
scaling_type = scaling_type1 if scaling_type1 is not None else scaling_type2
scaling_factor = self.config.rope_scaling.get("factor")
if scaling_type == "linear":
self.rotary_emb = {scaled_rope_function}(
self.head_dim,
dim = self.head_dim,
max_position_embeddings=self.max_position_embeddings,
scaling_factor=scaling_factor,
base=self.rope_theta,
)
elif scaling_type == "llama3":
self.rotary_emb = {extended_rope_function}(
self.head_dim,
dim = self.head_dim,
max_position_embeddings=self.max_position_embeddings,
base=self.rope_theta,
)

6
unsloth/models/gemma.py
View file

@ -205,9 +205,10 @@ class GemmaFixedRotaryEmbedding(torch.nn.Module):
# Fixes https://github.com/huggingface/transformers/pull/28837
# https://github.com/microsoft/DeepSpeed/issues/4932
# The precision of RoPE buffers is not correct, so we cast to int64.
def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None,
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__()
self.dim = dim
self.max_position_embeddings = max_position_embeddings
@ -266,9 +267,10 @@ class GemmaFixedLinearScalingRotaryEmbedding(GemmaFixedRotaryEmbedding):
# Fixes https://github.com/huggingface/transformers/pull/28837
# https://github.com/microsoft/DeepSpeed/issues/4932
# The precision of RoPE buffers is not correct, so we cast to int64.
def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0,
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)
pass

9
unsloth/models/llama.py
View file

@ -977,9 +977,10 @@ class LlamaRotaryEmbedding(torch.nn.Module):
# Fixes https://github.com/huggingface/transformers/pull/28837
# https://github.com/microsoft/DeepSpeed/issues/4932
# The precision of RoPE buffers is not correct, so we cast to int64.
def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None,
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__()
self.dim = dim
self.max_position_embeddings = max_position_embeddings
@ -1032,9 +1033,10 @@ class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding):
# Fixes https://github.com/huggingface/transformers/pull/28837
# https://github.com/microsoft/DeepSpeed/issues/4932
# The precision of RoPE buffers is not correct, so we cast to int64.
def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0,
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)
pass
@ -1059,9 +1061,10 @@ pass
# See https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/rotary_embedding.py#L736
# For Llama 3.1
class LlamaExtendedRotaryEmbedding(torch.nn.Module):
def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None,
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__()
self.dim = dim
self.max_position_embeddings = max_position_embeddings