Fix mask creation not being skipped under torch.compile

src/transformers/masking_utils.py

@@ -25,6 +25,7 @@
 from .utils.import_utils import (
     is_torch_flex_attn_available,
     is_torch_greater_or_equal,
+    is_torchdynamo_exporting,
     is_tracing,
 )
 
@@ -233,6 +234,19 @@ def fast_all(tensor: torch.BoolTensor) -> torch.BoolTensor:
     return tensor.sum() == tensor.numel()
 
 
+def _cannot_decide_skip_while_tracing(padding_mask: torch.Tensor | None) -> bool:
+    """Whether the skip conditions cannot be evaluated because we are tracing.
+
+    Export hard-codes the decision into the exported forward, which is in general wrong (see
+    https://github.com/pytorch/pytorch/issues/108108). Reading the values of a `padding_mask` is
+    data-dependent. Without one, the remaining conditions are static and dynamo simply guards on them.
+
+    NOTE: until torch 2.14 (pytorch#176499), dynamo also reported exporting under `torch.compile`, so
+    older versions keep the previous behavior of never skipping while compiling.
+    """
+    return is_torchdynamo_exporting() or (padding_mask is not None and is_tracing(padding_mask))
+
+
 def _ignore_causal_mask_sdpa(
     padding_mask: torch.Tensor | None,
     q_length: int,
@@ -253,11 +267,7 @@ def _ignore_causal_mask_sdpa(
         mask_indices = torch.arange(kv_length, device=padding_mask.device) + kv_offset
         padding_mask = padding_mask[:, mask_indices]
 
-    # When using `torch.export` or `torch.onnx.dynamo_export`, we must pass an example input, and `is_causal` behavior is
-    # hard-coded to the forward. If a user exports a model with query_length > 1, the exported model will hard-code `is_causal=True`
-    # which is in general wrong (see https://github.com/pytorch/pytorch/issues/108108). Thus, we only set
-    # `ignore_causal_mask = True` if we are not tracing
-    if is_tracing(padding_mask):
+    if _cannot_decide_skip_while_tracing(padding_mask):
         return False
     # In this case, we need to add special patterns to the mask no matter what, so we cannot use any of the later skip conditions
     if local_attention_size is not None and kv_length >= local_attention_size:
@@ -291,7 +301,7 @@ def _can_skip_bidirectional_mask_xpu(
     - Skip if no padding and no local attention constraint
     """
 
-    if is_tracing(padding_mask):
+    if _cannot_decide_skip_while_tracing(padding_mask):
         return False
 
     # Check local attention constraint (same as CUDA)
@@ -324,10 +334,8 @@ def _ignore_bidirectional_mask_sdpa(
         # - Skip if no padding and no local attention constraint
         return _can_skip_bidirectional_mask_xpu(padding_mask, kv_length, local_attention_size)
 
-    # When using `torch.export` or `torch.onnx.dynamo_export`, we need to avoid to check the contents of the mask;
-    # otherwise, we will encounter dynamic control flows
     if (
-        not is_tracing(padding_mask)
+        not _cannot_decide_skip_while_tracing(padding_mask)
         and (padding_mask is None or padding_mask.all())
         # in this case we need to add special patterns to the mask so cannot be skipped otherwise
         and (local_attention_size is None or kv_length < local_attention_size)

tests/utils/test_masking_utils.py

@@ -17,6 +17,7 @@
 from transformers.testing_utils import (
     is_torch_available,
     require_torch,
+    require_torch_greater_or_equal,
     torch_device,
 )
 
@@ -177,6 +178,51 @@ def test_nonpacked_sequence_mask_skip(self):
         # cannot be skipped under compile, should result into a triu mask
         self.assertTrue(torch.equal(~torch.ones(*causal_mask.shape).triu(diagonal=1).bool(), causal_mask))
 
+    # `torch.compiler.is_exporting()` was hard-coded to `True` by dynamo, i.e. it also reported exporting under
+    # `torch.compile`, until https://github.com/pytorch/pytorch/pull/176499 (torch 2.14). On older versions we keep
+    # the conservative behavior of never skipping while tracing, so the mask is still materialized.
+    @require_torch_greater_or_equal("2.14")
+    def test_mask_skip_without_padding_mask_under_compile(self):
+        """
+        Checks whether the mask creation can still be skipped under `torch.compile` if we have no padding mask at all.
+        Whether a padding mask is provided is a static property that dynamo guards on - only the checks reading its
+        values are data-dependent, and have to be skipped while tracing.
+        """
+        config = LlamaConfig()
+        config._attn_implementation = "sdpa"
+
+        batch_size = 2
+        sequence_length = 10
+        inputs_embeds = torch.empty((batch_size, sequence_length, 8), dtype=torch.float16, device=torch_device)
+        padded_mask = torch.ones(batch_size, sequence_length, dtype=torch.long, device=torch_device)
+        padded_mask[0, :3] = 0
+
+        def create_masks(attention_mask):
+            causal_mask = create_causal_mask(
+                config=config,
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                past_key_values=None,
+            )
+            bidirectional_mask = create_bidirectional_mask(
+                config=config,
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+            )
+            return causal_mask, bidirectional_mask
+
+        compiled_create_masks = torch.compile(create_masks, fullgraph=True)
+
+        # Without a padding mask, both masks are skipped in eager as well as under compile
+        self.assertEqual(create_masks(None), (None, None))
+        self.assertEqual(compiled_create_masks(None), (None, None))
+
+        # With a padding mask, the masks are materialized. Under compile, reading its values must be skipped
+        # instead of raising a data-dependent control flow error
+        for causal_mask, bidirectional_mask in (create_masks(padded_mask), compiled_create_masks(padded_mask)):
+            self.assertIsNotNone(causal_mask)
+            self.assertIsNotNone(bidirectional_mask)
+
     def test_chunked_mask_with_left_padding_and_large_prefill(self):
         # Make sure we have an attention_chunk_size in the config
         config = LlamaConfig(attention_chunk_size=3, attn_implementation="sdpa")