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

Fix ROPE extension issue and device mismatch (#840)

Browse source 
* 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 <xiaoyang@youzan.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
XiaoYang 2024-08-01 03:05:08 +08:00 • committed by GitHub
commit 2de142712d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 4 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

7
unsloth/models/llama.py
View file

@ -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