ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-07-04 19:25:43 +00:00
commit 97116919ea
2 changed files with 88 additions and 57 deletions
Download patch file Download diff file Expand all files Collapse all files

39
tests/utils/test_attn_mask_compat.py
View file

@ -31,17 +31,16 @@ compat = _load_compat_module()
def test_no_deprecation_warning_on_causal_mask():
with warnings.catch_warnings(record=True) as caught:
with warnings.catch_warnings(record = True) as caught:
warnings.simplefilter("always")
compat.AttentionMaskConverter(is_causal=True, sliding_window=3).to_causal_4d(
compat.AttentionMaskConverter(is_causal = True, sliding_window = 3).to_causal_4d(
1,
8,
8,
dtype=torch.float16,
dtype = torch.float16,
)
assert not any(
issubclass(w.category, FutureWarning)
and "modeling_attn_mask_utils" in str(w.message)
issubclass(w.category, FutureWarning) and "modeling_attn_mask_utils" in str(w.message)
for w in caught
)
@ -61,23 +60,23 @@ def test_causal_4d_matches_transformers(batch_size, query_length, sliding_window
with warnings.catch_warnings():
warnings.simplefilter("ignore", FutureWarning)
expected = legacy.AttentionMaskConverter(
is_causal=True,
sliding_window=sliding_window,
is_causal = True,
sliding_window = sliding_window,
).to_causal_4d(
batch_size,
query_length,
key_value_length,
dtype=dtype,
dtype = dtype,
)
actual = compat.AttentionMaskConverter(
is_causal=True,
sliding_window=sliding_window,
is_causal = True,
sliding_window = sliding_window,
).to_causal_4d(
batch_size,
query_length,
key_value_length,
dtype=dtype,
dtype = dtype,
)
if expected is None:
@ -95,7 +94,9 @@ def test_causal_4d_matches_transformers(batch_size, query_length, sliding_window
(torch.tensor([[1, 1, 1, 0, 0], [1, 1, 1, 1, 1]]), 0),
],
)
def test_prepare_4d_causal_attention_mask_for_sdpa_matches_transformers(attention_mask, past_length):
def test_prepare_4d_causal_attention_mask_for_sdpa_matches_transformers(
attention_mask, past_length
):
try:
legacy = importlib.import_module("transformers.modeling_attn_mask_utils")
except ImportError:
@ -103,7 +104,7 @@ def test_prepare_4d_causal_attention_mask_for_sdpa_matches_transformers(attentio
batch_size = 2 if attention_mask is not None else 1
query_length = 5
inputs_embeds = torch.zeros(batch_size, query_length, 16, dtype=torch.float32)
inputs_embeds = torch.zeros(batch_size, query_length, 16, dtype = torch.float32)
with warnings.catch_warnings():
warnings.simplefilter("ignore", FutureWarning)
@ -112,7 +113,7 @@ def test_prepare_4d_causal_attention_mask_for_sdpa_matches_transformers(attentio
(batch_size, query_length),
inputs_embeds,
past_length,
sliding_window=3,
sliding_window = 3,
)
actual = compat._prepare_4d_causal_attention_mask_for_sdpa(
@ -120,7 +121,7 @@ def test_prepare_4d_causal_attention_mask_for_sdpa_matches_transformers(attentio
(batch_size, query_length),
inputs_embeds,
past_length,
sliding_window=3,
sliding_window = 3,
)
if expected is None:
@ -135,14 +136,14 @@ def test_prepare_4d_attention_mask_for_sdpa_matches_transformers():
except ImportError:
pytest.skip("transformers.modeling_attn_mask_utils unavailable")
mask = torch.tensor([[1, 1, 0, 0], [1, 1, 1, 1]], dtype=torch.float32)
mask = torch.tensor([[1, 1, 0, 0], [1, 1, 1, 1]], dtype = torch.float32)
dtype = torch.float32
with warnings.catch_warnings():
warnings.simplefilter("ignore", FutureWarning)
expected = legacy._prepare_4d_attention_mask_for_sdpa(mask, dtype=dtype)
expected = legacy._prepare_4d_attention_mask_for_sdpa(mask, dtype = dtype)
actual = compat._prepare_4d_attention_mask_for_sdpa(mask, dtype=dtype)
actual = compat._prepare_4d_attention_mask_for_sdpa(mask, dtype = dtype)
if expected is None:
assert actual is None
@ -156,7 +157,7 @@ def test_repo_has_no_direct_deprecated_imports():
for path in model_dir.glob("*.py"):
if path.name == "_attn_mask_compat.py":
continue
text = path.read_text(encoding="utf-8")
text = path.read_text(encoding = "utf-8")
if "transformers.modeling_attn_mask_utils" in text:
offenders.append(str(path.relative_to(_REPO_ROOT)))
assert offenders == []

106
unsloth/models/_attn_mask_compat.py
View file

@ -38,7 +38,11 @@ class AttentionMaskConverter:
is_causal: bool
sliding_window: int | None = None
def __init__(self, is_causal: bool, sliding_window: int | None = None):
def __init__(
self,
is_causal: bool,
sliding_window: int | None = None,
):
self.is_causal = is_causal
self.sliding_window = sliding_window
@ -56,7 +60,9 @@ class AttentionMaskConverter:
device: Union[torch.device, str] = "cpu",
) -> torch.Tensor | None:
if not self.is_causal:
raise ValueError(f"Please use `to_causal_4d` only if {self.__class__} has `is_causal` set to True.")
raise ValueError(
f"Please use `to_causal_4d` only if {self.__class__} has `is_causal` set to True."
)
input_shape = (batch_size, query_length)
past_key_values_length = key_value_length - query_length
@ -66,9 +72,9 @@ class AttentionMaskConverter:
causal_4d_mask = self._make_causal_mask(
input_shape,
dtype,
device=device,
past_key_values_length=past_key_values_length,
sliding_window=self.sliding_window,
device = device,
past_key_values_length = past_key_values_length,
sliding_window = self.sliding_window,
)
return causal_4d_mask
@ -93,19 +99,23 @@ class AttentionMaskConverter:
causal_4d_mask = self._make_causal_mask(
input_shape,
dtype,
device=attention_mask_2d.device,
past_key_values_length=past_key_values_length,
sliding_window=self.sliding_window,
device = attention_mask_2d.device,
past_key_values_length = past_key_values_length,
sliding_window = self.sliding_window,
)
elif self.sliding_window is not None:
raise NotImplementedError("Sliding window is currently only implemented for causal masking")
raise NotImplementedError(
"Sliding window is currently only implemented for causal masking"
)
expanded_attn_mask = self._expand_mask(attention_mask_2d, dtype, tgt_len=input_shape[-1]).to(
attention_mask_2d.device
)
expanded_attn_mask = self._expand_mask(
attention_mask_2d, dtype, tgt_len = input_shape[-1]
).to(attention_mask_2d.device)
if causal_4d_mask is not None:
expanded_attn_mask = causal_4d_mask.masked_fill(expanded_attn_mask.bool(), torch.finfo(dtype).min)
expanded_attn_mask = causal_4d_mask.masked_fill(
expanded_attn_mask.bool(), torch.finfo(dtype).min
)
return expanded_attn_mask
@ -118,19 +128,22 @@ class AttentionMaskConverter:
sliding_window: int | None = None,
):
bsz, tgt_len = input_ids_shape
mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
mask_cond = torch.arange(mask.size(-1), device=device)
mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device = device)
mask_cond = torch.arange(mask.size(-1), device = device)
mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
mask = mask.to(dtype)
if past_key_values_length > 0:
mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
mask = torch.cat(
[torch.zeros(tgt_len, past_key_values_length, dtype = dtype, device = device), mask],
dim = -1,
)
if sliding_window is not None:
diagonal = past_key_values_length - sliding_window - 1
context_mask = torch.tril(torch.ones_like(mask, dtype=torch.bool), diagonal=diagonal)
context_mask = torch.tril(torch.ones_like(mask, dtype = torch.bool), diagonal = diagonal)
if is_torchdynamo_compiling():
mask = mask.clone()
mask.masked_fill_(context_mask, torch.finfo(dtype).min)
@ -138,27 +151,28 @@ class AttentionMaskConverter:
return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
@staticmethod
def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = None):
def _expand_mask(
mask: torch.Tensor,
dtype: torch.dtype,
tgt_len: int | None = None,
):
bsz, src_len = mask.size()
tgt_len = tgt_len if tgt_len is not None else src_len
expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
inverted_mask = torch.tensor(1.0, dtype=dtype) - expanded_mask
inverted_mask = torch.tensor(1.0, dtype = dtype) - expanded_mask
return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
@staticmethod
def _unmask_unattended(
expanded_mask: torch.FloatTensor,
min_dtype: float,
):
def _unmask_unattended(expanded_mask: torch.FloatTensor, min_dtype: float):
if expanded_mask.dtype == torch.bool:
raise ValueError(
"AttentionMaskConverter._unmask_unattended expects a float `expanded_mask`, got a BoolTensor."
)
return expanded_mask.mul(~torch.all(expanded_mask == min_dtype, dim=-1, keepdim=True))
return expanded_mask.mul(~torch.all(expanded_mask == min_dtype, dim = -1, keepdim = True))
@staticmethod
def _ignore_causal_mask_sdpa(
@ -199,24 +213,28 @@ def _prepare_4d_causal_attention_mask_for_sdpa(
past_key_values_length: int,
sliding_window: int | None = None,
):
attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
attn_mask_converter = AttentionMaskConverter(is_causal = True, sliding_window = sliding_window)
key_value_length = input_shape[-1] + past_key_values_length
is_tracing_ = is_tracing(inputs_embeds)
ignore_causal_mask = AttentionMaskConverter._ignore_causal_mask_sdpa(
attention_mask=attention_mask,
inputs_embeds=inputs_embeds,
past_key_values_length=past_key_values_length,
sliding_window=sliding_window,
attention_mask = attention_mask,
inputs_embeds = inputs_embeds,
past_key_values_length = past_key_values_length,
sliding_window = sliding_window,
)
if ignore_causal_mask:
expanded_4d_mask = None
elif attention_mask is None:
expanded_4d_mask = attn_mask_converter.to_causal_4d(
input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device
input_shape[0],
input_shape[-1],
key_value_length,
dtype = inputs_embeds.dtype,
device = inputs_embeds.device,
)
else:
if attention_mask.dim() == 4:
@ -225,27 +243,39 @@ def _prepare_4d_causal_attention_mask_for_sdpa(
expanded_4d_mask = attn_mask_converter.to_4d(
attention_mask,
input_shape[-1],
dtype=inputs_embeds.dtype,
key_value_length=key_value_length,
dtype = inputs_embeds.dtype,
key_value_length = key_value_length,
)
if not is_tracing_ and expanded_4d_mask is not None and expanded_4d_mask.device.type in ["cuda", "xpu"]:
if (
not is_tracing_
and expanded_4d_mask is not None
and expanded_4d_mask.device.type in ["cuda", "xpu"]
):
expanded_4d_mask = AttentionMaskConverter._unmask_unattended(
expanded_4d_mask, min_dtype=torch.finfo(inputs_embeds.dtype).min
expanded_4d_mask, min_dtype = torch.finfo(inputs_embeds.dtype).min
)
return expanded_4d_mask
def _prepare_4d_attention_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = None):
return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
def _prepare_4d_attention_mask(
mask: torch.Tensor,
dtype: torch.dtype,
tgt_len: int | None = None,
):
return AttentionMaskConverter._expand_mask(mask = mask, dtype = dtype, tgt_len = tgt_len)
def _prepare_4d_attention_mask_for_sdpa(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = None):
def _prepare_4d_attention_mask_for_sdpa(
mask: torch.Tensor,
dtype: torch.dtype,
tgt_len: int | None = None,
):
_, key_value_length = mask.shape
tgt_len = tgt_len if tgt_len is not None else key_value_length
if not is_tracing(mask) and torch.all(mask == 1):
return None
return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
return AttentionMaskConverter._expand_mask(mask = mask, dtype = dtype, tgt_len = tgt_len)