Repository: huggingface/transformers
Pull request: Allow bi-directional attention for all models

# What does this PR do?

Allow the `is_causal` kwarg and config attribute to make well-behaved decoder-only models act as encoders

Changed files:
filename: src/transformers/masking_utils.py
patch: @@ -865,6 +865,17 @@ def create_causal_mask(
             An optional mask function to combine with the causal mask function (by doing the intersection of both). This is
             useful to easily overlay another mask on top of the causal one, for example for image tokens handling.
     """
+    # Power feature: if `is_causal` is False, then fallback to bi-directional mask for bi-directional attention.
+    # It allows to use decoder-only models with bi-directional attention as well
+    if not getattr(config, "is_causal", True):
+        return create_bidirectional_mask(
+            config,
+            input_embeds,
+            attention_mask,
+            or_mask_function=or_mask_function,
+            and_mask_function=and_mask_function,
+        )
+
     # If we have an hybrid cache structure, here we want to create the mask for the full layers
     if hasattr(past_key_values, "is_sliding") and False in past_key_values.is_sliding:
         layer_idx = past_key_values.is_sliding.index(False)
@@ -1057,6 +1068,17 @@ def create_sliding_window_causal_mask(
             An optional mask function to combine with the sliding causal mask function (by doing the intersection of both). This is
             useful to easily overlay another mask on top of the sliding causal one, for example for image tokens handling.
     """
+    # Power feature: if `is_causal` is False, then fallback to bi-directional mask for bi-directional attention
+    # It allows to use decoder-only models with bi-directional attention as well
+    if not getattr(config, "is_causal", True):
+        return create_bidirectional_sliding_window_mask(
+            config,
+            input_embeds,
+            attention_mask,
+            or_mask_function=or_mask_function,
+            and_mask_function=and_mask_function,
+        )
+
     # If we have an hybrid cache structure, here we want to create the mask for the sliding layers
     if hasattr(past_key_values, "is_sliding") and True in past_key_values.is_sliding:
         layer_idx = past_key_values.is_sliding.index(True)

filename: src/transformers/utils/generic.py
patch: @@ -749,6 +749,8 @@ class TransformersKwargs(TypedDict, total=False):
             Maximum sequence length for key state.
         position_ids (`torch.LongTensor`, *optional*)
             Indices of positions of each input sequence tokens.
+        is_causal (`bool`, *optional*)
+            Can be set to False to enable bi-directional attention, i.e. use decoder Attention modules as encoders.
     """
 
     num_items_in_batch: Optional["torch.Tensor"]
@@ -760,6 +762,7 @@ class TransformersKwargs(TypedDict, total=False):
     max_length_q: int | None
     max_length_k: int | None
     position_ids: Optional["torch.LongTensor"]
+    is_causal: bool | None
 
 
 def is_timm_config_dict(config_dict: dict[str, Any]) -> bool:
@@ -919,6 +922,19 @@ def wrapper(self, *args, **kwargs):
             if return_dict is None:
                 return_dict = getattr(self.config, "return_dict", True)
 
+            # Maybe temporarily overwrite config value to create the correct mask - kwarg takes precedence
+            is_causal = kwargs.get("is_causal")
+            if is_causal is None:
+                is_causal = getattr(self.config, "is_causal", None)
+
+            if is_causal is not None:
+                is_causal_in_config = hasattr(self.config, "is_causal")
+                if is_causal_in_config:
+                    is_causal_original_value = self.config.is_causal
+                # Set it to both config and kwargs (it's needed in both, and can come from only 1 of the sources)
+                self.config.is_causal = is_causal
+                kwargs["is_causal"] = is_causal
+
             all_args = kwargs.copy()
             if "kwargs" in all_args:
                 for k, v in all_args["kwargs"].items():
@@ -1017,6 +1033,13 @@ def wrapped_forward(*args, **kwargs):
             for module, original_forward in monkey_patched_layers:
                 module.forward = original_forward
 
+            # Restore original config value
+            if is_causal is not None:
+                if is_causal_in_config:
+                    self.config.is_causal = is_causal_original_value
+                else:
+                    del self.config.is_causal
+
             # Inject collected outputs into model output
             for key in collected_outputs:
                 if key == "hidden_states":

filename: tests/utils/test_modeling_utils.py
patch: @@ -2304,6 +2304,81 @@ def test_composite_model_inherit_properties(self):
         self.assertEqual(model._keep_in_fp32_modules, {"linear", "head"})  # language model + composite model
         self.assertEqual(model._keep_in_fp32_modules_strict, {"simple"})  # vision model
 
+    @parameterized.expand([("sdpa",), ("flash_attention_2",)])
+    def test_decoder_only_model_can_be_used_as_encoder(self, attn_implementation: str):
+        """Test that most well-behaved decoder models can be used as encoders through the `is_causal` kwarg/config.
+        Note that it's enough to test it on Llama, as the entry points are all through general code
+        (masking_utils.py + `check_model_inputs` decorator). This makes it easier as the model need to use both the
+        mask API from masking_utils.py and the decorator as mentionned above, and we don't know what models follow that
+        standard exactly (so we cannot make it easily a common model test)."""
+        if attn_implementation == "flash_attention_2" and not is_flash_attn_2_available():
+            self.skipTest("FA2 not available")
+
+        from transformers import LlamaConfig, LlamaModel
+        from transformers.masking_utils import create_bidirectional_mask
+
+        config = LlamaConfig(
+            num_hidden_layers=2,
+            hidden_size=32,
+            intermediate_size=64,
+            vocab_size=100,
+            attn_implementation=attn_implementation,
+        )
+        model = LlamaModel(copy.deepcopy(config)).to(device=torch_device, dtype=torch.bfloat16)
+
+        # Create inputs, making sure we use padding to verify that mask creation accounts for it correctly
+        input_ids = torch.randint(5, 95, (2, 17), device=torch_device)
+        attention_mask = torch.ones_like(input_ids, device=torch_device)
+        attention_mask[1, 0:3] = 0
+
+        # The original `create_causal_mask` used in modeling_llama forward more kwargs than `create_bidirectional_mask`,
+        # so we need this one instead to absorb them
+        def create_bidirectional_mask_with_kwargs(
+            config,
+            input_embeds,
+            attention_mask,
+            encoder_hidden_states=None,
+            or_mask_function=None,
+            and_mask_function=None,
+            **kwargs,
+        ):
+            return create_bidirectional_mask(
+                config, input_embeds, attention_mask, encoder_hidden_states, or_mask_function, and_mask_function
+            )
+
+        # Explicitly monkey patch the mask creation function + forward the is_causal kwarg to get the expected result
+        # from the model behaving as encoder instead of decoder
+        with patch(
+            "transformers.models.llama.modeling_llama.create_causal_mask", new=create_bidirectional_mask_with_kwargs
+        ):
+            reference = model(input_ids, attention_mask=attention_mask, is_causal=False).last_hidden_state
+            without_kwarg = model(input_ids, attention_mask=attention_mask).last_hidden_state
+
+        # Here, since we have padding, the mask created should never be None. Since the mask is never None, the sdpa
+        # backend will always use `is_causal=False`, so both should be strictly equivalent
+        if attn_implementation == "sdpa":
+            torch.testing.assert_close(reference, without_kwarg)
+        # But FA2 relies solely on the `is_causal` kwarg to decide how to dispatch, as it will use varlen since we
+        # have padding, so both won't be equivalent at all
+        else:
+            # Everything should be different (we only test the maximum of the diff to avoid flakyness)
+            self.assertTrue(torch.abs(reference - without_kwarg).max() >= 1e-1)
+
+        # Now if we simply forward the kwarg with the usual mask function, it should still work the exact same
+        with_kwarg_only = model(input_ids, attention_mask=attention_mask, is_causal=False).last_hidden_state
+        torch.testing.assert_close(reference, with_kwarg_only)
+
+        # Now, if we use the usual forward, the model should behave normally as a decoder, and output should be
+        # completely different
+        as_decoder = model(input_ids, attention_mask=attention_mask).last_hidden_state
+        # Everything should be different (we only test the maximum of the diff to avoid flakyness)
+        self.assertTrue(torch.abs(reference - as_decoder).max() >= 1e-1)
+
+        # It should also work with it in the config
+        model.config.is_causal = False
+        with_config_only = model(input_ids, attention_mask=attention_mask).last_hidden_state
+        torch.testing.assert_close(reference, with_config_only)
+
 
 @slow
 @require_torch