[PT2] Serialize and load transformations (#3329)

### Changes

Implement `get_config` and `load_from_config` function for new tracing
Analog #2531


### Related tickets

152996

### Tests

tests/torch2/function_hook/test_serialization.py

Author: AlexanderDokuchaev

nncf/experimental/torch2/function_hook/serialization.py

@@ -0,0 +1,118 @@
+# 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.
+
+from typing import Any, Dict, List, TypedDict, TypeVar, cast
+from weakref import WeakKeyDictionary
+
+from torch import nn
+
+import nncf
+from nncf.experimental.torch2.function_hook.wrapper import get_hook_storage
+from nncf.experimental.torch2.function_hook.wrapper import wrap_model
+from nncf.torch.layer_utils import COMPRESSION_MODULES
+from nncf.torch.layer_utils import StatefullModuleInterface
+
+COMPRESSION_STATE_ATTR = "compression_state"
+TModel = TypeVar("TModel", bound=nn.Module)
+
+
+class S_COMMAND(TypedDict):
+    hook_names_in_model: List[str]
+    module_cls_name: str
+    module_config: Dict[str, Any]
+
+
+def get_config(model: nn.Module) -> Dict[str, Any]:
+    """
+    Returns serializable config which contains all information required to recover all additional modules placement.
+
+    :param model: The model to serialize.
+    :return: Serializable config.
+    """
+    hook_storage = get_hook_storage(model)
+
+    # Find shared modules
+    modules_map: WeakKeyDictionary[nn.Module, List[str]] = WeakKeyDictionary()
+    for name, module in hook_storage.named_modules(remove_duplicate=False):
+        splitted_name = name.split(".")
+        if len(splitted_name) != 3:
+            # Expected depths of target hook module is 3
+            # <3 - ModuleDicts in HookStorage, >3 - submodules of hooks
+            continue
+        if module not in modules_map:
+            modules_map[module] = []
+        modules_map[module].append(name)
+
+    # Generate serialized transformation commands
+    serialized_transformations: List[S_COMMAND] = []
+    for module, names in modules_map.items():
+        compression_module_name = module.__class__.__name__
+        if compression_module_name not in COMPRESSION_MODULES.registry_dict:
+            msg = (
+                f"Could not serialize compression module with name {compression_module_name}. "
+                "Please register your module in the COMPRESSION_MODULES registry."
+            )
+            raise nncf.InternalError(msg)
+        if not isinstance(module, StatefullModuleInterface):
+            msg = "Support only StatefullModuleInterface modules"
+            raise nncf.InternalError(msg)
+
+        serialized_transformations.append(
+            {
+                "hook_names_in_model": names,
+                "module_cls_name": compression_module_name,
+                "module_config": module.get_config(),
+            }
+        )
+
+    return {COMPRESSION_STATE_ATTR: serialized_transformations}
+
+
+def load_from_config(model: TModel, config: Dict[str, Any]) -> TModel:
+    """
+    Initialize model with compressed modules from config file.
+
+    .. code-block:: python
+
+        model = MyModel()
+        qmodel = nncf.quantize(model, ...)
+        torch.save(
+            {
+                "state_dict": qmodel.state_dict(),
+                "config": get_config(qmodel),
+            },
+            "ckpt.pth",
+        )
+        ...
+        ckpt = torch.load("ckpt.pth")
+        restored_model = load_from_config(MyModel(), ckpt["config"])
+        restored_model.load_state_dict(ckpt["state_dict"])
+
+    :param model: The original uncompressed model.
+    :param config: The configuration dictionary containing the compressed model information.
+    :return: The compressed model.
+    """
+    wrapped_model = wrap_model(model)
+    hook_storage = get_hook_storage(wrapped_model)
+    transformation_commands = cast(List[S_COMMAND], config[COMPRESSION_STATE_ATTR])
+    for command in transformation_commands:
+        module_cls = COMPRESSION_MODULES.get(command["module_cls_name"])
+        module = module_cls.from_config(command["module_config"])
+        for target_name in command["hook_names_in_model"]:
+            hook_type, hook_key, hook_id = target_name.split(".")
+            storage_dict = getattr(hook_storage, hook_type)
+            if hook_key not in storage_dict:
+                storage_dict[hook_key] = nn.ModuleDict()
+            if hook_id in storage_dict[hook_key]:
+                msg = f"{hook_id=} for {hook_type}.{hook_key} already registered"
+                raise nncf.InternalError(msg)
+            storage_dict[hook_key][hook_id] = module
+    return wrapped_model

tests/torch2/function_hook/helpers.py

@@ -14,6 +14,8 @@
 
 from nncf.experimental.torch2.function_hook.wrapper import register_post_function_hook
 from nncf.experimental.torch2.function_hook.wrapper import wrap_model
+from nncf.torch.layer_utils import COMPRESSION_MODULES
+from nncf.torch.layer_utils import StatefullModuleInterface
 
 
 class CallCount(torch.nn.Module):
@@ -160,3 +162,21 @@ def __init__(self):
     def forward(self, x):
         self.counter += 1
         return x + 1
+
+
+@COMPRESSION_MODULES.register()
+class HookWithState(torch.nn.Module, StatefullModuleInterface):
+    def __init__(self, state: str):
+        super().__init__()
+        self._state = state
+        self._dummy_param = torch.nn.Parameter(torch.tensor(1.0))
+
+    def forward(self, x):
+        return x + self._dummy_param
+
+    def get_config(self):
+        return self._state
+
+    @classmethod
+    def from_config(cls, state: str):
+        return cls(state)

tests/torch2/function_hook/test_serialization.py

@@ -0,0 +1,87 @@
+# 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.
+from pathlib import Path
+
+import pytest
+import torch
+from torch import nn
+
+import nncf
+from nncf.experimental.torch2.function_hook import get_hook_storage
+from nncf.experimental.torch2.function_hook import register_post_function_hook
+from nncf.experimental.torch2.function_hook import register_pre_function_hook
+from nncf.experimental.torch2.function_hook import wrap_model
+from nncf.experimental.torch2.function_hook.serialization import get_config
+from nncf.experimental.torch2.function_hook.serialization import load_from_config
+from tests.torch2.function_hook.helpers import HookWithState
+from tests.torch2.function_hook.helpers import SimpleModel
+
+
+@pytest.mark.parametrize("is_shared_hook", [True, False], ids=["shared_hook", "not_shared_hook"])
+def test_save_load(tmp_path: Path, is_shared_hook: bool):
+    model = wrap_model(SimpleModel())
+
+    hook1 = HookWithState("hook1")
+    hook2 = hook1 if is_shared_hook else HookWithState("hook2")
+
+    register_pre_function_hook(model, "conv1/conv2d/0", 0, hook1)
+    register_post_function_hook(model, "simple/conv/conv2d/0", 0, hook2)
+
+    state_dict = model.state_dict()
+    compression_config = get_config(model)
+
+    torch.save(
+        {
+            "model_state_dict": state_dict,
+            "compression_config": compression_config,
+        },
+        tmp_path / "checkpoint.pth",
+    )
+
+    ckpt = torch.load(tmp_path / "checkpoint.pth")
+    config = ckpt["compression_config"]
+    restored_model = load_from_config(SimpleModel(), config)
+    restored_model.load_state_dict(ckpt["model_state_dict"])
+
+    assert state_dict == restored_model.state_dict()
+
+    tensor = model.get_example_inputs()
+    ret_1 = model(tensor)
+    ret_2 = restored_model(tensor)
+    assert torch.allclose(ret_1[0], ret_2[0])
+    assert torch.allclose(ret_1[1], ret_2[1])
+
+    hook_storage = get_hook_storage(restored_model)
+    hook1 = hook_storage.get_submodule("pre_hooks.conv1/conv2d/0__0.0")
+    hook2 = hook_storage.get_submodule("post_hooks.simple/conv/conv2d/0__0.0")
+    assert (hook1 is hook2) == is_shared_hook
+
+
+def test_error_dublicate_names():
+    config = {
+        "compression_state": [
+            {
+                "hook_names_in_model": ["pre_hooks.conv1/conv2d/0__0.0", "pre_hooks.conv1/conv2d/0__0.0"],
+                "module_cls_name": "HookWithState",
+                "module_config": "hook1",
+            }
+        ]
+    }
+    with pytest.raises(nncf.InternalError, match="already registered"):
+        load_from_config(SimpleModel(), config)
+
+
+def test_error_not_registered_compression_modules():
+    model = wrap_model(SimpleModel())
+    register_pre_function_hook(model, "conv1/conv2d/0", 0, nn.ReLU())
+
+    with pytest.raises(nncf.InternalError, match="Please register your module in the COMPRESSION_MODULES registry."):
+        get_config(model)