encoder_models.py

# Copyright 2022 The HuggingFace and Meta Team.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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
import torch.nn.functional as F
from transformers.activations import ACT2FN

from .base import BetterTransformerBaseLayer


if TYPE_CHECKING:
    from transformers import PretrainedConfig


class AlbertLayerBetterTransformer(BetterTransformerBaseLayer, nn.Module):
    def __init__(self, albert_layer, config):
        r"""
        A simple conversion of the ALBERT layer to its `BetterTransformer` implementation.

        Args:
            albert_layer (`torch.nn.Module`):
                The original ALBERT Layer where the weights needs to be retrieved.
        """
        super().__init__(config)
        super(BetterTransformerBaseLayer, self).__init__()
        # In_proj layer
        self.in_proj_weight = nn.Parameter(
            torch.cat(
                [
                    albert_layer.attention.query.weight,
                    albert_layer.attention.key.weight,
                    albert_layer.attention.value.weight,
                ]
            )
        )
        self.in_proj_bias = nn.Parameter(
            torch.cat(
                [
                    albert_layer.attention.query.bias,
                    albert_layer.attention.key.bias,
                    albert_layer.attention.value.bias,
                ]
            )
        )

        # Out proj layer
        self.out_proj_weight = albert_layer.attention.dense.weight
        self.out_proj_bias = albert_layer.attention.dense.bias

        # Linear layer 1
        self.linear1_weight = albert_layer.ffn.weight
        self.linear1_bias = albert_layer.ffn.bias

        # Linear layer 2
        self.linear2_weight = albert_layer.ffn_output.weight
        self.linear2_bias = albert_layer.ffn_output.bias

        # Layer norm 1
        self.norm1_eps = albert_layer.attention.LayerNorm.eps
        self.norm1_weight = albert_layer.attention.LayerNorm.weight
        self.norm1_bias = albert_layer.attention.LayerNorm.bias

        # Layer norm 2
        self.norm2_eps = albert_layer.full_layer_layer_norm.eps
        self.norm2_weight = albert_layer.full_layer_layer_norm.weight
        self.norm2_bias = albert_layer.full_layer_layer_norm.bias

        # Model hyper parameters
        self.num_heads = albert_layer.attention.num_attention_heads
        self.embed_dim = albert_layer.attention.all_head_size

        # 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": ["attention.query.weight", "attention.key.weight", "attention.value.weight"],
            "in_proj_bias": ["attention.query.bias", "attention.key.bias", "attention.value.bias"],
            "out_proj_weight": "attention.dense.weight",
            "out_proj_bias": "attention.dense.bias",
            "linear1_weight": "ffn.weight",
            "linear1_bias": "ffn.bias",
            "linear2_weight": "ffn_output.weight",
            "linear2_bias": "ffn_output.bias",
            "norm1_eps": "attention.LayerNorm.eps",
            "norm1_weight": "attention.LayerNorm.weight",
            "norm1_bias": "attention.LayerNorm.bias",
            "norm2_eps": "full_layer_layer_norm.eps",
            "norm2_weight": "full_layer_layer_norm.weight",
            "norm2_bias": "full_layer_layer_norm.bias",
        }
        self.attention_head_size = config.hidden_size // config.num_attention_heads
        self.attention_probs_dropout_prob = config.attention_probs_dropout_prob
        self.hidden_dropout_prob = config.hidden_dropout_prob
        self.act_fn_callable = ACT2FN[self.act_fn]

        self.validate_bettertransformer()

    def forward(self, hidden_states, attention_mask, *_):
        if not self.training and not torch.is_autocast_enabled() and not torch.is_autocast_cpu_enabled():
            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.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 hidden_states.is_nested and self.is_last_layer:
                hidden_states = hidden_states.to_padded_tensor(0.0)
        else:
            qkv = F.linear(hidden_states, weight=self.in_proj_weight, bias=self.in_proj_bias)

            qkv = qkv.view(qkv.size()[:-1] + (3, self.num_heads, self.attention_head_size)).permute(2, 0, 3, 1, 4)
            query, key, value = qkv[0], qkv[1], qkv[2]

            # NOTE: In PyTorch 2.0, passing an attention_mask will automatically dispatch
            # to the "math" path and will NOT use flash attention / memory-efficient attention.
            # We should support xformers / Hazy-flash / rocm-flash directly and stop relying on PyTorch to do the work.
            if self.training:
                attention_mask = None
            attention_out = F.scaled_dot_product_attention(
                query,
                key,
                value,
                attn_mask=attention_mask,
                is_causal=False,
                dropout_p=self.attention_probs_dropout_prob if self.training else 0.0,
            )

            attention_out = attention_out.permute(0, 2, 1, 3).contiguous()
            new_attention_out_shape = attention_out.size()[:-2] + (self.num_heads * self.attention_head_size,)
            attention_out = attention_out.view(new_attention_out_shape)

            # BertSelfOutput
            attention_out = F.layer_norm(
                F.dropout(
                    F.linear(attention_out, self.out_proj_weight, self.out_proj_bias),
                    p=self.hidden_dropout_prob,
                    training=self.training,
                )
                + hidden_states,
                normalized_shape=self.norm1_weight.shape,
                weight=self.norm1_weight,
                bias=self.norm1_bias,
            )

            # BertIntermediate
            hidden_states = self.act_fn_callable(F.linear(attention_out, self.linear1_weight, self.linear1_bias))

            # BertOutput
            hidden_states = F.layer_norm(
                attention_out
                + F.dropout(
                    F.linear(hidden_states, self.linear2_weight, self.linear2_bias),
                    p=self.hidden_dropout_prob,
                    training=self.training,
                ),
                normalized_shape=self.norm2_weight.shape,
                weight=self.norm2_weight,
                bias=self.norm2_bias,
            )

        return (hidden_states,)


class BertLayerBetterTransformer(BetterTransformerBaseLayer, nn.Module):
    def __init__(self, bert_layer, config):
        r"""
        A simple conversion of the BERT layer to its `BetterTransformer` implementation.

        Args:
            bert_layer (`torch.nn.Module`):
                The original BERT Layer where the weights needs to be retrieved.
        """
        super().__init__(config)
        super(BetterTransformerBaseLayer, self).__init__()
        # In_proj layer
        self.in_proj_weight = nn.Parameter(
            torch.cat(
                [
                    bert_layer.attention.self.query.weight,
                    bert_layer.attention.self.key.weight,
                    bert_layer.attention.self.value.weight,
                ]
            )
        )
        self.in_proj_bias = nn.Parameter(
            torch.cat(
                [
                    bert_layer.attention.self.query.bias,
                    bert_layer.attention.self.key.bias,
                    bert_layer.attention.self.value.bias,
                ]
            )
        )

        # Out proj layer
        self.out_proj_weight = bert_layer.attention.output.dense.weight
        self.out_proj_bias = bert_layer.attention.output.dense.bias

        # Linear layer 1
        self.linear1_weight = bert_layer.intermediate.dense.weight
        self.linear1_bias = bert_layer.intermediate.dense.bias

        # Linear layer 2
        self.linear2_weight = bert_layer.output.dense.weight
        self.linear2_bias = bert_layer.output.dense.bias

        # Layer norm 1
        self.norm1_eps = bert_layer.attention.output.LayerNorm.eps
        self.norm1_weight = bert_layer.attention.output.LayerNorm.weight
        self.norm1_bias = bert_layer.attention.output.LayerNorm.bias

        # Layer norm 2
        self.norm2_eps = bert_layer.output.LayerNorm.eps
        self.norm2_weight = bert_layer.output.LayerNorm.weight
        self.norm2_bias = bert_layer.output.LayerNorm.bias

        # Model hyper parameters
        self.num_heads = bert_layer.attention.self.num_attention_heads
        self.embed_dim = bert_layer.attention.self.all_head_size

        # 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": [
                "attention.self.query.weight",
                "attention.self.key.weight",
                "attention.self.value.weight",
            ],
            "in_proj_bias": ["attention.self.query.bias", "attention.self.key.bias", "attention.self.value.bias"],
            "out_proj_weight": "attention.output.dense.weight",
            "out_proj_bias": "attention.output.dense.bias",
            "linear1_weight": "intermediate.dense.weight",
            "linear1_bias": "intermediate.dense.bias",
            "linear2_weight": "output.dense.weight",
            "linear2_bias": "output.dense.bias",
            "norm1_eps": "attention.output.LayerNorm.eps",
            "norm1_weight": "attention.output.LayerNorm.weight",
            "norm1_bias": "attention.output.LayerNorm.bias",
            "norm2_eps": "output.LayerNorm.eps",
            "norm2_weight": "output.LayerNorm.weight",
            "norm2_bias": "output.LayerNorm.bias",
        }
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.hidden_dropout_prob = config.hidden_dropout_prob
        self.attention_probs_dropout_prob = config.attention_probs_dropout_prob
        self.act_fn_callable = ACT2FN[self.act_fn]

        self.validate_bettertransformer()

    def forward(self, hidden_states, attention_mask, *_):
        # No check on output_attentions here as roformer relies on BertLayerBetterTransformer but does not pass output_attentions as keyword argument.
        if not self.training and not torch._C._is_any_autocast_enabled():
            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.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 hidden_states.is_nested and self.is_last_layer:
                hidden_states = hidden_states.to_padded_tensor(0.0)
        else:
            qkv = F.linear(hidden_states, weight=self.in_proj_weight, bias=self.in_proj_bias)

            qkv = qkv.view(qkv.size()[:-1] + (3, self.num_heads, self.attention_head_size)).permute(2, 0, 3, 1, 4)
            query, key, value = qkv[0], qkv[1], qkv[2]

            # NOTE: In PyTorch 2.0, passing an attention_mask will automatically dispatch
            # to the "math" path and will NOT use flash attention / memory-efficient attention.
            # We should support xformers / Hazy-flash / rocm-flash directly and stop relying on PyTorch to do the work.
            if self.training:
                attention_mask = None
            attention_out = F.scaled_dot_product_attention(
                query,
                key,
                value,
                attn_mask=attention_mask,
                is_causal=False,
                dropout_p=self.attention_probs_dropout_prob if self.training else 0.0,
            )

            attention_out = attention_out.permute(0, 2, 1, 3).contiguous()
            new_attention_out_shape = attention_out.size()[:-2] + (self.num_heads * self.attention_head_size,)
            attention_out = attention_out.view(new_attention_out_shape)

            # BertSelfOutput
            attention_out = F.layer_norm(
                F.dropout(
                    F.linear(attention_out, self.out_proj_weight, self.out_proj_bias),
                    p=self.hidden_dropout_prob,
                    training=self.training,
                )
                + hidden_states,
                normalized_shape=self.norm1_weight.shape,
                weight=self.norm1_weight,
                bias=self.norm1_bias,
            )

            # BertIntermediate
            hidden_states = self.act_fn_callable(F.linear(attention_out, self.linear1_weight, self.linear1_bias))

            # BertOutput
            hidden_states = F.layer_norm(
                attention_out
                + F.dropout(
                    F.linear(hidden_states, self.linear2_weight, self.linear2_bias),
                    p=self.hidden_dropout_prob,
                    training=self.training,
                ),
                normalized_shape=self.norm2_weight.shape,
                weight=self.norm2_weight,
                bias=self.norm2_bias,
            )

        return (hidden_states,)


class BartEncoderLayerBetterTransformer(BetterTransformerBaseLayer, nn.Module):
    def __init__(self, bart_layer, config):
        r"""
        A simple conversion of the `BartEncoderLayer` to its `BetterTransformer` implementation.

        Args:
            bart_layer (`torch.nn.Module`):
                The original `BartEncoderLayer` where the weights needs to be retrieved.
        """
        super().__init__(config)
        super(BetterTransformerBaseLayer, self).__init__()
        # In_proj layer
        self.in_proj_weight = nn.Parameter(
            torch.cat(
                [
                    bart_layer.self_attn.q_proj.weight,
                    bart_layer.self_attn.k_proj.weight,
                    bart_layer.self_attn.v_proj.weight,
                ]
            )
        )

        self.in_proj_bias = nn.Parameter(
            torch.cat(
                [
                    bart_layer.self_attn.q_proj.bias,
                    bart_layer.self_attn.k_proj.bias,
                    bart_layer.self_attn.v_proj.bias,
                ]
            )
        )

        # Out proj layer
        self.out_proj_weight = bart_layer.self_attn.out_proj.weight
        self.out_proj_bias = bart_layer.self_attn.out_proj.bias

        # Linear layer 1
        self.linear1_weight = bart_layer.fc1.weight
        self.linear1_bias = bart_layer.fc1.bias

        # Linear layer 2
        self.linear2_weight = bart_layer.fc2.weight
        self.linear2_bias = bart_layer.fc2.bias

        # Layer norm 1
        self.norm1_eps = bart_layer.self_attn_layer_norm.eps
        self.norm1_weight = bart_layer.self_attn_layer_norm.weight
        self.norm1_bias = bart_layer.self_attn_layer_norm.bias

        # Layer norm 2
        self.norm2_eps = bart_layer.final_layer_norm.eps
        self.norm2_weight = bart_layer.final_layer_norm.weight
        self.norm2_bias = bart_layer.final_layer_norm.bias

        # Model hyper parameters
        self.num_heads = bart_layer.self_attn.num_heads
        self.embed_dim = bart_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.original_layers_mapping = {
            "in_proj_weight": ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"],
            "in_proj_bias": ["self_attn.q_proj.bias", "self_attn.k_proj.bias", "self_attn.v_proj.bias"],
            "out_proj_weight": "self_attn.out_proj.weight",
            "out_proj_bias": "self_attn.out_proj.bias",
            "linear1_weight": "fc1.weight",
            "linear1_bias": "fc1.bias",
            "linear2_weight": "fc2.weight",
            "linear2_bias": "fc2.bias",
            "norm1_eps": "self_attn_layer_norm.eps",
            "norm1_weight": "self_attn_layer_norm.weight",
            "norm1_bias": "self_attn_layer_norm.bias",
            "norm2_eps": "final_layer_norm.eps",
            "norm2_weight": "final_layer_norm.weight",
            "norm2_bias": "final_layer_norm.bias",
        }
        self.dropout = config.attention_dropout
        self.activation_dropout = config.activation_dropout
        self.attention_head_size = config.d_model // config.encoder_attention_heads
        self.act_fn_callable = ACT2FN[self.act_fn]

        self.validate_bettertransformer()

    def forward(self, hidden_states, attention_mask, output_attentions: bool, position_bias=None, *_, **__):
        if output_attentions:
            raise ValueError("output_attentions=True can not be supported with BetterTransformer.")

        if not self.training and not torch.is_autocast_enabled() and not torch.is_autocast_cpu_enabled():
            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
                if len(attention_mask.shape) == 4:
                    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)
        else:
            qkv = F.linear(hidden_states, weight=self.in_proj_weight, bias=self.in_proj_bias)

            qkv = qkv.view(qkv.size()[:-1] + (3, self.num_heads, self.attention_head_size)).permute(2, 0, 3, 1, 4)
            query, key, value = qkv[0], qkv[1], qkv[2]

            # NOTE: In PyTorch 2.0, passing an attention_mask will automatically dispatch
            # to the "math" path and will NOT use flash attention / memory-efficient attention.
            # We should support xformers / Hazy-flash / rocm-flash directly and stop relying on PyTorch to do the work.
            if self.training:
                attention_mask = None
            attention_out = F.scaled_dot_product_attention(
                query,
                key,
                value,
                attn_mask=attention_mask,
                is_causal=False,
                dropout_p=self.dropout if self.training else 0.0,
            )

            attention_out = attention_out.permute(0, 2, 1, 3).contiguous()
            new_attention_out_shape = attention_out.size()[:-2] + (self.num_heads * self.attention_head_size,)
            attention_out = attention_out.view(new_attention_out_shape)

            # BertSelfOutput
            attention_out = F.layer_norm(
                F.dropout(
                    F.linear(attention_out, self.out_proj_weight, self.out_proj_bias),
                    p=self.dropout,
                    training=self.training,
                )
                + hidden_states,
                normalized_shape=self.norm1_weight.shape,
                weight=self.norm1_weight,
                bias=self.norm1_bias,
            )

            # One additional dropout compared to bert
            hidden_states = F.dropout(
                self.act_fn_callable(F.linear(attention_out, self.linear1_weight, self.linear1_bias)),
                p=self.activation_dropout,
                training=self.training,
            )

            hidden_states = F.layer_norm(
                attention_out
                + F.dropout(
                    F.linear(hidden_states, self.linear2_weight, self.linear2_bias),
                    p=self.dropout,
                    training=self.training,
                ),
                normalized_shape=self.norm2_weight.shape,
                weight=self.norm2_weight,
                bias=self.norm2_bias,
            )
        return (hidden_states,)


class MBartEncoderLayerBetterTransformer(BetterTransformerBaseLayer, nn.Module):
    def __init__(self, mbart_layer, config):
        r"""
        A simple conversion of the `MBartEncoderLayer` to its `BetterTransformer` implementation.
        Args:
            mbart_layer (`torch.nn.Module`):
                The original `MBartEncoderLayer` where the weights needs to be retrieved.
        """
        super().__init__(config)
        super(BetterTransformerBaseLayer, self).__init__()
        # In_proj layer
        self.in_proj_weight = nn.Parameter(
            torch.cat(
                [
                    mbart_layer.self_attn.q_proj.weight,
                    mbart_layer.self_attn.k_proj.weight,
                    mbart_layer.self_attn.v_proj.weight,
                ]
            )
        )

        self.in_proj_bias = nn.Parameter(
            torch.cat(
                [
                    mbart_layer.self_attn.q_proj.bias,
                    mbart_layer.self_attn.k_proj.bias,
                    mbart_layer.self_attn.v_proj.bias,
                ]
            )
        )

        # Out proj layer
        self.out_proj_weight = mbart_layer.self_attn.out_proj.weight
        self.out_proj_bias = mbart_layer.self_attn.out_proj.bias

        # Linear layer 1
        self.linear1_weight = mbart_layer.fc1.weight
        self.linear1_bias = mbart_layer.fc1.bias

        # Linear layer 2
        self.linear2_weight = mbart_layer.fc2.weight
        self.linear2_bias = mbart_layer.fc2.bias

        # Layer norm 1
        self.norm1_eps = mbart_layer.self_attn_layer_norm.eps
        self.norm1_weight = mbart_layer.self_attn_layer_norm.weight
        self.norm1_bias = mbart_layer.self_attn_layer_norm.bias

        # Layer norm 2
        self.norm2_eps = mbart_layer.final_layer_norm.eps
        self.norm2_weight = mbart_layer.final_layer_norm.weight
        self.norm2_bias = mbart_layer.final_layer_norm.bias

        # Model hyper parameters
        self.num_heads = mbart_layer.self_attn.num_heads
        self.embed_dim = mbart_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.original_layers_mapping = {
            "in_proj_weight": [
                "self_attn.q_proj.weight",
                "self_attn.k_proj.weight",
                "self_attn.v_proj.weight",
            ],
            "in_proj_bias": ["self_attn.q_proj.bias", "self_attn.k_proj.bias", "self_attn.v_proj.bias"],
            "out_proj_weight": "self_attn.out_proj.weight",
            "out_proj_bias": "self_attn.out_proj.bias",
            "linear1_weight": "fc1.weight",
            "linear1_bias": "fc1.bias",
            "linear2_weight": "fc2.weight",
            "linear2_bias": "fc2.bias",
            "norm1_weight": "self_attn_layer_norm.weight",
            "norm1_bias": "self_attn_layer_norm.bias",
            "norm1_eps": "self_attn_layer_norm.eps",
            "norm2_weight": "final_layer_norm.weight",
            "norm2_bias": "final_layer_norm.bias",
            "norm2_eps": "final_layer_norm.eps",
        }
        self.dropout = config.attention_dropout
        self.activation_dropout = config.activation_dropout
        self.attention_head_size = config.d_model // config.encoder_attention_heads
        self.act_fn_callable = ACT2FN[self.act_fn]

        self.validate_bettertransformer()

    def forward(self, hidden_states, attention_mask, output_attentions: bool, position_bias=None, *_, **__):
        if output_attentions:
            raise ValueError("output_attentions=True can not be supported with BetterTransformer.")

        if not self.training and not torch.is_autocast_enabled() and not torch.is_autocast_cpu_enabled():
            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
                if len(attention_mask.shape) == 4:
                    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)
        else:
            residual = hidden_states
            hidden_states = F.layer_norm(
                hidden_states,
                normalized_shape=self.norm1_weight.shape,
                weight=self.norm1_weight,
                bias=self.norm1_bias,
            )

            qkv = F.linear(hidden_states, weight=self.in_proj_weight, bias=self.in_proj_bias)
            qkv = qkv.view(qkv.size()[:-1] + (3, self.num_heads, self.attention_head_size)).permute(2, 0, 3, 1, 4)
            query, key, value = qkv[0], qkv[1], qkv[2]

            # NOTE: In PyTorch 2.0, passing an attention_mask will automatically dispatch
            # to the "math" path and will NOT use flash attention / memory-efficient attention.
            # We should support xformers / Hazy-flash / rocm-flash directly and stop relying on PyTorch to do the work.
            if self.training:
                attention_mask = None
            attention_out = F.scaled_dot_product_attention(
                query,
                key,
                value,
                attn_mask=attention_mask,
                is_causal=False,
                dropout_p=self.dropout if self.training else 0.0,
            )

            attention_out = attention_out.permute(0, 2, 1, 3).contiguous()
            new_attention_out_shape = attention_out.size()[:-2] + (self.num_heads * self.attention_head_size,)
            attention_out = attention_out.view(new_attention_out_shape)

            hidden_states = residual + F.dropout(
                F.linear(attention_out, self.out_proj_weight, self.out_proj_bias),
                p=self.dropout,
                training=self.training,
            )
            residual = hidden_states
            hidden_states = F.layer_norm(
                hidden_states,
                normalized_shape=self.norm2_weight.shape,
                weight=self.norm2_weight,
                bias=self.norm2_bias,
            )

            # One additional dropout compared to bert
            hidden_states = F.dropout(
                self.act_fn_callable(F.linear(hidden_states, self.linear1_weight, self.linear1_bias)),
                p=self.activation_dropout,
                training=self.training,
            )

            hidden_states = residual + F.dropout(
                F.linear(hidden_states, self.linear2_weight, self.linear2_bias),
                p=self.dropout,
                training=self.training,
            )

        return (hidden_states,)


class DistilBertLayerBetterTransformer(BetterTransformerBaseLayer, nn.Module):
    def __init__(self, bert_layer, config):
        r"""
        A simple conversion of the Distill-BERTLayer to its `BetterTransformer` implementation.

        Args:
            bert_layer (`torch.nn.Module`):
                The original Distill-BERT Layer where the weights needs to be retrieved.
        """
        super().__init__(config)
        super(BetterTransformerBaseLayer, self).__init__()
        # In_proj layer
        self.in_proj_weight = nn.Parameter(
            torch.cat(
                [
                    bert_layer.attention.q_lin.weight,
                    bert_layer.attention.k_lin.weight,
                    bert_layer.attention.v_lin.weight,
                ]
            )
        )
        self.in_proj_bias = nn.Parameter(
            torch.cat(
                [
                    bert_layer.attention.q_lin.bias,
                    bert_layer.attention.k_lin.bias,
                    bert_layer.attention.v_lin.bias,
                ]
            )
        )

        # Out proj layer
        self.out_proj_weight = bert_layer.attention.out_lin.weight
        self.out_proj_bias = bert_layer.attention.out_lin.bias

        # Linear layer 1
        self.linear1_weight = bert_layer.ffn.lin1.weight
        self.linear1_bias = bert_layer.ffn.lin1.bias

        # Linear layer 2
        self.linear2_weight = bert_layer.ffn.lin2.weight
        self.linear2_bias = bert_layer.ffn.lin2.bias

        # Layer norm 1
        self.norm1_eps = bert_layer.sa_layer_norm.eps
        self.norm1_weight = bert_layer.sa_layer_norm.weight
        self.norm1_bias = bert_layer.sa_layer_norm.bias

        # Layer norm 2
        self.norm2_eps = bert_layer.output_layer_norm.eps
        self.norm2_weight = bert_layer.output_layer_norm.weight
        self.norm2_bias = bert_layer.output_layer_norm.bias

        # Model hyper parameters
        self.num_heads = bert_layer.attention.n_heads
        self.embed_dim = bert_layer.attention.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.original_layers_mapping = {
            "in_proj_weight": ["attention.q_lin.weight", "attention.k_lin.weight", "attention.v_lin.weight"],
            "in_proj_bias": ["attention.q_lin.bias", "attention.k_lin.bias", "attention.v_lin.bias"],
            "out_proj_weight": "attention.out_lin.weight",
            "out_proj_bias": "attention.out_lin.bias",
            "linear1_weight": "ffn.lin1.weight",
            "linear1_bias": "ffn.lin1.bias",
            "linear2_weight": "ffn.lin2.weight",
            "linear2_bias": "ffn.lin2.bias",
            "norm1_weight": "sa_layer_norm.weight",
            "norm1_bias": "sa_layer_norm.bias",
            "norm2_weight": "output_layer_norm.weight",
            "norm2_bias": "output_layer_norm.bias",
        }
        self.attention_dropout = config.attention_dropout
        self.dropout = config.dropout
        self.attention_head_size = config.dim // config.n_heads
        self.act_fn_callable = ACT2FN[self.act_fn]

        self.validate_bettertransformer()

    def forward(self, hidden_states, attn_mask, output_attentions: bool, head_mask=None, *_):
        if output_attentions:
            raise ValueError("output_attentions=True can not be supported with BetterTransformer.")

        if not self.training and not torch.is_autocast_enabled() and not torch.is_autocast_cpu_enabled():
            if hidden_states.is_nested:
                attn_mask = None

            if attn_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
                attn_mask = attn_mask.bool()
                attn_mask = torch.reshape(attn_mask, (attn_mask.shape[0], attn_mask.shape[-1]))
                seqlen = attn_mask.shape[1]
                lengths = torch.sum(~attn_mask, 1)
                if not all(l == seqlen for l in lengths):
                    hidden_states = torch._nested_tensor_from_mask(hidden_states, attn_mask)
                attn_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,
                attn_mask,
            )
            if hidden_states.is_nested and self.is_last_layer:
                hidden_states = hidden_states.to_padded_tensor(0.0)
        else:
            qkv = F.linear(hidden_states, weight=self.in_proj_weight, bias=self.in_proj_bias)

            qkv = qkv.view(qkv.size()[:-1] + (3, self.num_heads, self.attention_head_size)).permute(2, 0, 3, 1, 4)
            query, key, value = qkv[0], qkv[1], qkv[2]

            # TODO: Kind of stupid to do that at each layer, should be fixed in transformers
            attn_mask = attn_mask.unsqueeze(1).unsqueeze(2).to(dtype=query.dtype)
            attn_mask = (1.0 - attn_mask) * torch.finfo(query.dtype).min

            # NOTE: In PyTorch 2.0, passing an attention_mask will automatically dispatch
            # to the "math" path and will NOT use flash attention / memory-efficient attention.
            # We should support xformers / Hazy-flash / rocm-flash directly and stop relying on PyTorch to do the work.
            if self.training:
                attn_mask = None
            attention_out = F.scaled_dot_product_attention(
                query,
                key,
                value,
                attn_mask=attn_mask,
                is_causal=False,
                dropout_p=self.attention_dropout if self.training else 0.0,
            )

            attention_out = attention_out.permute(0, 2, 1, 3).contiguous()
            new_attention_out_shape = attention_out.size()[:-2] + (self.num_heads * self.attention_head_size,)
            attention_out = attention_out.view(new_attention_out_shape)

            # BertSelfOutput
            attention_out = F.layer_norm(
                F.dropout(
                    F.linear(attention_out, self.out_proj_weight, self.out_proj_bias),
                    p=self.dropout,
                    training=self.training,
                )
                + hidden_states,
                normalized_shape=self.norm1_weight.shape,
                weight=self.norm1_weight,
                bias=self.norm1_bias,
            )

            # BertIntermediate
            hidden_states = self.act_fn_callable(F.linear(attention_out, self.linear1_weight, self.linear1_bias))

            # BertOutput
            hidden_states = F.layer_norm(
                attention_out
                + F.dropout(
                    F.linear(hidden_states, self.linear2_weight, self.linear2_bias),
                    p=self.dropout,
                    training=self.training,
                ),
                normalized_shape=self.norm2_weight.shape,
                weight=self.norm2_weight,
                bias=self.norm2_bias,
            )
        return (hidden_states,)


class ViTLayerBetterTransformer(BetterTransformerBaseLayer, nn.Module):
    def __init__(self, vit_layer, config):
        r"""
        A simple conversion of the ViTLayer to its `BetterTransformer` implementation.

        Args:
            vit_layer (`torch.nn.Module`):
                The original `ViTLayer` where the weights needs to be retrieved.
        """
        super().__init__(config)
        super(BetterTransformerBaseLayer, self).__init__()
        # In_proj layer
        self.in_proj_weight = nn.Parameter(
            torch.cat(
                [
                    vit_layer.attention.attention.query.weight,
                    vit_layer.attention.attention.key.weight,
                    vit_layer.attention.attention.value.weight,
                ]
            )
        )
        self.in_proj_bias = nn.Parameter(
            torch.cat(
                [
                    vit_layer.attention.attention.query.bias,
                    vit_layer.attention.attention.key.bias,
                    vit_layer.attention.attention.value.bias,
                ]
            )
        )

        # Out proj layer
        self.out_proj_weight = vit_layer.attention.output.dense.weight
        self.out_proj_bias = vit_layer.attention.output.dense.bias

        # Linear layer 1
        self.linear1_weight = vit_layer.intermediate.dense.weight
        self.linear1_bias = vit_layer.intermediate.dense.bias

        # Linear layer 2
        self.linear2_weight = vit_layer.output.dense.weight
        self.linear2_bias = vit_layer.output.dense.bias

        # Layer norm 1
        self.norm1_eps = vit_layer.layernorm_before.eps
        self.norm1_weight = vit_layer.layernorm_before.weight
        self.norm1_bias = vit_layer.layernorm_before.bias

        # Layer norm 2
        self.norm2_eps = vit_layer.layernorm_after.eps
        self.norm2_weight = vit_layer.layernorm_after.weight
        self.norm2_bias = vit_layer.layernorm_after.bias

        # Model hyper parameters
        self.num_heads = vit_layer.attention.attention.num_attention_heads
        self.embed_dim = int(vit_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.original_layers_mapping = {
            "in_proj_weight": [
                "attention.attention.query.weight",
                "attention.attention.key.weight",
                "attention.attention.value.weight",
            ],
            "in_proj_bias": [
                "attention.attention.query.bias",
                "attention.attention.key.bias",
                "attention.attention.value.bias",
            ],
            "out_proj_weight": "attention.output.dense.weight",
            "out_proj_bias": "attention.output.dense.bias",
            "linear1_weight": "intermediate.dense.weight",
            "linear1_bias": "intermediate.dense.bias",
            "linear2_weight": "output.dense.weight",
            "linear2_bias": "output.dense.bias",
            "norm1_weight": "layernorm_before.weight",
            "norm1_bias": "layernorm_before.bias",
            "norm2_weight": "layernorm_after.weight",
            "norm2_bias": "layernorm_after.bias",
        }

        self.validate_bettertransformer()

    def forward(self, hidden_states, output_attentions: bool, *_, **__):
        if output_attentions:
            raise ValueError("output_attentions=True can not be supported with BetterTransformer.")

        if not self.training and not torch.is_autocast_enabled() and not torch.is_autocast_cpu_enabled():
            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)
        else:
            raise NotImplementedError(
                "Training and Autocast are not implemented for BetterTransformer + ViT. Please open an issue."
            )
        return (hidden_states,)


class ViltLayerBetterTransformer(BetterTransformerBaseLayer, nn.Module):
    def __init__(self, vilt_layer, config):
        r"""
        A simple conversion of the VilTLayer to its `BetterTransformer` implementation.

        Args:
            vilt_layer (`torch.nn.Module`):
                The original `VilTLayer` where the weights needs to be retrieved.
        """
        super().__init__(config)
        super(BetterTransformerBaseLayer, self).__init__()
        # In_proj layer
        self.in_proj_weight = nn.Parameter(
            torch.cat(
                [
                    vilt_layer.attention.attention.query.weight,
                    vilt_layer.attention.attention.key.weight,
                    vilt_layer.attention.attention.value.weight,
                ]
            )
        )
        self.in_proj_bias = nn.Parameter(
            torch.cat(
                [
                    vilt_layer.attention.attention.query.bias,
                    vilt_layer.attention.attention.key.bias,
                    vilt_layer.attention.attention.value.bias,
                ]
            )
        )

        # Out proj layer
        self.out_proj_weight = vilt_layer.attention.output.dense.weight
        self.out_proj_bias = vilt_layer.attention.output.dense.bias

        # Linear layer 1
        self.linear1_weight = vilt_layer.intermediate.dense.weight
        self.linear1_bias = vilt_layer.intermediate.dense.bias

        # Linear layer 2
        self.linear2_weight = vilt_layer.output.dense.weight
        self.linear2_bias = vilt_layer.output.dense.bias

        # Layer norm 1
        self.norm1_eps = vilt_layer.layernorm_before.eps
        self.norm1_weight = vilt_layer.layernorm_before.weight
        self.norm1_bias = vilt_layer.layernorm_before.bias

        # Layer norm 2
        self.norm2_eps = vilt_layer.layernorm_after.eps
        self.norm2_weight = vilt_layer.layernorm_after.weight
        self.norm2_bias = vilt_layer.layernorm_after.bias

        # Model hyper parameters
        self.num_heads = vilt_layer.attention.attention.num_attention_heads
        self.embed_dim = int(vilt_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.original_layers_mapping = {
            "in_proj_weight": [
                "attention.attention.query.weight",
                "attention.attention.key.weight",
                "attention.attention.value.weight",
            ],
            "in_proj_bias": [
                "attention.attention.query.bias",
                "attention.attention.key.bias",
                "attention.attention.value.bias",
            ],
            "out_proj_weight": "attention.output.dense.weight",
            "out_proj_bias": "attention.output.dense.bias",
            "linear1_weight": "intermediate.dense.weight",
            "linear1_bias": "intermediate.dense.bias",
            "linear2_weight": "output.dense.weight",
            "linear2_bias": "output.dense.bias",
            "norm1_weight": "layernorm_before.weight",
            "norm1_bias": "layernorm_before.bias",
            "norm2_weight": "layernorm_after.weight",
            "norm2_bias": "layernorm_after.bias",
        }

        self.validate_bettertransformer()

    def forward(self, hidden_states, layer_head_mask, output_attentions: bool, *_, **__):
        if output_attentions:
            raise ValueError("output_attentions=True can not be supported with BetterTransformer.")

        if not self.training and not torch.is_autocast_enabled() and not torch.is_autocast_cpu_enabled():
            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)
        else:
            raise NotImplementedError(
                "Training and Autocast are not implemented for BetterTransformer + Vilt. Please open an issue."
            )
        return (hidden_states,)


class Wav2Vec2EncoderLayerBetterTransformer(BetterTransformerBaseLayer, nn.Module):
    def __init__(self, wav2vec2_layer, config):
        r"""
        A simple conversion of the Wav2Vec2EncoderLayer to its `BetterTransformer` implementation.

        Args:
            wav2vec2_layer (`torch.nn.Module`):
                The original `Wav2Vec2EncoderLayer` where the weights needs to be retrieved.
        """
        super().__init__(config)
        super(BetterTransformerBaseLayer, self).__init__()
        # In_proj layer
        self.in_proj_weight = nn.Parameter(
            torch.cat(
                [
                    wav2vec2_layer.attention.q_proj.weight,
                    wav2vec2_layer.attention.k_proj.weight,
                    wav2vec2_layer.attention.v_proj.weight,
                ]
            )
        )
        self.in_proj_bias = nn.Parameter(
            torch.cat(
                [
                    wav2vec2_layer.attention.q_proj.bias,
                    wav2vec2_layer.attention.k_proj.bias,
                    wav2vec2_layer.attention.v_proj.bias,
                ]
            )
        )

        # Out proj layer
        self.out_proj_weight = wav2vec2_layer.attention.out_proj.weight
        self.out_proj_bias = wav2vec2_layer.attention.out_proj.bias

        # Linear layer 1
        self.linear1_weight = wav2vec2_layer.feed_forward.intermediate_dense.weight
        self.linear1_bias = wav2vec2_layer.feed_forward.intermediate_dense.bias

        # Linear layer 2
        self.linear2_weight = wav2vec2_layer.feed_forward.output_dense.weight
        self.linear2_bias = wav2vec2_layer.feed_forward.output_dense.bias

        # Layer norm 1
        self.norm1_eps = wav2vec2_layer.layer_norm.eps
        self.norm1_weight = wav2vec2_layer.layer_norm.weight
        self.norm1_bias = wav2vec2_layer.layer_norm.bias

        # Layer norm 2
        self.norm2_eps = wav2vec2_layer.final_layer_norm.eps
        self.norm2_weight = wav2vec2_layer.final_layer_norm.weight
        self.norm2_bias = wav2vec2_layer.final_layer_norm.bias

        # Model hyper parameters
        self.num_heads = wav2vec2_layer.attention.num_heads
        self.embed_dim = wav2vec2_layer.attention.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.original_layers_mapping = {
            "in_proj_weight": ["attention.q_proj.weight", "attention.k_proj.weight", "attention.v_proj.weight"],
            "in_proj_bias": ["attention.q_proj.bias", "attention.k_proj.bias", "attention.v_proj.bias"],
            "out_proj_weight": "attention.out_proj.weight",
            "out_proj_bias": "attention.out_proj.bias",
            "linear1_weight": "feed_forward.intermediate_dense.weight",
            "linear1_bias": "feed_forward.intermediate_dense.bias",
            "linear2_weight": "feed_forward.output_dense.weight",
            "linear2_bias": "feed_forward.output_dense.bias",
            "norm1_weight": "layer_norm.weight",
            "norm1_bias": "layer_norm.bias",
            "norm1_eps": "layer_norm.eps",
            "norm2_weight": "final_layer_norm.weight",
            "norm2_bias": "final_layer_norm.bias",
            "norm2_eps": "final_layer_norm.eps",
        }

        if config.do_stable_layer_norm:
            self.norm_first = True

        self.validate_bettertransformer()

    def forward(self, hidden_states, attention_mask, output_attentions: bool, **__):
        if output_attentions:
            raise ValueError("output_attentions=True can not be supported with BetterTransformer.")

        if not self.training and not torch.is_autocast_enabled() and not torch.is_autocast_cpu_enabled():
            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.bool()
                if len(attention_mask.shape) == 4:
                    attention_mask = attention_mask.squeeze(1)[:, 0]
                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 hidden_states.is_nested and self.is_last_layer:
                hidden_states = hidden_states.to_padded_tensor(0.0)
        else:
            raise NotImplementedError(
                "Training and Autocast are not implemented for BetterTransformer + Wav2Vec2. Please open an issue."
            )
        return (hidden_states,)


class FSMTEncoderLayerBetterTransformer(BetterTransformerBaseLayer, nn.Module):
    def __init__(self, fsmt_layer, config):
        r"""
        A simple conversion of the FSMT Encoder layer to its `BetterTransformer` implementation.

        Args:
            fsmt_layer (`torch.nn.Module`):
                The original FSMT Layer where the weights needs to be retrieved.
        """
        super().__init__(config)
        super(BetterTransformerBaseLayer, self).__init__()
        # In_proj layer
        self.in_proj_weight = nn.Parameter(
            torch.cat(
                [
                    fsmt_layer.self_attn.q_proj.weight,
                    fsmt_layer.self_attn.k_proj.weight,
                    fsmt_layer.self_attn.v_proj.weight,
                ]
            )
        )
        self.in_proj_bias = nn.Parameter(
            torch.cat(
                [
                    fsmt_layer.self_attn.q_proj.bias,
                    fsmt_layer.self_attn.k_proj.bias,
                    fsmt_layer.self_attn.v_proj.bias,
                ]
            )
        )

        # Out proj layer
        self.out_proj_weight = fsmt_layer.self_attn.out_proj.weight
        self.out_proj_bias = fsmt_layer.self_attn.out_proj.bias

        # Linear layer 1
        self.linear1_weight = fsmt_layer.fc1.weight
        self.linear1_bias = fsmt_layer.fc1.bias

        # Linear layer 2
        self.linear2_weight = fsmt_layer.fc2.weight
        self.linear2_bias = fsmt_layer.fc2.bias

        # Layer norm 1
        self.norm1_eps = fsmt_layer.self_attn_layer_norm.eps
        self.norm1_weight = fsmt_layer.self_attn_layer_norm.weight
        self.norm1_bias = fsmt_layer.self_attn_layer_norm.bias

        # Layer norm 2
        self.norm2_eps = fsmt_layer.final_layer_norm.eps
        self.norm2_weight = fsmt_layer.final_layer_norm.weight
        self.norm2_bias = fsmt_layer.final_layer_norm.bias

        # Model hyper parameters
        self.num_heads = fsmt_layer.self_attn.num_heads
        self.embed_dim = fsmt_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.original_layers_mapping = {
            "in_proj_weight": ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"],
            "in_proj_bias": ["self_attn.q_proj.bias", "self_attn.k_proj.bias", "self_attn.v_proj.bias"],
            "out_proj_weight": "self_attn.out_proj.weight",
            "out_proj_bias": "self_attn.out_proj.bias",
            "linear1_weight": "fc1.weight",
            "linear1_bias": "fc1.bias",
            "linear2_weight": "fc2.weight",
            "linear2_bias": "fc2.bias",
            "norm1_weight": "self_attn_layer_norm.weight",
            "norm1_bias": "self_attn_layer_norm.bias",
            "norm2_weight": "final_layer_norm.weight",
            "norm2_bias": "final_layer_norm.bias",
        }

        self.validate_bettertransformer()

    def forward(self, hidden_states, attention_mask, output_attentions: bool, position_bias=None, *_, **__):
        if output_attentions:
            raise ValueError("output_attentions=True can not be supported with BetterTransformer.")

        if not self.training and not torch.is_autocast_enabled() and not torch.is_autocast_cpu_enabled():
            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.bool()
                attention_mask = torch.reshape(attention_mask, (attention_mask.shape[0], attention_mask.shape[-1]))

                # FSMT swaps the first two axis before calling the encoder stack
                # Reference: https://github.com/huggingface/transformers/blob/699e90437f984d69ad3c9b891dd2e9d0fc2cffe4/src/transformers/models/fsmt/modeling_fsmt.py#L508
                if hidden_states.shape[0] != attention_mask.shape[0]:
                    hidden_states = hidden_states.transpose(1, 0)
                    original_shape = hidden_states.shape

                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)
        else:
            raise NotImplementedError(
                "Training and Autocast are not implemented for BetterTransformer + FSMT. Please open an issue."
            )

        return (hidden_states, attention_mask)


class ProphetNetEncoderLayerBetterTransformer(BetterTransformerBaseLayer, nn.Module):
    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)
        super(BetterTransformerBaseLayer, self).__init__()
        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, output_attentions: bool, *_, **__):
        if output_attentions:
            raise ValueError("output_attentions=True can not be supported with BetterTransformer.")

        if not self.training and not torch.is_autocast_enabled() and not torch.is_autocast_cpu_enabled():
            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)
        else:
            raise ValueError(
                "Training and Autocast are not implemented for BetterTransformer + ProphetNet. Please open an issue."
            )

        return (hidden_states,)


class CLIPLayerBetterTransformer(BetterTransformerBaseLayer, nn.Module):
    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)
        super(BetterTransformerBaseLayer, self).__init__()
        # 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.original_layers_mapping = {
            "in_proj_weight": ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"],
            "in_proj_bias": ["self_attn.q_proj.bias", "self_attn.k_proj.bias", "self_attn.v_proj.bias"],
            "out_proj_weight": "self_attn.out_proj.weight",
            "out_proj_bias": "self_attn.out_proj.bias",
            "linear1_weight": "mlp.fc1.weight",
            "linear1_bias": "mlp.fc1.bias",
            "linear2_weight": "mlp.fc2.weight",
            "linear2_bias": "mlp.fc2.bias",
            "norm1_eps": "layer_norm1.eps",
            "norm1_weight": "layer_norm1.weight",
            "norm1_bias": "layer_norm1.bias",
            "norm2_eps": "layer_norm2.eps",
            "norm2_weight": "layer_norm2.weight",
            "norm2_bias": "layer_norm2.bias",
        }

        self.validate_bettertransformer()

    def forward(self, hidden_states, attention_mask, causal_attention_mask, output_attentions: bool, *_, **__):
        if output_attentions:
            raise ValueError("output_attentions=True can not be supported with BetterTransformer.")

        if not self.training and not torch.is_autocast_enabled() and not torch.is_autocast_cpu_enabled():
            # we expect attention_mask to be None in the vision model
            if attention_mask is not None or causal_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,
            )
        else:
            raise NotImplementedError(
                "Training and Autocast are not implemented for BetterTransformer + CLIP. Please open an issue."
            )

        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