From e160305b669c70b90496e4d7c71ba5e6e494a1a8 Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Wed, 4 Jun 2025 14:58:50 -0500 Subject: [PATCH] Update prepare 4d causal attention call (#2678) --- unsloth/models/llama.py | 32 ++++++++++++++++++++++++-------- 1 file changed, 24 insertions(+), 8 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 53c956ce1a..cba8a68565 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -121,16 +121,32 @@ def _fast_prepare_inputs_for_generation(self, input_ids, attention_mask=None, ** base_model = getattr(base_model, base_model.base_model_prefix) if hasattr(base_model, "_prepare_4d_causal_attention_mask_with_cache_position"): + def needs_device_kw(fn) -> bool: + try: + sig = inspect.signature(inspect.unwrap(fn)) + return "device" in sig.parameters + except: + # transformers <= 4.51.3 includes device arg but > 4.51.3 does not + return transformers_version < Version("4.52.0") + + kwargs = { + "sequence_length": 1, + "target_length": cache_length, + "dtype": self.dtype, + "cache_position": torch.arange(cache_length, cache_length+1, device=input_ids.device), + "batch_size": bs, + "config": self.config, + "past_key_values": past_key_values, + } + try: + if needs_device_kw(base_model._prepare_4d_causal_attention_mask_with_cache_position): + kwargs["device"] = input_ids.device + except: + print(f"Unsloth: Could not inspect signature of {base_model._prepare_4d_causal_attention_mask_with_cache_position}") + attention_mask = base_model._prepare_4d_causal_attention_mask_with_cache_position( attention_mask, - sequence_length=1, - target_length=cache_length, - dtype=self.dtype, - device=input_ids.device, - cache_position=torch.arange(cache_length, cache_length+1, device=input_ids.device), - batch_size=bs, - config=self.config, - past_key_values=past_key_values, + **kwargs, ) else: attention_mask = attention_mask[:,[-1]]