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 <xiaoyang@youzan.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>

* Update gemma.py

---------

Co-authored-by: XiaoYang <xyangk@gmail.com>
Co-authored-by: xiaoyang <xiaoyang@youzan.com>
This commit is contained in:
Daniel Han 2024-07-31 12:10:33 -07:00 committed by GitHub
commit 4e570be9ae
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 6 additions and 4 deletions

View file

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

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