Remove ambiguous `padding_mask` and instead use a 2D->4D Attn Mask Mapper (#26792)

* [Attn Mask Converter] refactor attn mask

* up

* Apply suggestions from code review
Co-authored-by: fxmarty <9808326+fxmarty@users.noreply.github.com>

* improve

* rename

* better cache

* renaming

* improve more

* improve

* fix bug

* finalize

* make style & make fix-copies

* correct more

* start moving attention_mask

* fix llama

* improve falcon

* up

* improve more

* improve more

* Update src/transformers/models/owlv2/modeling_owlv2.py

* make style

* make style

* rename to converter

* Apply suggestions from code review

---------
Co-authored-by: fxmarty <9808326+fxmarty@users.noreply.github.com>

src/transformers/models/deprecated/open_llama/modeling_open_llama.py

@@ -560,7 +560,7 @@ class OpenLlamaModel(OpenLlamaPreTrainedModel):
     def set_input_embeddings(self, value):
         self.embed_tokens = value
 
-    # Copied from transformers.models.llama.modeling_llama.LlamaModel._prepare_decoder_attention_mask
+    # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
     def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
         # create causal mask
         # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]

src/transformers/models/persimmon/modeling_persimmon.py

@@ -548,7 +548,6 @@ class PersimmonModel(PersimmonPreTrainedModel):
         config: PersimmonConfig
     """
 
-    # Copied from transformers.models.llama.modeling_llama.LlamaModel.__init__ with LLAMA->PERSIMMON,Llama->Persimmon,PersimmonRMSNorm->nn.LayerNorm,norm->final_layernorm,rms_final_layernorm_eps->layer_norm_eps
     def __init__(self, config: PersimmonConfig):
         super().__init__(config)
         self.padding_idx = config.pad_token_id

tests/models/llama/test_modeling_llama.py

@@ -39,6 +39,149 @@ if is_torch_available():
         LlamaModel,
         LlamaTokenizer,
     )
+    from transformers.models.llama.modeling_llama import AttnMaskConverter
+
+
+@require_torch
+class AttentionMaskTester(unittest.TestCase):
+    def check_non_causal(self, bsz, q_len, kv_len, mask_2d, mask_4d):
+        mask_indices = (mask_2d != 1)[:, None].broadcast_to((bsz, q_len, kv_len))
+        mask_4d_values = mask_4d[:, 0][mask_indices]
+        is_inf = mask_4d_values == -float("inf")
+        is_min = mask_4d_values == torch.finfo(mask_4d.dtype).min
+        assert torch.logical_or(is_inf, is_min).all()
+
+    def check_to_4d(self, mask_converter, q_len, kv_len, additional_mask=None, bsz=3):
+        mask_2d = torch.ones((bsz, kv_len), device=torch_device, dtype=torch.long)
+
+        if additional_mask is not None:
+            for bsz_idx, seq_idx in additional_mask:
+                mask_2d[bsz_idx, seq_idx] = 0
+
+        mask_4d = mask_converter.to_4d(mask_2d, query_length=q_len, key_value_length=kv_len)
+
+        assert mask_4d.shape == (bsz, 1, q_len, kv_len)
+
+        context = mask_converter.sliding_window
+        if mask_converter.is_causal and context is None:
+            # k * (k+1) / 2 tokens are masked in triangualar masks
+            num_tokens_masked = bsz * (q_len * (q_len - 1) // 2)
+
+            if 0 not in mask_2d:
+                assert (mask_4d != 0).sum().cpu().item() == num_tokens_masked
+            if 0 in mask_2d:
+                # at least causal mask + maybe more
+                assert (mask_4d != 0).sum().cpu().item() >= num_tokens_masked
+                self.check_non_causal(bsz, q_len, kv_len, mask_2d, mask_4d)
+        elif not mask_converter.is_causal and context is None:
+            if 0 not in mask_2d:
+                assert (mask_4d != 0).sum().cpu().item() == 0
+            if 0 in mask_2d:
+                self.check_non_causal(bsz, q_len, kv_len, mask_2d, mask_4d)
+        elif mask_converter.is_causal and context is not None:
+            # k * (k+1) / 2 tokens are masked in triangualar masks
+            num_tokens_masked = (q_len * (q_len - 1) // 2) + self.compute_num_context_mask(kv_len, context, q_len)
+            num_tokens_masked = bsz * num_tokens_masked
+
+            if 0 not in mask_2d:
+                assert (mask_4d != 0).sum().cpu().item() == num_tokens_masked
+            if 0 in mask_2d:
+                # at least causal mask + maybe more
+                assert (mask_4d != 0).sum().cpu().item() >= num_tokens_masked
+                self.check_non_causal(bsz, q_len, kv_len, mask_2d, mask_4d)
+
+    def check_to_causal(self, mask_converter, q_len, kv_len, bsz=3):
+        mask_4d = mask_converter.to_causal_4d(bsz, query_length=q_len, key_value_length=kv_len, device=torch_device)
+
+        if q_len == 1 and mask_converter.sliding_window is None:
+            # no causal mask if q_len is 1
+            assert mask_4d is None
+            return
+
+        context = mask_converter.sliding_window
+        if mask_converter.is_causal and context is None:
+            # k * (k+1) / 2 tokens are masked in triangualar masks
+            num_tokens_masked = bsz * (q_len * (q_len - 1) // 2)
+
+            assert (mask_4d != 0).sum().cpu().item() == num_tokens_masked
+        elif not mask_converter.is_causal and context is None:
+            assert (mask_4d != 0).sum().cpu().item() == 0
+        elif mask_converter.is_causal and context is not None:
+            # k * (k+1) / 2 tokens are masked in triangualar masks
+            num_tokens_masked = (q_len * (q_len - 1) // 2) + self.compute_num_context_mask(kv_len, context, q_len)
+            num_tokens_masked = bsz * num_tokens_masked
+
+            assert (mask_4d != 0).sum().cpu().item() == num_tokens_masked
+
+    def compute_num_context_mask(self, kv_len, context, q_len):
+        # This function computes the # of attention tokens that are added for
+        # the sliding window
+        c_mask_len = kv_len - context
+        num_mask_triangle = c_mask_len * (c_mask_len + 1) // 2
+        cut_mask_len = max(c_mask_len - q_len, 0)
+        num_cut_mask = cut_mask_len * (cut_mask_len + 1) // 2
+        return num_mask_triangle - num_cut_mask
+
+    def test_2d_to_4d_causal(self):
+        mask_converter = AttnMaskConverter(is_causal=True)
+
+        # auto-regressive use case
+        self.check_to_4d(mask_converter, q_len=1, kv_len=7)
+        # special auto-regressive case
+        self.check_to_4d(mask_converter, q_len=3, kv_len=7)
+        # non auto-regressive case
+        self.check_to_4d(mask_converter, q_len=7, kv_len=7)
+
+        # same with extra attention masks
+        self.check_to_4d(mask_converter, q_len=1, kv_len=7, additional_mask=[(0, 2), (1, 3), (2, 0)])
+        self.check_to_4d(mask_converter, q_len=3, kv_len=7, additional_mask=[(0, 2), (1, 3), (2, 0)])
+        self.check_to_4d(mask_converter, q_len=7, kv_len=7, additional_mask=[(0, 2), (1, 3), (2, 0)])
+
+    def test_2d_to_4d(self):
+        torch.ones((3, 7), device=torch_device, dtype=torch.long)
+        mask_converter = AttnMaskConverter(is_causal=False)
+
+        # non auto-regressive case
+        self.check_to_4d(mask_converter, q_len=7, kv_len=7)
+
+        # same with extra attention masks
+        self.check_to_4d(mask_converter, q_len=7, kv_len=7, additional_mask=[(0, 2), (1, 3), (2, 0)])
+
+    def test_2d_to_4d_causal_sliding(self):
+        torch.ones((3, 7), device=torch_device, dtype=torch.long)
+        mask_converter = AttnMaskConverter(is_causal=True, sliding_window=5)
+
+        # auto-regressive use case
+        self.check_to_4d(mask_converter, q_len=1, kv_len=7)
+        # special auto-regressive case
+        self.check_to_4d(mask_converter, q_len=3, kv_len=7)
+        # non auto-regressive case
+        self.check_to_4d(mask_converter, q_len=7, kv_len=7)
+
+        # same with extra attention masks
+        self.check_to_4d(mask_converter, q_len=1, kv_len=7, additional_mask=[(0, 2), (1, 3), (2, 0)])
+        self.check_to_4d(mask_converter, q_len=3, kv_len=7, additional_mask=[(0, 2), (1, 3), (2, 0)])
+        self.check_to_4d(mask_converter, q_len=7, kv_len=7, additional_mask=[(0, 2), (1, 3), (2, 0)])
+
+    def test_causal_mask(self):
+        mask_converter = AttnMaskConverter(is_causal=True)
+
+        # auto-regressive use case
+        self.check_to_causal(mask_converter, q_len=1, kv_len=7)
+        # special auto-regressive case
+        self.check_to_causal(mask_converter, q_len=3, kv_len=7)
+        # non auto-regressive case
+        self.check_to_causal(mask_converter, q_len=7, kv_len=7)
+
+    def test_causal_mask_sliding(self):
+        mask_converter = AttnMaskConverter(is_causal=True, sliding_window=3)
+
+        # auto-regressive use case
+        self.check_to_causal(mask_converter, q_len=1, kv_len=7)
+        # special auto-regressive case
+        self.check_to_causal(mask_converter, q_len=3, kv_len=7)
+        # non auto-regressive case
+        self.check_to_causal(mask_converter, q_len=7, kv_len=7)
 
 
 class LlamaModelTester: