[PT2] rework ForwardWithHooks (#3362)

### Changes

Rework ForwardWithHooks to be more flexible to patching forward
function.
Now model collects inside ForwardWithHooks and not used original forward
to find model.
Add `_has_compatible_shallow_copy_type` and `_parse_to` to ignored
function names.
Fix tests on cuda.

### Reason for changes

```python 
model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-GPTNeoXForCausalLM", torch_dtype=torch.bfloat16, device_map="auto")
model = wrap_model(model)

#    model = cast(nn.Module, self._func.__self__)  # type: ignore[attr-defined]
# AttributeError: 'functools.partial' object has no attribute '__self__'. Did you mean: '__call__'?
```

Author: AlexanderDokuchaev

nncf/experimental/torch2/function_hook/graph/build_graph_mode.py

@@ -373,6 +373,6 @@ def build_graph(model: nn.Module, *args: Any, **kwargs: Any) -> nx.MultiDiGraph:
             with GraphBuilderMode(model=model, hook_storage=get_hook_storage(model)) as ctx:
                 args, kwargs = ctx.process_model_inputs(args, kwargs)
                 wrapped_forward = cast(ForwardWithHooks, model.forward)
-                outputs = wrapped_forward._func(*args, **kwargs)
+                outputs = wrapped_forward.orig_forward(*args, **kwargs)
                 outputs = ctx.process_model_outputs(outputs)
     return ctx.graph

nncf/experimental/torch2/function_hook/hook_executor_mode.py

@@ -41,6 +41,8 @@
     "size",
     "is_floating_point",
     "_set_grad_enabled",
+    "_parse_to",
+    "_has_compatible_shallow_copy_type",
 ]
 
 

nncf/experimental/torch2/function_hook/wrapper.py

@@ -14,7 +14,7 @@
 import inspect
 import types
 from types import MethodType
-from typing import Any, Callable, Dict, Tuple, TypeVar, cast
+from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, cast
 
 from torch import nn
 
@@ -31,70 +31,76 @@
 class ForwardWithHooks:
     """Class to wrap forward function of nn.Module, to forward function of the model with enabled FunctionHookMode"""
 
-    __slots__ = "_func", "__dict__", "__weakref__"
-    _func: Callable[..., Any]
-
-    def __new__(cls, orig_forward: Callable[..., Any]) -> ForwardWithHooks:
-        if not callable(orig_forward):
-            msg = "the first argument must be callable"
-            raise TypeError(msg)
+    __slots__ = "_orig_forward", "_model", "__dict__", "__weakref__"
+    _orig_forward: Callable[..., Any]
+    _model: nn.Module
 
-        if isinstance(orig_forward, ForwardWithHooks):
+    def __new__(cls, model: nn.Module, orig_forward: Optional[Callable[..., Any]] = None) -> ForwardWithHooks:
+        if isinstance(model.forward, ForwardWithHooks):
             msg = "Func already wrapped"
             raise TypeError(msg)
 
         self = super().__new__(cls)
 
-        self._func = orig_forward
+        self._orig_forward = model.forward if orig_forward is None else orig_forward
+        self._model = model
         return self
 
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
-        model = cast(nn.Module, self._func.__self__)  # type: ignore[attr-defined]
-        with FunctionHookMode(model=model, hook_storage=get_hook_storage(model)) as ctx:
+        with FunctionHookMode(model=self.model, hook_storage=get_hook_storage(self.model)) as ctx:
             args, kwargs = ctx.process_model_inputs(args, kwargs)
-            outputs = self._func(*args, **kwargs)
+            outputs = self.orig_forward(*args, **kwargs)
             outputs = ctx.process_model_outputs(outputs)
             return outputs
 
     def __repr__(self) -> str:
-        return f"ForwardWithHooks.{repr(self._func)}"
+        return f"ForwardWithHooks.{repr(self.orig_forward)}"
 
-    def __reduce__(self) -> Tuple[Callable[..., Any], Tuple[Any, ...], Tuple[Any, ...]]:
-        return type(self), (self._func,), (self._func, self.__dict__ or None)
+    def __reduce__(self) -> Any:
+        return type(self), (self.model, self.orig_forward), (self.model, self.orig_forward, self.__dict__ or None)
 
-    def __setstate__(self, state: Tuple[Any, Any]) -> None:
+    def __setstate__(self, state: Any) -> None:
         if not isinstance(state, tuple):
             msg = "argument to __setstate__ must be a tuple"
             raise TypeError(msg)
-        if len(state) != 2:
-            msg = f"expected 2 items in state, got {len(state)}"
+        if len(state) != 3:
+            msg = f"expected 3 items in state, got {len(state)}"
             raise TypeError(msg)
-        func, namespace = state
-        if not callable(func) or (namespace is not None and not isinstance(namespace, dict)):
+        model, orig_forward, namespace = state
+        if not callable(orig_forward) or (namespace is not None and not isinstance(namespace, dict)):
             msg = "invalid partial state"
             raise TypeError(msg)
 
         if namespace is None:
             namespace = {}
 
+        self._model = model
+        self._orig_forward = orig_forward
         self.__dict__ = namespace
-        self._func = func
 
     @property
     def __code__(self) -> types.CodeType:
         return self.__call__.__code__
 
     @property
     def __globals__(self) -> Dict[str, Any]:
-        return self._func.__globals__
+        return self.orig_forward.__globals__
 
     @property
     def __name__(self) -> str:
-        return self._func.__name__
+        return self.orig_forward.__name__
 
     @property
     def __signature__(self) -> inspect.Signature:
-        return inspect.signature(self._func)
+        return inspect.signature(self.orig_forward)
+
+    @property
+    def orig_forward(self) -> Callable[..., Any]:
+        return self._orig_forward
+
+    @property
+    def model(self) -> nn.Module:
+        return self._model
 
 
 class ReplicateForDataParallel:
@@ -135,16 +141,16 @@ def __call__(self, *args: Any, **kwargs: Any) -> nn.Module:
             raise nncf.InternalError(msg)
 
         if not (
-            isinstance(saved_forward_with_hooks._func, types.MethodType)
-            and saved_forward_with_hooks._func.__func__ is module.__class__.forward
+            isinstance(saved_forward_with_hooks.orig_forward, types.MethodType)
+            and saved_forward_with_hooks.orig_forward.__func__ is module.__class__.forward
         ):
             msg = "Not supported overridden forward method of original module"
             raise nncf.InternalError(msg)
 
         module.__dict__.pop("forward")
 
-        replica: nn.Module = self._func(*args, **kwargs)
-        replica.forward = ForwardWithHooks(replica.forward)
+        replica: nn.Module = self.func(*args, **kwargs)
+        replica.forward = ForwardWithHooks(replica)
         module.forward = saved_forward_with_hooks
 
         return replica
@@ -193,7 +199,7 @@ def wrap_model(model: TModel) -> TModel:
     :param model: The nn.Module to be wrapped.
     :return: The modified model with the custom behavior injected.
     """
-    model.forward = ForwardWithHooks(model.forward)
+    model.forward = ForwardWithHooks(model)
     model._replicate_for_data_parallel = ReplicateForDataParallel(model._replicate_for_data_parallel)  # type: ignore
     model.add_module(ATR_HOOK_STORAGE, HookStorage())
     return model
@@ -236,7 +242,6 @@ def register_pre_function_hook(model: nn.Module, op_name: str, port_id: int, hoo
     :param op_name: The name of the operation associated with the hook.
     :param port_id: The port ID associated with the hook.
     :param hook: The pre-function hook module to be executed.
-
     :return: A handle that can be used to remove the hook later.
     """
     hook_storage = get_hook_storage(model)

tests/torch2/function_hook/test_train.py

@@ -87,7 +87,7 @@ def patched_forward(*args, **kwargs):
     wrapped_model.forward = patched_forward
     optimizer = torch.optim.Adam(wrapped_model.parameters(), lr=0.1)
     parallel_model = torch.nn.DataParallel(wrapped_model)
-    with pytest.raises(nncf.InternalError, match="Not supported overridden forward method"):
+    with pytest.raises(nncf.InternalError, match="Not supported overridden forward"):
         run_one_epoch(parallel_model, optimizer, use_cuda=True)
 
 

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,13 +27,40 @@
 ADD_VALUE = 2.0
 
 
-def test_wrapper():
+@pytest.mark.parametrize("forward_type", ["origin", "partial", "bound", "fn"])
+def test_wrapper(forward_type: str):
     example_input = helpers.ConvModel.get_example_inputs()
     model = helpers.ConvModel()
     model.eval()
-    ret = model(example_input)
-    wrapped = wrap_model(model)
-    wrapped_ret = wrapped(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)
+    elif forward_type == "fn":
+        old_forward = model.forward
+
+        def new_forward(x):
+            return old_forward(x)
+
+        model.forward = new_forward
+
+    with torch.no_grad():
+        ret = model(example_input)
+        wrapped = wrap_model(model)
+        wrapped_ret = wrapped(example_input)
+
     torch.testing.assert_close(ret, wrapped_ret)