From 86884ab446b284d302bf558a868d1f56c0fc1e0c Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Mon, 28 Jul 2025 15:34:49 +0530 Subject: [PATCH] Fixup multi GPU workload. (#3049) * sync all instead * sync after move and rope init instead * sync after rope inside * Return new tensors and no sync * Sync only current stream * Fixup mask for xformers * sync for prefill only * clean up --- unsloth/models/llama.py | 2 ++ unsloth/models/mistral.py | 27 ++++++++++++++++++++------- 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 3c0d5012ae..e7d9084ad8 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -499,6 +499,8 @@ def LlamaAttention_fast_forward( # else inplace_rope_embedding(Q, K, cos, sin, position_ids) # ) Q, K = fast_rope_embedding(Q, K, cos, sin) + # synchronize before cat to avoid race condition + torch.cuda.current_stream(Q.device).synchronize() if past_key_value is not None: K = torch.cat([past_key_value[0], K], dim = 2) diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index 68d4ba43fb..ef046d1f3c 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -190,7 +190,8 @@ def MistralForCausalLM_fast_forward( bsz, q_len = input_ids.shape sliding_window = getattr(self.config, "sliding_window", None) - if HAS_XFORMERS and attention_mask is None: + if HAS_XFORMERS: + # Always create causal mask for xformers if sliding_window is None or sliding_window == "null" or sliding_window <= 0: causal_mask = xformers.attn_bias.LowerTriangularMask() elif q_len <= sliding_window: @@ -200,12 +201,13 @@ def MistralForCausalLM_fast_forward( .from_seqlens([q_len]*bsz)\ .make_local_attention(window_size = sliding_window) - elif not HAS_XFORMERS and attention_mask is None: + # If attention_mask exists, it will be handled in the attention forward + + else: + # Not using xformers - need to create attention masks if sliding_window is None or sliding_window == "null" or sliding_window <= 0 or q_len <= sliding_window: # Fully causal mask - mask = torch.full((q_len, q_len), -torch.inf, device=input_ids.device) - mask = torch.triu(mask, diagonal=1) - attention_mask = mask.expand(bsz, 1, q_len, q_len) + causal_mask_values = torch.triu(torch.full((q_len, q_len), -torch.inf, device=input_ids.device), diagonal=1) else: # Sliding window attention q_indices = torch.arange(q_len, device=input_ids.device).view(-1, 1) @@ -214,8 +216,19 @@ def MistralForCausalLM_fast_forward( causal_bool_mask = k_indices <= q_indices window_bool_mask = (q_indices - k_indices) < sliding_window - mask = torch.where(causal_bool_mask & window_bool_mask, 0.0, -torch.inf) - attention_mask = mask[None, None, :, :].expand(bsz, 1, q_len, q_len) + causal_mask_values = torch.where(causal_bool_mask & window_bool_mask, 0.0, -torch.inf) + + # Combine with existing attention_mask if present + if attention_mask is None: + attention_mask = causal_mask_values[None, None, :, :].expand(bsz, 1, q_len, q_len) + else: + # attention_mask should be [bsz, 1, q_len, q_len] or broadcastable + # Add causal mask to existing attention mask + if attention_mask.dim() == 2: + # [bsz, seq_len] -> [bsz, 1, 1, seq_len] + attention_mask = attention_mask[:, None, None, :] + attention_mask = attention_mask.expand(bsz, 1, q_len, q_len) + attention_mask = attention_mask + causal_mask_values[None, None, :, :] attention_mask = attention_mask.to(dtype=_get_dtype(self.config.torch_dtype))