position_ids

This commit is contained in:
Daniel Han-Chen 2024-02-24 14:13:31 +11:00
commit 8920cda9a7
2 changed files with 2 additions and 2 deletions

View file

@ -271,7 +271,7 @@ class FastGemmaModel(FastLlamaModel):
# https://github.com/huggingface/transformers/pull/27931
# https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py
import transformers.models.gemma.modeling_gemma
# transformers.models.gemma.modeling_gemma.GemmaRotaryEmbedding = LlamaRotaryEmbedding
transformers.models.gemma.modeling_gemma.GemmaRotaryEmbedding = LlamaRotaryEmbedding
return
pass

View file

@ -821,7 +821,7 @@ class LlamaRotaryEmbedding(torch.nn.Module):
self.register_buffer("sin_cached", emb.sin().to(dtype=dtype, device=device, non_blocking=True), persistent=False)
pass
def forward(self, x, seq_len=None):
def forward(self, x, position_ids, seq_len=None):
# x: [bs, num_attention_heads, seq_len, head_size]
if seq_len > self.max_seq_len_cached:
self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)