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:
parent
dfca5516e7
commit
4e570be9ae
2 changed files with 6 additions and 4 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue