[PT2] Context manager to disable tracing (#3361)

### Changes

Add context manager to disable tracing for some part of code in a model

Author: AlexanderDokuchaev

nncf/experimental/torch2/function_hook/hook_executor_mode.py

@@ -517,3 +517,17 @@ def disable(self) -> Iterator[None]:
         self.enabled = False
         yield
         self.enabled = ret
+
+
+@contextmanager
+def disable_function_hook_mode() -> Iterator[None]:
+    """
+    Temporarily disables the function tracing and execution hooks within a context.
+    """
+    enabled_modes = torch.overrides._get_current_function_mode_stack()  # type: ignore[no-untyped-call]
+    state = {(mode, mode.enabled) for mode in enabled_modes if isinstance(mode, FunctionHookMode)}
+    for mode, _ in state:
+        mode.enabled = False
+    yield
+    for mode, enabled in state:
+        mode.enabled = enabled

nncf/torch/dynamic_graph/context.py

@@ -25,6 +25,8 @@
 from nncf.common.utils.api_marker import api
 from nncf.common.utils.debug import is_debug
 from nncf.common.utils.patcher import PATCHER
+from nncf.experimental.common.check_feature import is_experimental_torch_tracing_enabled
+from nncf.experimental.torch2.function_hook.hook_executor_mode import disable_function_hook_mode
 from nncf.torch.dynamic_graph.graph import DynamicGraph
 from nncf.torch.dynamic_graph.graph import DynamicGraphNode
 from nncf.torch.dynamic_graph.graph import DynamicGraphNodeParameters
@@ -505,9 +507,15 @@ def disable_tracing(method):
     Patch a method so that it will be executed within no_nncf_trace context
     :param method: A method to patch.
     """
+    if is_experimental_torch_tracing_enabled():
 
-    def no_nncf_trace_wrapper(self, fn, *args, **kwargs):
-        with no_nncf_trace():
-            return fn(*args, **kwargs)
+        def no_nncf_trace_wrapper(self, fn, *args, **kwargs):
+            with disable_function_hook_mode():
+                return fn(*args, **kwargs)
+    else:
+
+        def no_nncf_trace_wrapper(self, fn, *args, **kwargs):
+            with no_nncf_trace():
+                return fn(*args, **kwargs)
 
     PATCHER.patch(method, no_nncf_trace_wrapper)

nncf/torch/dynamic_graph/patch_pytorch.py

@@ -181,11 +181,18 @@ class MagicFunctionsToPatch:
 
 @api(canonical_alias="nncf.torch.register_operator")
 def register_operator(name=None):
-    def wrap(operator):
-        op_name = name
-        if op_name is None:
-            op_name = operator.__name__
-        return wrap_operator(operator, PatchedOperatorInfo(op_name, NamespaceTarget.EXTERNAL))
+    if is_experimental_torch_tracing_enabled():
+
+        def wrap(operator):
+            # Skip wrapping operator for tracing by TorchFunctionMode
+            return operator
+    else:
+
+        def wrap(operator):
+            op_name = name
+            if op_name is None:
+                op_name = operator.__name__
+            return wrap_operator(operator, PatchedOperatorInfo(op_name, NamespaceTarget.EXTERNAL))
 
     return wrap
 

tests/torch2/data/function_hook/disable_tracing/graph.dot

@@ -0,0 +1,9 @@
+digraph {
+rankdir=TB;
+0 [label="{type: input|name: x|dtype: torch.float32|shape: (1, 1)}", fillcolor="#adadad", fontcolor="#000000", shape=record, style="filled,rounded"];
+1 [label="{type: function_call|op_name: /add/0|fn_name: add|args: [\nTensorMeta(dtype=torch.float32, shape=(1, 1)),\n1,\n]|kwargs: \{\}}", fillcolor="#caffbf", fontcolor="#000000", shape=record, style="filled,rounded"];
+2 [label="{type: function_call|op_name: /sub/0|fn_name: sub|args: [\nTensorMeta(dtype=torch.float32, shape=(1, 1)),\n1,\n]|kwargs: \{\}}", fillcolor="#ffadad", fontcolor="#000000", shape=record, style="filled,rounded"];
+3 [label="{type: output|name: output|dtype: torch.float32|shape: (1, 1)}", fillcolor="#adadad", fontcolor="#000000", shape=record, style="filled,rounded"];
+0 -> 1 [label="(1, 1)\n0 → 0"];
+2 -> 3 [label="(1, 1)\n0 → 0"];
+}

tests/torch2/function_hook/test_disable_tracing.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2025 Intel Corporation
+# 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.
+
+import torch
+from torch import nn
+from torch.overrides import _get_current_function_mode_stack
+
+from nncf.experimental.torch2.function_hook.graph.build_graph_mode import build_graph
+from nncf.experimental.torch2.function_hook.graph.graph_visualization import to_pydot
+from nncf.experimental.torch2.function_hook.hook_executor_mode import FunctionHookMode
+from nncf.experimental.torch2.function_hook.hook_executor_mode import disable_function_hook_mode
+from nncf.experimental.torch2.function_hook.wrapper import get_hook_storage
+from nncf.experimental.torch2.function_hook.wrapper import wrap_model
+from nncf.torch import disable_tracing
+from tests.cross_fw.shared.paths import TEST_ROOT
+from tests.torch2.utils import compare_with_reference_file
+
+REF_DIR = TEST_ROOT / "torch2" / "data" / "function_hook" / "disable_tracing"
+
+
+def test_disable_function_hook_mode():
+    model = wrap_model(nn.Conv2d(1, 1, 1))
+    with FunctionHookMode(model=model, hook_storage=get_hook_storage(model)) as ctx:
+        assert ctx.enabled
+        with disable_function_hook_mode():
+            assert not ctx.enabled
+        assert ctx.enabled
+
+
+class ModelNoTrace(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        x = x + 1
+        x = self.foo(x)
+        x = x - 1
+        return x
+
+    def foo(self, x):
+        mode = _get_current_function_mode_stack()
+        assert len(mode) == 1
+        assert isinstance(mode[0], FunctionHookMode)
+        assert not mode[0].enabled
+        return x - 1
+
+
+disable_tracing(ModelNoTrace.foo)
+
+
+def test_build_graph_with_disable_tracing(regen_ref_data):
+    model = wrap_model(ModelNoTrace())
+    graph = build_graph(model, torch.randn(1, 1))
+    dot_graph = to_pydot(graph)
+    compare_with_reference_file(str(dot_graph), REF_DIR / "graph.dot", regen_ref_data)

tests/torch2/function_hook/test_train.py

@@ -70,7 +70,7 @@ def patched_forward(*args, **kwargs):
     # In overridden bound method __self__ links to original model for all replicas
     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 overwriting forward method, expected ForwardWithHooks"):
+    with pytest.raises(nncf.InternalError, match="Not supported overridden forward"):
         run_one_epoch(parallel_model, optimizer, use_cuda=True)
 
 
@@ -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 overwriting forward method of original module"):
+    with pytest.raises(nncf.InternalError, match="Not supported overridden forward method"):
         run_one_epoch(parallel_model, optimizer, use_cuda=True)