From 5cfea201294d66646ec65ca46fd1277f2b042530 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 29 Jan 2024 17:35:19 +1100 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 01d0af67f2..36bb58f3d9 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -488,15 +488,15 @@ def LlamaModel_fast_forward( # Fix up attention mask by setting elements to 0 # Specifically for DPO - try: - if self._has_no_labels and attention_mask is not None: - inputs_requires_grad = inputs_embeds.requires_grad - if inputs_requires_grad: inputs_embeds.requires_grad_(False) - inputs_embeds *= attention_mask.unsqueeze(0).transpose(0, 1).transpose(1, 2) - if inputs_requires_grad: inputs_embeds.requires_grad_(True) - pass - except: - print(inputs_embeds.shape, attention_mask.shape) + if self._has_no_labels and attention_mask is not None and \ + attention_mask.shape[1] == seq_length: + # 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) + inputs_embeds *= attention_mask.unsqueeze(0).transpose(0, 1).transpose(1, 2) + if inputs_requires_grad: inputs_embeds.requires_grad_(True) + pass # Ignore attention_mask if attention_mask is None: