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

dim fix

Browse source
This commit is contained in:
Daniel Han 2025-02-01 19:27:33 -08:00
commit 578ed29771
2 changed files with 2 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

1
unsloth/models/gemma.py
View file

@ -224,7 +224,6 @@ class GemmaFixedRotaryEmbedding(torch.nn.Module):
self.base = base
# Dynamic RoPE we first set it to a max of 4 * 8192 tokens then we iteratively grow this
self.current_rope_size = min(4 * 8192, self.max_position_embeddings)
print(dim, max_position_embeddings, base)
# Build here to make `torch.jit.trace` work.
self._set_cos_sin_cache(seq_len=self.current_rope_size, device=device, dtype=torch.get_default_dtype())

3
unsloth/models/llama.py
View file

@ -1159,7 +1159,8 @@ class LlamaRotaryEmbedding(torch.nn.Module):
# [TODO] Hack to pass in config - need to remove later
base = config.rope_theta
partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
dim = int((config.hidden_size // config.num_attention_heads))
dim = getattr(config, "head_dim", None)
if dim is None: dim = int((config.hidden_size // config.num_attention_heads))
device = "cuda"
max_position_embeddings = config.max_position_embeddings
pass