From 95f4dc6137472ac039e901bc6571327cd267ec6c Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Fri, 4 Jul 2025 03:54:27 -0500 Subject: [PATCH] only warn about prepare causal attention mask when transformers<=4.52.4 (#2867) --- unsloth/models/llama.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index d0ff413925..7d56dac2ec 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -152,12 +152,13 @@ def _fast_prepare_inputs_for_generation(self, input_ids, attention_mask=None, ** ) else: attention_mask = attention_mask[:,[-1]] - logger.warning_once( - f"{self.__class__.__name__} has no `_prepare_4d_causal_attention_mask_with_cache_position` method " - "defined in its base modeling class. Compiled forward passes will be sub-optimal. If you're " - "writing code, see Llama for an example implementation. If you're a user, please report this " - "issue on GitHub." - ) + if transformers_version <= Version("4.52.4"): + logger.warning_once( + f"{self.__class__.__name__} has no `_prepare_4d_causal_attention_mask_with_cache_position` method " + "defined in its base modeling class. Compiled forward passes will be sub-optimal. If you're " + "writing code, see Llama for an example implementation. If you're a user, please report this " + "issue on GitHub." + ) if "cache_position" in kwargs: kwargs["position_ids"] = kwargs["cache_position"]