suppport override forward by partial

Author: AlexanderDokuchaev

nncf/experimental/torch2/function_hook/wrapper.py

@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+import functools
 import inspect
 import types
 from types import MethodType
@@ -49,7 +50,12 @@ def __new__(cls, orig_forward: Callable[..., Any]) -> ForwardWithHooks:
         return self
 
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
-        model = cast(nn.Module, self._func.__self__)  # type: ignore[attr-defined]
+        if hasattr(self._func, "__self__"):
+            # For bound method module is stored in the __self__
+            model = cast(nn.Module, self._func.__self__)  # type: ignore[attr-defined]
+        elif isinstance(self._func, functools.partial):
+            # For using partial to override the forward method, module is stored in the args[0] and mapped to self arg.
+            model = cast(nn.Module, self._func.args[0])  # type: ignore[attr-defined]
         with FunctionHookMode(model=model, hook_storage=get_hook_storage(model)) as ctx:
             args, kwargs = ctx.process_model_inputs(args, kwargs)
             outputs = self._func(*args, **kwargs)

tests/torch2/function_hook/test_wrapper.py

@@ -9,7 +9,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import types
 from copy import deepcopy
+from functools import partial
 from pathlib import Path
 
 import onnxruntime as ort
@@ -25,11 +27,29 @@
 ADD_VALUE = 2.0
 
 
-def test_wrapper():
+@pytest.mark.parametrize("forward_type", ["origin", "partial", "methodtype"])
+def test_wrapper(forward_type: str):
     example_input = helpers.ConvModel.get_example_inputs()
     model = helpers.ConvModel()
     model.eval()
     ret = model(example_input)
+
+    model._old_forward = model.forward
+
+    if forward_type == "partial":
+        # Like in accelerate module
+        def new_forward(self, x):
+            return self._old_forward(x)
+
+        model.forward = partial(new_forward, model)
+
+    elif forward_type == "methodtype":
+
+        def new_forward(self, x):
+            return model._old_forward(x)
+
+        model.forward = types.MethodType(new_forward, model)
+
     wrapped = wrap_model(model)
     wrapped_ret = wrapped(example_input)
     torch.testing.assert_close(ret, wrapped_ret)