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

fix(attn-mask-compat): preserve JIT/FX tracing detection in fallback

Browse source 
Codex review follow-up on PR #6880 (P2). The previous fallback
(``def is_tracing(tensor=None): return is_torchdynamo_compiling()``)
was strictly less conservative than what upstream
``transformers==4.51.3`` did inline before ``is_tracing`` was added
to ``transformers.utils.import_utils``. The pre-4.52 upstream
expression was:

    is_tracing = torch.jit.is_tracing() or isinstance(
        inputs_embeds, torch.fx.Proxy
    ) or is_torchdynamo_compiling()

The previous fallback only consulted Dynamo. That meant callers
tracing or exporting under transformers 4.51.x would silently hit
the data-dependent ``torch.all(attention_mask == 1)`` branch in
``_ignore_causal_mask_sdpa`` (and the equivalent in
``_prepare_4d_attention_mask_for_sdpa``) — failing on proxy control
flow or baking the wrong SDPA causal-mask path.

The local fallback now mirrors the legacy upstream expression:

- ``torch.jit.is_tracing()`` for ``torch.jit.trace`` /
  ``torch.jit.script`` flows.
- ``isinstance(tensor, torch.fx.Proxy)`` for ``symbolic_trace`` and
  ``torch.export`` paths that don't go through Dynamo.
- ``is_torchdynamo_compiling()`` for ``torch.compile`` and
  ``torch._dynamo`` paths.

The CUDA stream capture, FakeTensor, and JAX (torchax) checks the
modern ``is_tracing`` does are out of scope: those need newer
``import_utils`` helpers, and the conservative dynamo fallback is
the right choice when those helpers aren't available. This matches
the behavior of the upstream ``is_tracing`` helpers that landed in
4.52 in the first place — those were new additions, not behaviour
that pre-existed in 4.51.x.

Tests:
- ``test_import_falls_back_when_is_tracing_missing`` now exercises
  all three branches (dynamo idle → False; ``torch.fx.Proxy`` arg →
  True; patched ``torch.jit.is_tracing()`` → True).
- All 26 tests in ``tests/utils/test_attn_mask_compat.py`` pass.
- ``ruff check`` and ``ruff format`` clean on both files.

Signed-off-by: Taranum Wasu <taranumwasu@Taranums-MacBook-Pro.local>
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Taranum Wasu 2026-07-08 03:57:58 +05:30
commit 881cea038e
2 changed files with 91 additions and 53 deletions
Download patch file Download diff file Expand all files Collapse all files

65
tests/utils/test_attn_mask_compat.py
View file

@ -33,13 +33,13 @@ 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)
@ -62,23 +62,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:
@ -106,7 +106,7 @@ def test_prepare_4d_causal_attention_mask_for_sdpa_matches_transformers(
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)
@ -115,7 +115,7 @@ def test_prepare_4d_causal_attention_mask_for_sdpa_matches_transformers(
(batch_size, query_length),
inputs_embeds,
past_length,
sliding_window = 3,
sliding_window=3,
)
actual = compat._prepare_4d_causal_attention_mask_for_sdpa(
@ -123,7 +123,7 @@ def test_prepare_4d_causal_attention_mask_for_sdpa_matches_transformers(
(batch_size, query_length),
inputs_embeds,
past_length,
sliding_window = 3,
sliding_window=3,
)
if expected is None:
@ -138,14 +138,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
@ -159,7 +159,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 == []
@ -174,8 +174,14 @@ def test_import_falls_back_when_is_tracing_missing():
on the lower bound tested in CI (`__from_pyproject__` matrix cell).
Reload the module with `is_tracing` removed from the namespace and confirm
the local fallback is used (returns False when dynamo is idle, without
raising).
the local fallback is used. The fallback must mirror the legacy
`transformers==4.51.3` inline expression
(``torch.jit.is_tracing() or isinstance(tensor, torch.fx.Proxy) or
is_torchdynamo_compiling()``) so the data-dependent ``torch.all(...)``
branches in the mask helpers continue to be skipped during JIT trace,
symbolic trace, and Dynamo compilation — otherwise tracing/exporting
these models on transformers 4.51.x either fails on proxy control flow
or bakes the wrong SDPA causal-mask path.
"""
fake_import_utils = types.ModuleType("transformers.utils.import_utils")
@ -203,7 +209,28 @@ def test_import_falls_back_when_is_tracing_missing():
"would not exercise the fallback path"
)
# Falls back to the local definition.
# Dynamo idle and no JIT/FX active → False.
assert reloaded.is_tracing() is False
# Sanity: accepts an optional tensor positional arg without raising.
assert reloaded.is_tracing(torch.zeros(1)) is False
# ``torch.fx.Proxy`` should be detected even when Dynamo is idle, since
# symbolic_trace / export-only paths don't go through dynamo. Construct
# the Proxy from a real fx.Graph node (passing a Tensor directly to
# ``Proxy(...)`` is a common foot-gun that raises AttributeError).
fx_graph = torch.fx.Graph()
fx_node = fx_graph.create_node("call_function", torch.zeros, (torch.zeros(1).shape,))
proxy = torch.fx.Proxy(fx_node)
assert reloaded.is_tracing(proxy) is True
# ``torch.jit.is_tracing()`` should be detected via patch.
with mock.patch("torch.jit.is_tracing", return_value=True):
assert reloaded.is_tracing() is True
# Dynamo compilation is also covered (the fallback calls
# ``is_torchdynamo_compiling`` from the module-level import, which is
# bound at fallback-definition time — exactly the same import binding
# that the real ``is_tracing`` uses). We don't re-test the dynamo path
# here because it's already exercised by the upstream test suite, and
# patching the import after the module is loaded would not affect the
# closure's reference.

79
unsloth/models/_attn_mask_compat.py
View file

@ -36,21 +36,32 @@ try:
# `is_tracing` was added to `transformers.utils.import_utils` in 4.52
# (commit that introduced `_prepare_4d_attention_mask_for_sdpa` rewrites).
# Unsloth's declared lower bound is `transformers>=4.51.3`, so import
# defensively and fall back to a conservative local implementation when
# the symbol is not exported — matching what upstream Transformers did
# before `is_tracing` existed (only `is_torchdynamo_compiling`).
# defensively and fall back to a local implementation when the symbol
# is not exported. The fallback mirrors the tracing-detection expression
# that upstream `transformers.modeling_attn_mask_utils` used inline before
# `is_tracing` was added — `torch.jit.is_tracing() or
# isinstance(inputs_embeds, torch.fx.Proxy) or is_torchdynamo_compiling()`
# — so the data-dependent `torch.all(...)` branches in the mask helpers
# continue to be skipped during JIT trace / symbolic trace / Dynamo
# compilation, preserving the SDPA path selection.
from transformers.utils.import_utils import is_tracing # type: ignore[attr-defined]
except ImportError:
def is_tracing(tensor = None) -> bool: # type: ignore[no-redef]
def is_tracing(tensor=None) -> bool: # type: ignore[no-redef]
"""Local fallback for transformers < 4.52.
Returns True only when Dynamo is actively compiling. Other tracing
backends (JIT, CUDA stream capture, FakeTensor, JAX) are not
detectable via `import_utils` in these older releases; we conservatively
treat them as "not tracing", matching the pre-`is_tracing` upstream
behavior where these checks were guarded by `is_torchdynamo_compiling`.
Returns True when the active context is any of: ``torch.jit.trace``,
``torch.fx.symbolic_trace``, or Dynamo compilation. Other tracing
backends that the modern ``transformers.utils.import_utils.is_tracing``
detects (CUDA stream capture, FakeTensor, JAX via torchax) cannot be
detected without newer ``import_utils`` helpers; for those we fall
back to the dynamo check, which matches the conservative pre-4.52
upstream behavior on the supported lower bound.
"""
if torch.jit.is_tracing():
return True
if tensor is not None and isinstance(tensor, torch.fx.Proxy):
return True
return is_torchdynamo_compiling()
@ -93,9 +104,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
@ -120,9 +131,9 @@ 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(
@ -130,7 +141,7 @@ class AttentionMaskConverter:
)
expanded_attn_mask = self._expand_mask(
attention_mask_2d, dtype, tgt_len = input_shape[-1]
attention_mask_2d, dtype, tgt_len=input_shape[-1]
).to(attention_mask_2d.device)
if causal_4d_mask is not None:
@ -149,22 +160,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,
[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)
@ -193,7 +204,7 @@ class AttentionMaskConverter:
"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(
@ -234,17 +245,17 @@ 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:
@ -254,8 +265,8 @@ def _prepare_4d_causal_attention_mask_for_sdpa(
input_shape[0],
input_shape[-1],
key_value_length,
dtype = inputs_embeds.dtype,
device = inputs_embeds.device,
dtype=inputs_embeds.dtype,
device=inputs_embeds.device,
)
else:
if attention_mask.dim() == 4:
@ -264,8 +275,8 @@ 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 (
@ -274,7 +285,7 @@ def _prepare_4d_causal_attention_mask_for_sdpa(
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
@ -285,7 +296,7 @@ def _prepare_4d_attention_mask(
dtype: torch.dtype,
tgt_len: int | None = 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)
def _prepare_4d_attention_mask_for_sdpa(
@ -299,4 +310,4 @@ def _prepare_4d_attention_mask_for_sdpa(
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)