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

Update llama.py

Browse source
This commit is contained in:
Daniel Han-Chen 2024-03-03 03:35:21 +11:00
commit 9a2e791e6c
1 changed files with 12 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

14
unsloth/models/llama.py
View file

@ -511,7 +511,12 @@ def LlamaModel_fast_forward(
# Mormalized from Gemma
if self.config.model_type == "gemma":
inputs_requires_grad = inputs_embeds.requires_grad
if inputs_requires_grad: inputs_embeds.requires_grad_(False)
if not inputs_embeds.is_leaf:
inputs_embeds = inputs_embeds.detach()
inputs_requires_grad = True
elif inputs_requires_grad:
inputs_embeds.requires_grad_(False)
pass
inputs_embeds *= math_sqrt(self.config.hidden_size)
if inputs_requires_grad: inputs_embeds.requires_grad_(True)
pass
@ -522,7 +527,12 @@ def LlamaModel_fast_forward(
# Careful for inference the attention_mask is size (1, kv_seq_len)
# Whilst the input_embeds is size (1, 1, 4096)
inputs_requires_grad = inputs_embeds.requires_grad
if inputs_requires_grad: inputs_embeds.requires_grad_(False)
if not inputs_embeds.is_leaf:
inputs_embeds = inputs_embeds.detach()
inputs_requires_grad = True
elif inputs_requires_grad:
inputs_embeds.requires_grad_(False)
pass
inputs_embeds *= attention_mask.unsqueeze(0).transpose(0, 1).transpose(1, 2)
if inputs_requires_grad: inputs_embeds.requires_grad_(True)
pass