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 16:31:01 +11:00
commit 2f73cb4049
1 changed files with 10 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

17
unsloth/models/llama.py
View file

@ -486,13 +486,6 @@ def LlamaModel_fast_forward(
if inputs_embeds is None:
inputs_embeds = self.embed_tokens(input_ids)
# 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:
padding_mask = None
@ -525,6 +518,15 @@ def LlamaModel_fast_forward(
use_cache = False
pass
# Fix up attention mask by setting elements to 0
# Specifically for DPO
if self._has_no_labels and attention_mask is not None:
inputs_requires_grad = hidden_states.requires_grad
if inputs_requires_grad: hidden_states.requires_grad_(False)
hidden_states *= attention_mask.unsqueeze(0).transpose(0, 1).transpose(1, 2)
if inputs_requires_grad: hidden_states.requires_grad_(True)
pass
# decoder layers
all_hidden_states = () if output_hidden_states else None
all_self_attns = () if output_attentions else None
@ -624,6 +626,7 @@ def LlamaForCausalLM_fast_forward(
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
self.model._has_no_labels = labels is None
outputs = self.model(
input_ids=input_ids,
causal_mask=causal_mask,