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-29 17:35:19 +11:00
commit 5cfea20129
1 changed files with 9 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

18
unsloth/models/llama.py
View file

@ -488,15 +488,15 @@ def LlamaModel_fast_forward(
# Fix up attention mask by setting elements to 0 # Fix up attention mask by setting elements to 0
# Specifically for DPO # Specifically for DPO
try: if self._has_no_labels and attention_mask is not None and \
if self._has_no_labels and attention_mask is not None: attention_mask.shape[1] == seq_length:
inputs_requires_grad = inputs_embeds.requires_grad # Careful for inference the attention_mask is size (1, kv_seq_len)
if inputs_requires_grad: inputs_embeds.requires_grad_(False) # Whilst the input_embeds is size (1, 1, 4096)
inputs_embeds *= attention_mask.unsqueeze(0).transpose(0, 1).transpose(1, 2) inputs_requires_grad = inputs_embeds.requires_grad
if inputs_requires_grad: inputs_embeds.requires_grad_(True) if inputs_requires_grad: inputs_embeds.requires_grad_(False)
pass inputs_embeds *= attention_mask.unsqueeze(0).transpose(0, 1).transpose(1, 2)
except: if inputs_requires_grad: inputs_embeds.requires_grad_(True)
print(inputs_embeds.shape, attention_mask.shape) pass
# Ignore attention_mask # Ignore attention_mask
if attention_mask is None: if attention_mask is None: