From 4e570be9ae4ced8cdc64e498125708e34942befc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 31 Jul 2024 12:10:33 -0700 Subject: [PATCH] Fix RoPE extension (#846) * bugs * Update _utils.py * flash-attn softcapping * Update gemma2.py * Update gemma2.py * Update gemma2.py * Update gemma2.py * Update mapper.py * Update README.md * Update _utils.py * Fix ROPE extension issue and device mismatch (#840) * When an exception has been assigned using as target, it is cleared at the end of the except clause.(https://docs.python.org/3/reference/compound_stmts.html#the-try-statement) * Update loader.py * round up to extend rope size * inv_freq.device changed, make sure they are on the same device --------- Co-authored-by: xiaoyang Co-authored-by: Daniel Han * Update gemma.py --------- Co-authored-by: XiaoYang Co-authored-by: xiaoyang --- unsloth/models/gemma.py | 3 ++- unsloth/models/llama.py | 7 ++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index e3f1e615db..a0894ec7a0 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -14,6 +14,7 @@ from .llama import * from ._utils import __version__ +import math try: from transformers.models.gemma.modeling_gemma import ( @@ -256,7 +257,7 @@ class GemmaFixedRotaryEmbedding(torch.nn.Module): def extend_rope_embedding(self, x, seq_len): if seq_len <= self.current_rope_size: return # Iteratively grow by increments of 8192 - self.current_rope_size = int(round(seq_len / 8192)) * 8192 + self.current_rope_size = math.ceil(seq_len / 8192) * 8192 self._set_cos_sin_cache(self.current_rope_size, device = "cuda:0", dtype = x.dtype) pass pass diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index b5244ed4ee..e6c9280bc5 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -14,6 +14,7 @@ import torch import gc +import math from typing import Optional, Tuple, List, Union from ._utils import * from ._utils import __version__ @@ -1036,7 +1037,7 @@ class LlamaRotaryEmbedding(torch.nn.Module): def extend_rope_embedding(self, x, seq_len): if seq_len <= self.current_rope_size: return # Iteratively grow by increments of 8192 - self.current_rope_size = int(round(seq_len / 8192)) * 8192 + self.current_rope_size = math.ceil(seq_len / 8192) * 8192 self._set_cos_sin_cache(self.current_rope_size, device = "cuda:0", dtype = x.dtype) pass pass @@ -1109,7 +1110,7 @@ class LlamaExtendedRotaryEmbedding(torch.nn.Module): # in FP32. They are applied (multiplied) in FP32 as well. self.current_rope_size = seq_len - t = torch.arange(self.current_rope_size, device="cpu", dtype=torch.int64).float() + t = torch.arange(self.current_rope_size, device=self.inv_freq.device, dtype=torch.int64).float() freqs = torch.outer(t, self.inv_freq) # Different from paper, but it uses a different permutation in order to obtain the same calculation @@ -1158,7 +1159,7 @@ class LlamaExtendedRotaryEmbedding(torch.nn.Module): def extend_rope_embedding(self, x, seq_len): if seq_len <= self.current_rope_size: return # Iteratively grow by increments of 8192 - self.current_rope_size = int(round(seq_len / 8192)) * 8192 + self.current_rope_size = math.ceil(seq_len / 8192) * 8192 self._set_cos_sin_cache(self.current_rope_size, device = "cuda:0", dtype = x.dtype) pass pass