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-01-28 15:13:54 +11:00
commit 6aa46ffff6
1 changed files with 6 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

9
unsloth/models/llama.py
View file

@ -486,9 +486,12 @@ def LlamaModel_fast_forward(
if inputs_embeds is None:
inputs_embeds = self.embed_tokens(input_ids)
inputs_embeds.requires_grad_(False)
inputs_embeds *= attention_mask.unsqueeze(0).transpose(0, 1).transpose(1, 2)
inputs_embeds.requires_grad_(True)
# Fix up attention mask by setting elements to 0
inputs_requires_grad = inputs_embeds.requires_grad
if inputs_requires_grad: inputs_embeds.requires_grad_(False)
inputs_embeds[attention_mask == 0] = -torch.inf
if inputs_requires_grad: inputs_embeds.requires_grad_(True)
#
# Ignore attention_mask
if attention_mask is None: