Add bettertransformer reverse transform (huggingface#868)

* add reverse

* all tests should pass

* add adapted from

* fix tests

* fix test

Author: fxmarty

optimum/bettertransformer/models/base.py

@@ -11,7 +11,7 @@
 # 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, Optional
+from typing import TYPE_CHECKING
 
 
 if TYPE_CHECKING:
@@ -20,7 +20,7 @@
 import torch
 import torch.nn as nn
 
-from ...utils import logging, recurse_setattr
+from ...utils import logging, recurse_getattr, recurse_setattr
 
 
 KNOWN_ACTIVATION_ATTRIBUTES = ["hidden_act", "activation", "act_fn", "activation_function"]
@@ -38,7 +38,6 @@ class BetterTransformerBaseLayer(nn.Module):
     def __init__(
         self,
         config: "PretrainedConfig",
-        orig_layer: Optional[nn.Module] = None,
     ):
         r"""
         Base layer for `BetterTransformer` integration. This class is used to wrap all the necessary
@@ -57,6 +56,7 @@ def __init__(
         self.embed_dim = None
         self.num_layers = None
         self.original_layers_mapping = {}
+        self.module_mapping = None
         self.is_decoder = False
         # Some models does not have some attributes thus needs to be ignored
         # e.g. whisper does not have self_attn.k_proj.bias but has self_attn.v_proj.bias & self_attn.q_proj.bias
@@ -84,14 +84,6 @@ def __init__(
                 self.num_layers = getattr(config, attr)
                 break
 
-        # TODO: re-enable once fixed
-        # if orig_layer is not None:
-        #     # Last step, store the old module skeleton by copying the old module and putting
-        #     # it on the `meta` device.
-        #     self.orig_layer = deepcopy(orig_layer).to("meta")
-        # else:
-        #     self.orig_layer = orig_layer
-
     def validate_bettertransformer(self):
         r"""
         A wrapper function to validate the `BetterTransformer` implementation. Implements most relevant checks
@@ -147,32 +139,45 @@ def forward_checker(self, *args, **kwargs):
                 " Please use `model.eval()` before running the model.",
             )
 
-    def _revert_back_to_original_module(self):
-        r"""
-        A wrapper function to replace the current layer with the previous non-BetterTransformer
-        layer.
-        """
+    def _revert(self, module: torch.nn.Module) -> torch.nn.Module:
+        if self.module_mapping is not None:
+            if "" in self.module_mapping.values():
+                for bt_module_attr_name, value in self.module_mapping.items():
+                    if value == "":
+                        module = getattr(self, bt_module_attr_name)
+                        return module
+            else:
+                raise NotImplementedError("replacing a submodule in revert is not supported")
+
         for modified_layer_key_names, original_layer_key_names in self.original_layers_mapping.items():
             if isinstance(original_layer_key_names, list):
                 current_weight = getattr(self, modified_layer_key_names)
 
                 # Split the current weight n chunks - this is useful to split
                 # the qkv layers into q, k, v layers for example.
                 split_index = current_weight.shape[0] // len(original_layer_key_names)
-                for i, module in enumerate(original_layer_key_names):
+                for i, subparam_name in enumerate(original_layer_key_names):
+                    if recurse_getattr(module, subparam_name) is None:
+                        # this is for example the case if bias=False is set for a nn.Linear layer
+                        continue
+
                     if module not in self.keys_to_ignore:
-                        recurse_setattr(
-                            self.orig_layer,
-                            module,
-                            nn.Parameter(current_weight[i * split_index : (i + 1) * split_index]),
-                        )
+                        parameter = current_weight[i * split_index : (i + 1) * split_index]
+                        if isinstance(recurse_getattr(module, subparam_name), torch.nn.Parameter):
+                            parameter = torch.nn.Parameter(parameter)
+                        recurse_setattr(module, subparam_name, parameter)
             elif isinstance(original_layer_key_names, str):
-                if module not in self.keys_to_ignore:
-                    recurse_setattr(self.orig_layer, original_layer_key_names, getattr(self, modified_layer_key_names))
+                if recurse_getattr(module, original_layer_key_names) is None:
+                    # this is for example the case if bias=False is set for a nn.Linear layer
+                    continue
+
+                parameter = getattr(self, modified_layer_key_names)
+                if isinstance(recurse_getattr(module, original_layer_key_names), torch.nn.Parameter):
+                    parameter = torch.nn.Parameter(parameter)
+                recurse_setattr(module, original_layer_key_names, parameter)
             else:
                 raise ValueError(
                     f"Invalid type {type(modified_layer_key_names)} for `original_layers_mapping`",
                     " please use either `str` or `list`.",
                 )
-
-        return self.orig_layer
+        return module