From bb37fb71c627038d489cacbfb4f0b5dc37d6b0f6 Mon Sep 17 00:00:00 2001
From: Edd <68678137+Erland366@users.noreply.github.com>
Date: Fri, 25 Jul 2025 14:38:20 +0800
Subject: [PATCH] Fix Llama and Gemma inference (#3034)

* Fix Llama and Gemma inference

* Add simple quality life for CUDA link error (which is not captured since we bypass all error)
---
 unsloth/models/falcon_h1.py |  2 +-
 unsloth/models/llama.py     | 10 +++++-----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/unsloth/models/falcon_h1.py b/unsloth/models/falcon_h1.py
index 0db9c1ca4e..5643e3dfe9 100644
--- a/unsloth/models/falcon_h1.py
+++ b/unsloth/models/falcon_h1.py
@@ -51,7 +51,7 @@ try:
     from transformers.models.falcon_h1.modeling_falcon_h1 import (
         FalconH1Attention,
     )
-except:
+except ModuleNotFoundError:
     # if we are on a old version of transformers technically it should fail in the try except above
     # but if somehow we make it here, we need to raise an error since FalconH1Attention is not available
     # or renamed
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 9b920fdac9..3c0d5012ae 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -780,7 +780,7 @@ def LlamaModel_fast_forward(
     # Fix up attention mask by setting elements to 0
     # Specifically for DPO
     if getattr(self, "_has_no_labels", False) is True and (attention_mask is not None) and (past_key_values is None) and \
-        (not train_embed_tokens):
+        (not train_embed_tokens) and self.training:
         # 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
@@ -865,21 +865,21 @@ def LlamaModel_fast_forward(
 
             # https://github.com/pytorch/pytorch/issues/103749
             # Need to convert to float and not using bool
-            attention_mask = (1.0 - attention_mask.float()) * torch.finfo(inputs_embeds.dtype).min
+            # attention_mask = (1.0 - attention_mask.float()) * torch.finfo(inputs_embeds.dtype).min
             dynamic_SWA_mask = _prepare_4d_causal_attention_mask_for_sdpa(
                 attention_mask,
                 (batch_size, seq_length),
                 inputs_embeds,
                 past_key_values_length,
                 sliding_window = self.config.sliding_window,
-            )[0][0]
+            )
             dynamic_GA_mask = _prepare_4d_causal_attention_mask_for_sdpa(
                 attention_mask,
                 (batch_size, seq_length),
                 inputs_embeds,
                 past_key_values_length,
                 sliding_window = None,
-            )[0][0]
+            )
             use_static_mask = False
 
         elif not hasattr(self, "SWA_mask"):
@@ -953,7 +953,7 @@ def LlamaModel_fast_forward(
         else:
             layer_outputs = decoder_layer(
                 hidden_states,
-                causal_mask=mask,
+                causal_mask         = mask,
                 attention_mask      = attention_mask,
                 position_ids        = position_ids,
                 past_key_value      = past_key_value,