Add BetterTransformer support for FlavaModel

docs/source/_toctree.yml

@@ -114,7 +114,7 @@
     - local: bettertransformer/tutorials/contribute
       title: How to add support for new architectures?
     title: Tutorials
-  title: BetterTransformer integration
+  title: BetterTransformer
   isExpanded: false
 - sections:
   - local: utils/dummy_input_generators

docs/source/bettertransformer/overview.mdx

@@ -28,11 +28,13 @@ The list of supported model below:
 - [BERT](https://arxiv.org/abs/1810.04805)
 - [BERT-generation](https://arxiv.org/abs/1907.12461)
 - [CamemBERT](https://arxiv.org/abs/1911.03894)
+- [CLIP](https://arxiv.org/abs/2103.00020)
 - [Data2VecText](https://arxiv.org/abs/2202.03555)
 - [DistilBert](https://arxiv.org/abs/1910.01108)
 - [DeiT](https://arxiv.org/abs/2012.12877)
 - [Electra](https://arxiv.org/abs/2003.10555)
 - [Ernie](https://arxiv.org/abs/1904.09223)
+- [Flava](https://arxiv.org/abs/2112.04482)
 - [FSMT](https://arxiv.org/abs/1907.06616)
 - [HuBERT](https://arxiv.org/pdf/2106.07447.pdf)
 - [LayoutLM](https://arxiv.org/abs/1912.13318)
@@ -52,6 +54,8 @@ The list of supported model below:
 - [XLMRoberta](https://arxiv.org/abs/1911.02116)
 - [YOLOS](https://arxiv.org/abs/2106.00666)
 
+Note that for encoder-decoder models, only the encoder part is supported by PyTorch's BetterTransformer for now.
+
 Let us know by opening an issue in 🤗 Optimum if you want more models to be supported, or check out the contribution guideline if you want to add it by yourself!
 
 ### Quick usage

docs/source/onnxruntime/package_reference/modeling_ort.mdx

@@ -20,6 +20,10 @@ specific language governing permissions and limitations under the License.
 
 [[autodoc]] onnxruntime.ORTModelForCausalLM
 
+## ORTModelForCustomTasks
+
+[[autodoc]] onnxruntime.ORTModelForCustomTasks
+
 ## ORTModelForFeatureExtraction
 
 [[autodoc]] onnxruntime.ORTModelForFeatureExtraction

optimum/bettertransformer/models/__init__.py

@@ -17,7 +17,9 @@
     AlbertLayerBetterTransformer,
     BartEncoderLayerBetterTransformer,
     BertLayerBetterTransformer,
+    CLIPLayerBetterTransformer,
     DistilBertLayerBetterTransformer,
+    FlavaLayerBetterTransformer,
     FSMTEncoderLayerBetterTransformer,
     MBartEncoderLayerBetterTransformer,
     ViltLayerBetterTransformer,
@@ -75,6 +77,14 @@
     # FSMTModel:
     "EncoderLayer": FSMTEncoderLayerBetterTransformer,
     "ViltLayer": ViltLayerBetterTransformer,
+    # Flava:
+    "FlavaLayer": FlavaLayerBetterTransformer,
+    # CLIP
+    "CLIPEncoderLayer": CLIPLayerBetterTransformer,
+}
+
+EXCLUDE_FROM_TRANSFORM = {
+    "clip": ["text_model"],  # text model uses causal attention, that is most likely not supported in BetterTransformer
 }
 
 

optimum/bettertransformer/models/base.py

@@ -14,12 +14,18 @@
 import torch
 import torch.nn as nn
 
+from ...utils import logging
+
 
 KNOWN_ACTIVATION_ATTRIBUTES = ["hidden_act", "activation", "act_fn", "activation_function"]
 KNOWN_POS_EMB_ATTRIBUTES = ["position_embedding_type"]
 KNOWN_NUM_LAYERS = ["num_hidden_layers", "num_layers", "encoder_layers", "n_layers"]
 
 SUPPORTED_ACTIVATION_FUNCTIONS = ["gelu", "relu", "gelu_new"]
+USE_AT_OWN_RISK_ACTIVATION_FUNCTIONS = ["quick_gelu"]
+
+
+logger = logging.get_logger(__name__)
 
 
 class BetterTransformerBaseLayer(nn.Module):
@@ -39,6 +45,10 @@ def __init__(self, config):
                 self.act_fn = getattr(config, attr)
                 break
 
+        # if act_fn not found in the config, fall back to the private `_get_activation_function` if available
+        if self.act_fn is None and hasattr(self, "_get_activation_function"):
+            self.act_fn = self._get_activation_function(config)
+
         # Get pos emb type
         for attr in KNOWN_POS_EMB_ATTRIBUTES:
             if hasattr(config, attr):
@@ -77,7 +87,12 @@ def validate_bettertransformer(self):
             raise ValueError("norm1_eps and norm2_eps must be equal for `BetterTransformer` integration.")
 
         # Check activation function
-        if self.act_fn not in SUPPORTED_ACTIVATION_FUNCTIONS:
+        if self.act_fn in USE_AT_OWN_RISK_ACTIVATION_FUNCTIONS:
+            logger.warning(
+                f"Overridding {self.act_fn} activation with gelu. Use the transformed model at your own risk, the output logits could be significantly different."
+            )
+            self.act_fn = "gelu"
+        elif self.act_fn not in SUPPORTED_ACTIVATION_FUNCTIONS:
             raise ValueError(
                 f"Activation function {self.act_fn} not supported" " for `BetterTransformer` integration."
             )

optimum/bettertransformer/models/encoder_models.py

@@ -11,12 +11,18 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from typing import TYPE_CHECKING
+
 import torch
 import torch.nn as nn
 
 from .base import BetterTransformerBaseLayer
 
 
+if TYPE_CHECKING:
+    from transformers import PretrainedConfig
+
+
 class AlbertLayerBetterTransformer(BetterTransformerBaseLayer):
     def __init__(self, albert_layer, config):
         r"""
@@ -1095,3 +1101,206 @@ def forward(self, hidden_states, attention_mask, position_bias=None, *_, **__):
         elif hidden_states.is_nested and self.is_last_layer:
             hidden_states = hidden_states.to_padded_tensor(0.0, original_shape)
         return (hidden_states, attention_mask)
+
+
+class CLIPLayerBetterTransformer(BetterTransformerBaseLayer):
+    def __init__(self, layer, config):
+        r"""
+        A simple conversion of the CLIPEncoderLayer to its `BetterTransformer` implementation.
+        **The implementation is valid only for the vision model, that does not use `causal_attention_mask`.**
+        Args:
+            layer (`torch.nn.Module`):
+                The original `CLIPEncoderLayer` where the weights needs to be retrieved.
+        """
+        super().__init__(config)
+        # In_proj layer
+        self.in_proj_weight = nn.Parameter(
+            torch.cat(
+                [
+                    layer.self_attn.q_proj.weight,
+                    layer.self_attn.k_proj.weight,
+                    layer.self_attn.v_proj.weight,
+                ]
+            )
+        )
+        self.in_proj_bias = nn.Parameter(
+            torch.cat(
+                [
+                    layer.self_attn.q_proj.bias,
+                    layer.self_attn.k_proj.bias,
+                    layer.self_attn.v_proj.bias,
+                ]
+            )
+        )
+
+        # Out proj layer
+        self.out_proj_weight = layer.self_attn.out_proj.weight
+        self.out_proj_bias = layer.self_attn.out_proj.bias
+
+        # Linear layer 1
+        self.linear1_weight = layer.mlp.fc1.weight
+        self.linear1_bias = layer.mlp.fc1.bias
+
+        # Linear layer 2
+        self.linear2_weight = layer.mlp.fc2.weight
+        self.linear2_bias = layer.mlp.fc2.bias
+
+        # Layer norm 1
+        self.norm1_eps = layer.layer_norm1.eps
+        self.norm1_weight = layer.layer_norm1.weight
+        self.norm1_bias = layer.layer_norm1.bias
+
+        # Layer norm 2
+        self.norm2_eps = layer.layer_norm2.eps
+        self.norm2_weight = layer.layer_norm2.weight
+        self.norm2_bias = layer.layer_norm2.bias
+
+        # Model hyper parameters
+        self.num_heads = layer.self_attn.num_heads
+        self.embed_dim = layer.self_attn.embed_dim
+
+        # Last step: set the last layer to `False` -> this will be set to `True` when converting the model
+        self.is_last_layer = False
+        self.norm_first = True
+
+        self.validate_bettertransformer()
+
+    def forward(self, hidden_states, attention_mask, *_, **__):
+        r"""
+        This is just a wrapper around the forward function proposed in:
+        https://github.com/huggingface/transformers/pull/19553
+        """
+        super().forward_checker()
+
+        # we expect attention_mask to be None in the vision model
+        if attention_mask is not None:
+            raise ValueError(
+                "Please do not use attention masks when using `BetterTransformer` converted vision models"
+            )
+
+        hidden_states = torch._transformer_encoder_layer_fwd(
+            hidden_states,
+            self.embed_dim,
+            self.num_heads,
+            self.in_proj_weight,
+            self.in_proj_bias,
+            self.out_proj_weight,
+            self.out_proj_bias,
+            self.use_gelu,
+            self.norm_first,
+            self.norm1_eps,
+            self.norm1_weight,
+            self.norm1_bias,
+            self.norm2_weight,
+            self.norm2_bias,
+            self.linear1_weight,
+            self.linear1_bias,
+            self.linear2_weight,
+            self.linear2_bias,
+            attention_mask,
+        )
+
+        return (hidden_states,)
+
+    def _get_activation_function(self, config: "PretrainedConfig"):
+        if hasattr(config, "vision_config") and hasattr(config, "text_config"):
+            assert config.vision_config.hidden_act == config.text_config.hidden_act
+            return config.vision_config.hidden_act
+        else:
+            return config.hidden_act
+
+
+class FlavaLayerBetterTransformer(BetterTransformerBaseLayer):
+    def __init__(self, flava_layer, config):
+        r"""
+        A simple conversion of the FlavaLayer to its `BetterTransformer` implementation.
+
+        Args:
+            flava_layer (`torch.nn.Module`):
+                The original `FlavaLayer` where the weights needs to be retrieved.
+        """
+        super().__init__(config.image_config)
+        # In_proj layer
+        self.in_proj_weight = nn.Parameter(
+            torch.cat(
+                [
+                    flava_layer.attention.attention.query.weight,
+                    flava_layer.attention.attention.key.weight,
+                    flava_layer.attention.attention.value.weight,
+                ]
+            )
+        )
+        self.in_proj_bias = nn.Parameter(
+            torch.cat(
+                [
+                    flava_layer.attention.attention.query.bias,
+                    flava_layer.attention.attention.key.bias,
+                    flava_layer.attention.attention.value.bias,
+                ]
+            )
+        )
+
+        # Out proj layer
+        self.out_proj_weight = flava_layer.attention.output.dense.weight
+        self.out_proj_bias = flava_layer.attention.output.dense.bias
+
+        # Linear layer 1
+        self.linear1_weight = flava_layer.intermediate.dense.weight
+        self.linear1_bias = flava_layer.intermediate.dense.bias
+
+        # Linear layer 2
+        self.linear2_weight = flava_layer.output.dense.weight
+        self.linear2_bias = flava_layer.output.dense.bias
+
+        # Layer norm 1
+        self.norm1_eps = flava_layer.layernorm_before.eps
+        self.norm1_weight = flava_layer.layernorm_before.weight
+        self.norm1_bias = flava_layer.layernorm_before.bias
+
+        # Layer norm 2
+        self.norm2_eps = flava_layer.layernorm_after.eps
+        self.norm2_weight = flava_layer.layernorm_after.weight
+        self.norm2_bias = flava_layer.layernorm_after.bias
+
+        # Model hyper parameters
+        self.num_heads = flava_layer.attention.attention.num_attention_heads
+        self.embed_dim = int(flava_layer.attention.attention.attention_head_size * self.num_heads)
+
+        # Last step: set the last layer to `False` -> this will be set to `True` when converting the model
+        self.is_last_layer = False
+        self.norm_first = True
+
+        self.validate_bettertransformer()
+
+    def forward(self, hidden_states, *_, **__):
+        r"""
+        This is just a wrapper around the forward function proposed in:
+        https://github.com/huggingface/transformers/pull/19553
+        """
+        super().forward_checker()
+        attention_mask = None
+
+        hidden_states = torch._transformer_encoder_layer_fwd(
+            hidden_states,
+            self.embed_dim,
+            self.num_heads,
+            self.in_proj_weight,
+            self.in_proj_bias,
+            self.out_proj_weight,
+            self.out_proj_bias,
+            self.use_gelu,
+            self.norm_first,
+            self.norm1_eps,
+            self.norm1_weight,
+            self.norm1_bias,
+            self.norm2_weight,
+            self.norm2_bias,
+            self.linear1_weight,
+            self.linear1_bias,
+            self.linear2_weight,
+            self.linear2_bias,
+            attention_mask,
+        )
+        if hidden_states.is_nested and self.is_last_layer:
+            hidden_states = hidden_states.to_padded_tensor(0.0)
+        return (hidden_states,)