From 0bb588bf8d1ab4272405634485cf43ee8f33b84a Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 14 Nov 2024 01:11:16 -0800
Subject: [PATCH] Update llama.py

---
 unsloth/models/llama.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 47a57024a2..f4609b81e9 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -738,6 +738,10 @@ def LlamaModel_fast_forward(
                 past_key_values_length,
                 sliding_window = None,
             )
+            # Fixes https://github.com/unslothai/unsloth/issues/853
+            # Unsloth needs a 2D mask, not a [2, 1, n, n] mask!
+            if self.SWA_mask.dim() == 4: self.SWA_mask = self.SWA_mask[0][0]
+            if self. GA_mask.dim() == 4: self. GA_mask = self. GA_mask[0][0]
         elif not hasattr(self, "SWA_mask"):
             if HAS_FLEX_ATTENTION:
                 # Use Flex Attention instead!