From 881cea038e5abe8b36cbc47037457c4ab268feed Mon Sep 17 00:00:00 2001 From: Taranum Wasu Date: Wed, 8 Jul 2026 03:57:58 +0530 Subject: [PATCH] fix(attn-mask-compat): preserve JIT/FX tracing detection in fallback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Co-authored-by: Cursor --- tests/utils/test_attn_mask_compat.py | 65 ++++++++++++++++------- unsloth/models/_attn_mask_compat.py | 79 ++++++++++++++++------------ 2 files changed, 91 insertions(+), 53 deletions(-) diff --git a/tests/utils/test_attn_mask_compat.py b/tests/utils/test_attn_mask_compat.py index a1959714b3..2c46f4cdc8 100644 --- a/tests/utils/test_attn_mask_compat.py +++ b/tests/utils/test_attn_mask_compat.py @@ -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. diff --git a/unsloth/models/_attn_mask_compat.py b/unsloth/models/_attn_mask_compat.py index bba78697f5..3b4608654e 100644 --- a/unsloth/models/_attn_mask_compat.py +++ b/unsloth/models/_attn_mask_compat.py @@ -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)