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

fix(gemma4): cast RoPE offset to int before mx.arange() (#4901)

Browse source 
* fix(gemma4): cast RoPE offset to int before mx.arange()

* fix(gemma4): use zero-based arange + offset to avoid CPU-GPU sync
This commit is contained in:
Théophile Lafargue 2026-04-09 14:56:42 +02:00 • committed by GitHub
commit f1041f885f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

2
unsloth/models/gemma4_text.py
View file

@ -247,7 +247,7 @@ class ProportionalRoPE(nn.Module):
def __call__(self, x: mx.array, offset: int = 0) -> mx.array:
# x shape: (B, n_heads, L, head_dim)
seq_len = x.shape[-2]
positions = mx.arange(offset, offset + seq_len, dtype = mx.float32)
positions = mx.arange(seq_len, dtype = mx.float32) + offset
# (L, head_dim//2)
freqs = mx.outer(positions, self._inv_freq)