Add BetterTransformer support for FlavaModel

[katiele47]

What does this PR do?

Fixes #20372

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?
• [Other] Applied make style

Questions

• Should I make this PR (and future PRs that aren't ready for review but are worth discussing) into a draft PR? There are some questions I'd like to confirm with you in advance.
• Would it be necessary to add the BetterTransformersFlavaTest or should I reuse BetterTransformersViltTest for FlavaModel? If so, should I update the name to BetterTransformersVisionTextTest?
• I ran into the following issues when testing the conversion of FlavaModel with both pytest and running this script

from transformers import AutoModel
from optimum.bettertransformer import BetterTransformer

model_id = "hf-internal-testing/tiny-random-FlavaModel"
model = AutoModel.from_pretrained(model_id)
bt_model = BetterTransformer.transform(model)

Error:

 ~/optimum/optimum/bettertransformer/models/base.py:81 in validate_bettertransformer    │
│                                                                                                  │
│    78 │   │                                                                                      │
│    79 │   │   # Check activation function                                                        │
│    80 │   │   if self.act_fn not in SUPPORTED_ACTIVATION_FUNCTIONS:                              │
│ ❱  81 │   │   │   raise ValueError(                                                              │
│    82 │   │   │   │   f"Activation function {self.act_fn} not supported" " for `BetterTransfor   │
│    83 │   │   │   )                                                                              │
│    84 │   │   self.use_gelu = (self.act_fn == "gelu") or (self.act_fn == "gelu_new")             │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
ValueError: Activation function None not supported for `BetterTransformer` integration.

I couldn't find where the self.act_fn get defined inside the forward method after following this guide. Would appreciate any help!

[younesbelkada]

Hi @katiele47 !
Thanks so much for your very clean and great PR! Everything looks perfect 🔥
Regarding FLAVA, I quickly dived into it and the fix seems to be relatively simple. The model has 2 components, image and text encoders. Each of these submodules has its own config file that you can access with .image_config or .text_config. I think that the fix that I propose here is fine, as for both sub-configs the same hyper parameters seems to be used (i.e. same hidden_size, same hidden_act, etc.). Check for example this config file

After accepting my suggestion, you can run pytest tests/bettertransformer/test_bettertransformer_vision.py:: BetterTransformersFlavaTest to make sure the test pass ;)

[katiele47]

Thanks a lot for your insight @younesbelkada! I will run the test and update this PR

[katiele47]

The test script now worked with the following output, but I encountered 2 failed tests from BetterTransformersViltTest and BetterTransformerFlavaTest. Do you think I should use a different resized image here? It seems weird how the BetterTransformersViltTest previously passed with the current input image.

FAILED tests/bettertransformer/test_bettertransformer_vision.py::BetterTransformersViLTTest::test_logits - ValueError: Input image size (224*224) doesn't match model (30*30).
FAILED tests/bettertransformer/test_bettertransformer_vision.py::BetterTransformersFlavaTest::test_logits - ValueError: Input image size (224*224) doesn't match model (30*30).

Test script output:

FlavaModel(
  (text_model): FlavaTextModel(
    (embeddings): FlavaTextEmbeddings(
      (word_embeddings): Embedding(1124, 32, padding_idx=0)
      (position_embeddings): Embedding(512, 32)
      (token_type_embeddings): Embedding(2, 32)
      (LayerNorm): LayerNorm((32,), eps=1e-12, elementwise_affine=True)
      (dropout): Dropout(p=0.0, inplace=False)
    )
    (encoder): FlavaEncoder(
      (layer): ModuleList(
        (0): FlavaLayerBetterTransformer()
        (1): FlavaLayerBetterTransformer()
        (2): FlavaLayerBetterTransformer()
        (3): FlavaLayerBetterTransformer()
        (4): FlavaLayerBetterTransformer()
      )
    )
    (layernorm): LayerNorm((32,), eps=1e-12, elementwise_affine=True)
    (pooler): FlavaPooler(
      (dense): Linear(in_features=32, out_features=32, bias=True)
      (activation): Tanh()
    )
  )
  (image_model): FlavaImageModel(
    (embeddings): FlavaImageEmbeddings(
      (patch_embeddings): PatchEmbeddings(
        (projection): Conv2d(3, 32, kernel_size=(2, 2), stride=(2, 2))
      )
      (dropout): Dropout(p=0.0, inplace=False)
    )
    (encoder): FlavaEncoder(
      (layer): ModuleList(
        (0): FlavaLayerBetterTransformer()
        (1): FlavaLayerBetterTransformer()
        (2): FlavaLayerBetterTransformer()
        (3): FlavaLayerBetterTransformer()
        (4): FlavaLayerBetterTransformer()
      )
    )
    (layernorm): LayerNorm((32,), eps=1e-12, elementwise_affine=True)
    (pooler): FlavaPooler(
      (dense): Linear(in_features=32, out_features=32, bias=True)
      (activation): Tanh()
    )
  )
  (multimodal_model): FlavaMultimodalModel(
    (encoder): FlavaEncoder(
      (layer): ModuleList(
        (0): FlavaLayerBetterTransformer()
        (1): FlavaLayerBetterTransformer()
        (2): FlavaLayerBetterTransformer()
        (3): FlavaLayerBetterTransformer()
        (4): FlavaLayerBetterTransformer()
      )
    )
    (layernorm): LayerNorm((32,), eps=1e-12, elementwise_affine=True)
    (pooler): FlavaPooler(
      (dense): Linear(in_features=32, out_features=32, bias=True)
      (activation): Tanh()
    )
  )
  (image_projection): Linear(in_features=32, out_features=32, bias=True)
  (text_projection): Linear(in_features=32, out_features=32, bias=True)
  (image_to_mm_projection): Linear(in_features=32, out_features=32, bias=True)
  (text_to_mm_projection): Linear(in_features=32, out_features=32, bias=True)
)

[younesbelkada]

Thanks so much for digging into the issue!
I think that there is a problem with the feature extractor in hf-internal-testing/tiny-random-FlavaModel , can you please try with a new model that I created? (Suggestion below) Thanks!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[katiele47]

Thanks again for your time and suggestion @younesbelkada! I have made the changes accordingly, and all tests passed. I'll keep an eye on any arising CircleCI issues in the meantime.

[younesbelkada]

Hi @katiele47 !
I left a couple of comments, this should fix the failing tests! ;)
Thanks again

[katiele47]

@younesbelkada Sorry for the multiple file changes, I think it's due to merging conflicts with upstream main, but I'm not entirely sure. Please let me know if this affect the original flava changes.

[katiele47]

Hi @younesbelkada I have created a new PR here that replaces this outdated PR. Please review and let me know if I should close this one. Thanks a lot!