From 37d0dedde5032e33216d6bc903111fbadd7ae64d Mon Sep 17 00:00:00 2001 From: Lei Zhenyuan Date: Tue, 22 Jul 2025 18:27:57 +0800 Subject: [PATCH 1/3] fix for casual mask (#3011) --- unsloth/models/llama.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 8d985aa9d2..27b92282c3 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -342,9 +342,9 @@ def LlamaAttention_fast_forward_inference( A = torch_matmul(A, Vnn, out = Qn) else: if SDPA_HAS_GQA: - A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = False, enable_gqa = True) + A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = is_causal, enable_gqa = True) else: - A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = False) + A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = is_causal) pass A = A.transpose(1, 2) A = A.reshape(bsz, 1, attention_size) From 7d8acbdf86c44e161905201060769fbab5a79257 Mon Sep 17 00:00:00 2001 From: Lei Zhenyuan Date: Tue, 22 Jul 2025 18:33:26 +0800 Subject: [PATCH 2/3] [intel] add for intel path for llama.py (#3012) * fix for intel path * remove unuse code * Update unsloth/models/llama.py --------- Co-authored-by: Daniel Han --- unsloth/models/llama.py | 27 ++++++++++++++++----------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 27b92282c3..de53704e66 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -92,6 +92,11 @@ def clean_gpu_cache(): else: torch.cuda.empty_cache() +if DEVICE_TYPE == "xpu": + get_current_device = torch.xpu.current_device +else: + get_current_device = torch.cuda.current_device + def original_apply_qkv(self, X): Q = self.q_proj(X) K = self.k_proj(X) @@ -1365,8 +1370,8 @@ class LlamaRotaryEmbedding(torch.nn.Module): self._set_cos_sin_cache(seq_len=self.current_rope_size, device=torch.device(device_idx), dtype=torch.get_default_dtype()) # dummy so that patch_utils doesn't fail for now - self.cos_cached = torch.empty(1, device=torch.cuda.current_device(), dtype=torch.get_default_dtype()) - self.sin_cached = torch.empty(1, device=torch.cuda.current_device(), dtype=torch.get_default_dtype()) + self.cos_cached = torch.empty(1, device=get_current_device(), dtype=torch.get_default_dtype()) + self.sin_cached = torch.empty(1, device=get_current_device(), dtype=torch.get_default_dtype()) pass def _set_cos_sin_cache(self, seq_len, device, dtype): @@ -1402,7 +1407,7 @@ class LlamaRotaryEmbedding(torch.nn.Module): def get_cached(self, seq_len = None, device_index = None): if device_index is None: - device_index = torch.cuda.current_device() + device_index = get_current_device() return self.multi_gpu_cos_cached[device_index], self.multi_gpu_sin_cached[device_index] pass @@ -1484,8 +1489,8 @@ class LlamaExtendedRotaryEmbedding(torch.nn.Module): self._set_cos_sin_cache(seq_len=self.current_rope_size, device=torch.device(device_idx), dtype=torch.get_default_dtype()) # dummy so that patch_utils doesn't fail for now - self.cos_cached = torch.empty(1, device=torch.cuda.current_device(), dtype=torch.get_default_dtype()) - self.sin_cached = torch.empty(1, device=torch.cuda.current_device(), dtype=torch.get_default_dtype()) + self.cos_cached = torch.empty(1, device=get_current_device(), dtype=torch.get_default_dtype()) + self.sin_cached = torch.empty(1, device=get_current_device(), dtype=torch.get_default_dtype()) pass def _set_cos_sin_cache(self, seq_len, device, dtype): @@ -1518,7 +1523,7 @@ class LlamaExtendedRotaryEmbedding(torch.nn.Module): def get_cached(self, seq_len = None, device_index = None): if device_index is None: - device_index = torch.cuda.current_device() + device_index = get_current_device() return self.multi_gpu_cos_cached[device_index], self.multi_gpu_sin_cached[device_index] pass @@ -1631,10 +1636,10 @@ class LongRopeRotaryEmbedding(torch.nn.Module): self.multi_gpu_short_sin_cached[device_idx] = sin_cached # dummy so that patch_utils doesn't fail for now - self.short_cos_cached = torch.empty(1, device=torch.cuda.current_device(), dtype=torch.get_default_dtype()) - self.short_sin_cached = torch.empty(1, device=torch.cuda.current_device(), dtype=torch.get_default_dtype()) - self.long_cos_cached = torch.empty(1, device=torch.cuda.current_device(), dtype=torch.get_default_dtype()) - self.long_sin_cached = torch.empty(1, device=torch.cuda.current_device(), dtype=torch.get_default_dtype()) + self.short_cos_cached = torch.empty(1, device=get_current_device(), dtype=torch.get_default_dtype()) + self.short_sin_cached = torch.empty(1, device=get_current_device(), dtype=torch.get_default_dtype()) + self.long_cos_cached = torch.empty(1, device=get_current_device(), dtype=torch.get_default_dtype()) + self.long_sin_cached = torch.empty(1, device=get_current_device(), dtype=torch.get_default_dtype()) pass def _set_cos_sin_cache(self, seq_len, device, dtype): @@ -1675,7 +1680,7 @@ class LongRopeRotaryEmbedding(torch.nn.Module): def get_cached(self, seq_len = None, device_index = None): if device_index is None: - device_index = torch.cuda.current_device() + device_index = get_current_device() if seq_len is not None and seq_len < self.original_max_position_embeddings: return self.multi_gpu_short_cos_cached[device_index], self.multi_gpu_short_sin_cached[device_index] return self.multi_gpu_long_cos_cached[device_index], self.multi_gpu_long_sin_cached[device_index] From 14c4bab15c496fcc59c199b70e0e568880170915 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 22 Jul 2025 03:34:51 -0700 Subject: [PATCH 3/3] Update llama.py --- unsloth/models/llama.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index de53704e66..9b920fdac9 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -86,16 +86,13 @@ from triton import __version__ as triton_version HAS_XFORMERS = xformers is not None BlockDiagonalCausalMask = xformers.attn_bias.BlockDiagonalCausalMask if HAS_XFORMERS else None -def clean_gpu_cache(): - if DEVICE_TYPE == "xpu": - torch.xpu.empty_cache() - else: - torch.cuda.empty_cache() - if DEVICE_TYPE == "xpu": + clean_gpu_cache = torch.xpu.empty_cache get_current_device = torch.xpu.current_device else: + clean_gpu_cache = torch.cuda.empty_cache get_current_device = torch.cuda.current_device +pass def original_apply_qkv(self, X): Q = self.q_proj(X)