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,
 )
 
@@ -257,7 +258,10 @@ def _ignore_causal_mask_sdpa(
     # 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):
+    # Under `torch.compile` we can still skip, but only if we do not have to read the values of the `padding_mask`.
+    # NOTE: this requires torch>=2.14. Before pytorch#176499, dynamo replaced `torch.compiler.is_exporting()` by a
+    # constant `True`, so older versions keep the previous behavior of never skipping while compiling.
+    if is_torchdynamo_exporting() or (padding_mask is not None and is_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 +295,7 @@ def _can_skip_bidirectional_mask_xpu(
     - Skip if no padding and no local attention constraint
     """
 
-    if is_tracing(padding_mask):
+    if is_torchdynamo_exporting() or (padding_mask is not None and is_tracing(padding_mask)):
         return False
 
     # Check local attention constraint (same as CUDA)
@@ -327,7 +331,7 @@ def _ignore_bidirectional_mask_sdpa(
     # 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 (is_torchdynamo_exporting() or (padding_mask is not None and is_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))
 
+    # Before https://github.com/pytorch/pytorch/pull/176499 (torch 2.14), dynamo replaced calls to
+    # `torch.compiler.is_exporting()` by a constant `True`, so we treated `torch.compile` as export and never
+    # skipped. On older versions the mask is therefore 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")