Repository: huggingface/optimum

Issue:
Community contribution - `BetterTransformer` integration for more models!
## `BetterTransformer` integration for more models!

`BetterTransformer` API provides faster inference on CPU & GPU through a simple interface! 

Models can benefit from very interesting speedups using a one liner and by making sure to install the latest version of PyTorch. A complete guideline on how to convert a new model has been created on the [BetterTransformer documentation](https://huggingface.co/docs/optimum/bettertransformer/tutorials/contribute)!

Here is a list of models that could be potentially supported, pick one of the architecture below and let's discuss about the conversion! 

Text models 🖊️  :

- [x] FSMT - [FSMTEncoderLayer](https://github.com/huggingface/transformers/blob/95754b47a6d4fbdad3440a45762531e8c471c528/src/transformers/models/fsmt/modeling_fsmt.py#L397) / @Sumanth077 https://github.com/huggingface/optimum/pull/494
- [ ] MobileBERT - [MobileBertLayer](https://github.com/huggingface/transformers/blob/95754b47a6d4fbdad3440a45762531e8c471c528/src/transformers/models/mobilebert/modeling_mobilebert.py#L498) / @raghavanone https://github.com/huggingface/optimum/pull/506
- [x] MBart - [MBartEncoderLayer](https://github.com/huggingface/transformers/blob/95754b47a6d4fbdad3440a45762531e8c471c528/src/transformers/models/mbart/modeling_mbart.py#L296) + [M2M100EncoderLayer](https://github.com/huggingface/transformers/blob/95754b47a6d4fbdad3440a45762531e8c471c528/src/transformers/models/m2m_100/modeling_m2m_100.py#L345) / https://github.com/huggingface/optimum/pull/516 @ravenouse 
- [x] ProphetNet - [ProphetNetEncoderLayer](https://github.com/huggingface/transformers/blob/95754b47a6d4fbdad3440a45762531e8c471c528/src/transformers/models/prophetnet/modeling_prophetnet.py#L1130)
- [x] RemBert - [RemBertLayer](https://github.com/huggingface/transformers/blob/95754b47a6d4fbdad3440a45762531e8c471c528/src/transformers/models/rembert/modeling_rembert.py#L415)
- [x] RocBert - [RocBertLayer](https://github.com/huggingface/transformers/blob/95754b47a6d4fbdad3440a45762531e8c471c528/src/transformers/models/roc_bert/modeling_roc_bert.py#LL519C7-L519C19)
- [x] RoFormer - [RoFormerLayer](https://github.com/huggingface/transformers/blob/95754b47a6d4fbdad3440a45762531e8c471c528/src/transformers/models/roformer/modeling_roformer.py#L448)
- [x] Tapas - [TapasLayer](https://github.com/huggingface/transformers/blob/95754b47a6d4fbdad3440a45762531e8c471c528/src/transformers/models/tapas/modeling_tapas.py#L524) / https://github.com/huggingface/optimum/pull/520


Vision models 📷  :

- [x] Blip - [BlipLayer](https://github.com/huggingface/transformers/blob/fcf813417aa34f3a0ea7d283f7d4f6b0834cf098/src/transformers/models/blip/modeling_blip.py#L372)
- [ ] Detr - [DetrLayer](https://github.com/huggingface/transformers/blob/95754b47a6d4fbdad3440a45762531e8c471c528/src/transformers/models/detr/modeling_detr.py#L610)
- [ ] Flava - [FlavaLayer](https://github.com/huggingface/transformers/blob/95754b47a6d4fbdad3440a45762531e8c471c528/src/transformers/models/flava/modeling_flava.py#L597)
- [ ] GLPN - [GLPNLayer](https://github.com/huggingface/transformers/blob/95754b47a6d4fbdad3440a45762531e8c471c528/src/transformers/models/glpn/modeling_glpn.py#L292) | Cannot be supported
- [x] ViLT - [ViLTLayer](https://github.com/huggingface/transformers/blob/95754b47a6d4fbdad3440a45762531e8c471c528/src/transformers/models/vilt/modeling_vilt.py#L472) / https://github.com/huggingface/optimum/pull/508

Audio models 🔉  :

- [ ] Speech2Text - [Speech2TextLayer](https://github.com/huggingface/transformers/blob/95754b47a6d4fbdad3440a45762531e8c471c528/src/transformers/models/speech_to_text/modeling_speech_to_text.py#L350)
- [ ] NEW: Audio Speech Transformer - [ASTLayer](https://github.com/huggingface/transformers/blob/f2e7d270ec795be09e6187dd2459edb43bd861c1/src/transformers/models/audio_spectrogram_transformer/modeling_audio_spectrogram_transformer.py#L274)

Let us also know if you think that some architectures can be supported that we missed. Note that for encoder-decoder based models below, we expect to convert the encoder only.

**Support for decoder-based models coming soon!**

cc @michaelbenayoun @fxmarty 

https://github.com/huggingface/transformers/issues/20372


Discussion:
Hi @younesbelkada would love to contribute to this Issue and can work on FSMT. 
Hey @Sumanth077 , thanks a bunch for your interest in this issue! 🚀  Would love to assist you for the integration and let's try to make this happen! 
I have updated the table above, and attaching you the contribution tutorial [here](https://huggingface.co/docs/optimum/bettertransformer/tutorials/contribute) ;) 
Would you mind forking this repo and start opening a draft pull request so that I can start guiding you there? 
Also please do not hesitate to ping us here for any issue you are facing for the integration 💪 
Thankyou for the reply @younesbelkada. Just opened a Draft Pull Request, haven't made any significant changes. 

In the Step 1: Identifying the source layer to change and in the BETTER_TRANFORMER_LAYERS_MAPPING_DICT, I couldn't find a mapping between the Module for the FSMT that can be converted to its BetterTransformer equivalent.

Should I start creating that. Would love your assistance
Hi @Sumanth077 , I have just replied on your PR, let's continue the discussion there ;) 
Hi, I would like to contribute as well. This would be my first contribution to open source, so I might need some hand holding 🤚

I followed [the documentation](https://huggingface.co/docs/optimum/bettertransformer/tutorials/contribute) and the progress made on FSMT in huggingface/optimum#494 to better understand the task.

I looked into ViLT via 

> `model = AutoModel.from_pretrained("dandelin/vilt-b32-mlm")` 

and as I understand the documentation, this should be the source layer to make changes to, including its attributes: 
> `      (0): ViltLayer(
        (attention): ViltAttention(
          (attention): ViltSelfAttention(
            (query): Linear(in_features=768, out_features=768, bias=True)
            (key): Linear(in_features=768, out_features=768, bias=True)
            (value): Linear(in_features=768, out_features=768, bias=True)
            (dropout): Dropout(p=0.0, inplace=False)
          )
          (output): ViltSelfOutput(
            (dense): Linear(in_features=768, out_features=768, bias=True)
            (dropout): Dropout(p=0.0, inplace=False)
          )
        )
        (intermediate): ViltIntermediate(
          (dense): Linear(in_features=768, out_features=3072, bias=True)
          (intermediate_act_fn): GELUActivation()
        )
        (output): ViltOutput(
          (dense): Linear(in_features=3072, out_features=768, bias=True)
          (dropout): Dropout(p=0.0, inplace=False)
        )
        (layernorm_before): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
        (layernorm_after): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
      )`

I could give the ViLTLayer a go, if it's ok with you @younesbelkada 🙂


Hi @ka00ri !
Thanks a lot for your message and interest in contributing! Would love to assist you for integrating `ViLT` into `BetterTransformer` 💪 
That is correct, this layer has to be the source layer to change! 
Would you mind opening a PR and tag us (myself, @michaelbenayoun & @fxmarty ) ? Thanks a bunch!
Hello, apologies for the delay, but I just opened up a draft PR to start discussion on how to add Better Transformer support for the ProphetNet encoder layer. I had a couple of questions about how to do this, so I was wondering who would would be the best person to reach out to regarding this. @michaelbenayoun @fxmarty @younesbelkada 
Hi @adit299 , thanks for adding the support for this architecture! Feel free to ask any question in the PR you opened.
Hi @younesbelkada, could I pick up the RoFormer?

@younesbelkada  doing Detr - [DetrLayer](https://github.com/huggingface/transformers/blob/95754b47a6d4fbdad3440a45762531e8c471c528/src/transformers/models/detr/modeling_detr.py#L610)
Hello @JanFidor 
Yes sure! 
@soma2000-lang perfect, let us know when you open a PR 💪 !
@younesbelkada Hi, thanks for responding, I'm not 100% certain, but I think RemBert, RoFormer and RocBert are already implemented, as they're already added to __init__.py, overview.mdx and the test_file, if that's the case, the list of models left to implement would need to be updated, let me know if you agree!
I see, thanks for clarifying. I will double check that and let you know
Thanks for letting me know! Indeed these are already implemented
I can propose you to add BetterTransformer support for `Blip` (updated the table above)
Thanks for the suggestion, I'll get on it!

Resolving patch:
diff --git a/optimum/bettertransformer/models/__init__.py b/optimum/bettertransformer/models/__init__.py
--- a/optimum/bettertransformer/models/__init__.py
+++ b/optimum/bettertransformer/models/__init__.py
@@ -29,6 +29,7 @@
     DistilBertLayerBetterTransformer,
     FSMTEncoderLayerBetterTransformer,
     MBartEncoderLayerBetterTransformer,
+    ProphetNetEncoderLayerBetterTransformer,
     ViltLayerBetterTransformer,
     ViTLayerBetterTransformer,
     Wav2Vec2EncoderLayerBetterTransformer,
@@ -74,6 +75,7 @@ class BetterTransformerManager:
         "opt": {"OPTAttention": OPTAttentionLayerBetterTransformer},
         "pegasus": {"PegasusAttention": BartAttentionLayerBetterTransformer},
         "rembert": {"RemBertLayer": BertLayerBetterTransformer},
+        "prophetnet": {"ProphetNetEncoderLayer": ProphetNetEncoderLayerBetterTransformer},
         "roberta": {"RobertaLayer": BertLayerBetterTransformer},
         "roc_bert": {"RoCBertLayer": BertLayerBetterTransformer},
         "roformer": {"RoFormerLayer": BertLayerBetterTransformer},
diff --git a/optimum/bettertransformer/models/encoder_models.py b/optimum/bettertransformer/models/encoder_models.py
--- a/optimum/bettertransformer/models/encoder_models.py
+++ b/optimum/bettertransformer/models/encoder_models.py
@@ -1275,6 +1275,139 @@ def forward(self, hidden_states, attention_mask, position_bias=None, *_, **__):
         return (hidden_states, attention_mask)
 
 
+class ProphetNetEncoderLayerBetterTransformer(BetterTransformerBaseLayer):
+    def __init__(self, prophetnet_layer, config):
+        r"""
+        A simple conversion of the ProphetNet Encoder layer to its `BetterTransformer` implementation.
+
+        Args:
+            prophet_net_layer (`torch.nn.Module`):
+                The original ProphetNet Layer where the weights needs to be retrieved.
+        """
+        super().__init__(config)
+        self.config = config
+        # In_proj layer
+        self.in_proj_weight = nn.Parameter(
+            torch.cat(
+                [
+                    prophetnet_layer.self_attn.query_proj.weight,
+                    prophetnet_layer.self_attn.key_proj.weight,
+                    prophetnet_layer.self_attn.value_proj.weight,
+                ]
+            )
+        )
+        self.in_proj_bias = nn.Parameter(
+            torch.cat(
+                [
+                    prophetnet_layer.self_attn.query_proj.bias,
+                    prophetnet_layer.self_attn.key_proj.bias,
+                    prophetnet_layer.self_attn.value_proj.bias,
+                ]
+            )
+        )
+
+        # Out proj layer
+        self.out_proj_weight = prophetnet_layer.self_attn.out_proj.weight
+        self.out_proj_bias = prophetnet_layer.self_attn.out_proj.bias
+
+        # Linear layer 1
+        self.linear1_weight = prophetnet_layer.feed_forward.intermediate.weight
+        self.linear1_bias = prophetnet_layer.feed_forward.intermediate.bias
+
+        # Linear layer 2
+        self.linear2_weight = prophetnet_layer.feed_forward.output.weight
+        self.linear2_bias = prophetnet_layer.feed_forward.output.bias
+
+        # Layer norm 1
+        self.norm1_eps = prophetnet_layer.self_attn_layer_norm.eps
+        self.norm1_weight = prophetnet_layer.self_attn_layer_norm.weight
+        self.norm1_bias = prophetnet_layer.self_attn_layer_norm.bias
+
+        # Layer norm 2
+        self.norm2_eps = prophetnet_layer.feed_forward_layer_norm.eps
+        self.norm2_weight = prophetnet_layer.feed_forward_layer_norm.weight
+        self.norm2_bias = prophetnet_layer.feed_forward_layer_norm.bias
+
+        # Model hyper parameters
+        self.num_heads = prophetnet_layer.self_attn.num_attn_heads
+        self.embed_dim = prophetnet_layer.self_attn.head_dim * 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.original_layers_mapping = {
+            "in_proj_weight": [
+                "self_attn.query_proj.weight",
+                "self_attn.key_proj.weight",
+                "self_attn.value_proj.weight",
+            ],
+            "in_proj_bias": ["self_attn.query_proj.bias", "self_attn.key_proj.bias", "self_attn.value_proj.bias"],
+            "out_proj_weight": "self_attn.out_proj.weight",
+            "out_proj_bias": "self_attn.out_proj.bias",
+            "linear1_weight": "feed_forward.intermediate.weight",
+            "linear1_bias": "feed_forward.intermediate.bias",
+            "linear2_weight": "feed_forward.output.weight",
+            "linear2_bias": "feed_forward.output.bias",
+            "norm1_weight": "self_attn_layer_norm.weight",
+            "norm1_bias": "self_attn_layer_norm.bias",
+            "norm2_weight": "feed_forward_layer_norm.weight",
+            "norm2_bias": "feed_forward_layer_norm.bias",
+        }
+
+        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()
+
+        if not hasattr(hidden_states, "original_shape"):
+            original_shape = hidden_states.shape
+        else:
+            original_shape = hidden_states.original_shape
+
+        if hidden_states.is_nested:
+            attention_mask = None
+
+        if attention_mask is not None:
+            # attention mask comes in with values 0 and -inf. we convert to torch.nn.TransformerEncoder style bool mask
+            # 0->false->keep this token -inf->true->mask this token
+            attention_mask = attention_mask.squeeze(1)[:, 0]
+            attention_mask = attention_mask.bool()
+            attention_mask = torch.reshape(attention_mask, (attention_mask.shape[0], attention_mask.shape[-1]))
+            hidden_states = torch._nested_tensor_from_mask(hidden_states, ~attention_mask)
+            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 not self.is_last_layer:
+            hidden_states.original_shape = original_shape
+        elif hidden_states.is_nested and self.is_last_layer:
+            hidden_states = hidden_states.to_padded_tensor(0.0, original_shape)
+        return (hidden_states,)
+
+
 class CLIPLayerBetterTransformer(BetterTransformerBaseLayer):
     def __init__(self, layer, config):
         r"""


Test patch:
diff --git a/optimum/utils/testing_utils.py b/optimum/utils/testing_utils.py
--- a/optimum/utils/testing_utils.py
+++ b/optimum/utils/testing_utils.py
@@ -13,13 +13,13 @@
 # limitations under the License.
 
 
-import collections
 import importlib.util
 import itertools
 import os
 import subprocess
 import sys
 import unittest
+from collections.abc import MutableMapping
 from typing import Any, Callable, Dict, Iterable, Optional, Tuple
 
 import torch
@@ -35,7 +35,7 @@ def flatten_dict(dictionary: Dict):
     items = []
     for k, v in dictionary.items():
         new_key = k
-        if isinstance(v, collections.MutableMapping):
+        if isinstance(v, MutableMapping):
             items.extend(flatten_dict(v).items())
         else:
             items.append((new_key, v))
diff --git a/tests/bettertransformer/test_encoder_decoder.py b/tests/bettertransformer/test_encoder_decoder.py
--- a/tests/bettertransformer/test_encoder_decoder.py
+++ b/tests/bettertransformer/test_encoder_decoder.py
@@ -43,6 +43,7 @@ class BetterTransformersEncoderDecoderTest(BetterTransformersTestMixin, unittest
         "marian",
         "mbart",
         "pegasus",
+        "prophetnet",
         "t5",
     ]
 
diff --git a/tests/bettertransformer/testing_utils.py b/tests/bettertransformer/testing_utils.py
--- a/tests/bettertransformer/testing_utils.py
+++ b/tests/bettertransformer/testing_utils.py
@@ -54,6 +54,7 @@
     "mbart": "hf-internal-testing/tiny-random-mbart",
     "opt": "hf-internal-testing/tiny-random-OPTModel",
     "pegasus": "hf-internal-testing/tiny-random-PegasusModel",
+    "prophetnet": "hirotasoshu/tiny-random-prophetnet",  # the other tiny ones have a too small max_position_embeddings
     "rembert": "hf-internal-testing/tiny-random-rembert",
     "roberta": "hf-internal-testing/tiny-random-RobertaModel",
     "rocbert": "hf-internal-testing/tiny-random-RoCBertModel",