Shared param (#1664)

* enable shared parameter support for ipex.optimize

* auto format flake8

* manually fix format

* fix deepspeed prepack

* only use module_cls instead of use module in ParameterWrapper

* use module_para/optimizer_para instead of current_para/low_precision_cache in ParamWrapper

* auto format flake8

* manually fix flake8 error

* add TODO for mkl linear investigation

* fix setattr

* use parameter/master_parameter instead of module_parameter/optimizer_parameter

* Revert "fix setattr"

This reverts commit 177f3ff8dc8747519d8f5e4cff0374016a724dd9.

* Revert "manually fix flake8 error"

This reverts commit 50aecff126aeacac5ebcad9263cabd0a4cbc842a.

* Revert "auto format flake8"

This reverts commit 2284fb8e4da9481e64bc9bd5326a8a330275f9de.

* Revert "manually fix format"

This reverts commit f560177e52f5a3b60dcced74f8a947c6ceba0aaa.

* clean up codes

* clean up code and fix

* clean up code

Author: zhuhaozhe

intel_extension_for_pytorch/frontend.py

@@ -50,80 +50,51 @@ def _copy_model_and_optimizer(model, optimizer):
         return new_model, optimizer
     else:
         new_optimizer = copy.deepcopy(optimizer)
-        new_optimizer.state.clear()
         dic_param = {}
+        dic_param_for_master_case = {}
         for k, value in zip(model.parameters(), new_model.parameters()):
             dic_param[k] = value
+        if hasattr(optimizer, "params_attr"):
+            params_attr = getattr(optimizer, "params_attr")
+            param_key_pair = {}
+            if len(params_attr) != 0:
+                new_params_attr = copy.deepcopy(params_attr)
+                for (k1, v1), (k2, v2) in zip(
+                    params_attr.items(), new_params_attr.items()
+                ):
+                    if v1.master_parameter is None:
+                        v2.parameter = dic_param[v1.parameter]
+                    else:
+                        dic_param_for_master_case[k1] = k2
+                    param_key_pair[k1] = k2
+                if len(dic_param_for_master_case) != 0:
+                    dic_param = dic_param_for_master_case
+                for k, v in param_key_pair.items():
+                    new_params_attr[dic_param[k]] = new_params_attr.pop(v)
+                setattr(new_optimizer, "params_attr", new_params_attr)
 
+        new_optimizer.state.clear()
         # deep copy param_groups
         for group1, group2 in zip(optimizer.param_groups, new_optimizer.param_groups):
             for i, p in enumerate(group1["params"]):
-                # for the p not in the dic_param case, the new optimizer state will be updated
-                # in _deep_copy_params_attr because the param here in optimizer state is the master
-                # parameter of the model, which has ever optimized by ipex.optimize
                 if p in dic_param:
                     new_model_param = dic_param[p]
                     group2["params"][i] = new_model_param
                     new_optimizer.state[new_model_param] = copy.deepcopy(
                         optimizer.state[p]
                     )
 
-        # deep copy params_attr for reentrancy of ipex.optimize
-        def _deep_copy_params_attr(old_module, new_module):
+        def _attach_master_weight_split_attr(old_module, new_module):
             if hasattr(old_module, "master_weight_split"):
                 setattr(
                     new_module, "master_weight_split", old_module.master_weight_split
                 )
-                master_weight_split = getattr(new_module, "master_weight_split")
-
-                for name, param in old_module.named_parameters():
-                    if master_weight_split:
-                        attr_name = name + "_trail"
-                        if param in optimizer.params_attr:
-                            new_optimizer.params_attr[
-                                getattr(new_module, name)
-                            ] = optimizer.params_attr[param]
-                            new_optimizer.params_attr[getattr(new_module, name)][
-                                "trail"
-                            ] = getattr(new_module, attr_name)
-                    else:
-                        attr_name = "master_" + name
-                        old_master_param = getattr(old_module, attr_name)
-                        new_master_param = getattr(new_module, attr_name)
-                        if old_master_param in optimizer.params_attr:
-                            new_optimizer.params_attr[
-                                new_master_param
-                            ] = optimizer.params_attr[old_master_param]
-                            if (
-                                "bf16_param"
-                                in new_optimizer.params_attr[new_master_param]
-                            ):
-                                new_optimizer.params_attr[new_master_param][
-                                    "bf16_param"
-                                ] = getattr(new_module, name)
-                            if (
-                                "fp16_param"
-                                in new_optimizer.params_attr[new_master_param]
-                            ):
-                                new_optimizer.params_attr[new_master_param][
-                                    "fp16_param"
-                                ] = getattr(new_module, name)
-
-                        # deep copy new optimizer state for master parameter
-                        new_optimizer.state[new_master_param] = copy.deepcopy(
-                            optimizer.state[old_master_param]
-                        )
-
             for (_, old_child), (_, new_child) in zip(
                 old_module.named_children(), new_module.named_children()
             ):
-                _deep_copy_params_attr(old_child, new_child)
-
-        if hasattr(optimizer, "params_attr"):
-            params_attr = {}
-            setattr(new_optimizer, "params_attr", params_attr)
-            _deep_copy_params_attr(model, new_model)
+                _attach_master_weight_split_attr(old_child, new_child)
 
+        _attach_master_weight_split_attr(model, new_model)
         return new_model, new_optimizer
 
 
@@ -587,31 +558,30 @@ def xpu_check_channel_last():
         utils._weight_prepack.record_input_shape_for_prepack(
             optimized_model, sample_input
         )
-
+    params_attr = {}
     if not model.training:
         if opt_properties.conv_bn_folding:
             try:
-                optimized_model = optimization.fuse(optimized_model, inplace=inplace)
+                optimized_model = optimization.fuse(optimized_model, inplace=True)
             except:  # noqa E722
                 warnings.warn(
                     "Conv BatchNorm folding failed during the optimize process."
                 )
         if opt_properties.linear_bn_folding:
             try:
-                optimized_model = linear_bn_fuse(optimized_model, inplace=inplace)
+                optimized_model = linear_bn_fuse(optimized_model, inplace=True)
             except BaseException:
                 warnings.warn(
                     "Linear BatchNorm folding failed during the optimize process."
                 )
         if opt_properties.replace_dropout_with_identity:
             utils._model_convert.replace_dropout_with_identity(optimized_model)
-        if dtype == torch.bfloat16:
-            optimized_model = utils._model_convert.convert_module_data_type(
-                optimized_model, torch.bfloat16
-            )
-        if dtype == torch.half:
-            optimized_model = utils._model_convert.convert_module_data_type(
-                optimized_model, torch.half
+        if dtype in (
+            torch.bfloat16,
+            torch.float16,
+        ):
+            params_attr, optimized_model = utils._model_convert.convert_model_data_type(
+                optimized_model, dtype
             )
 
     if opt_properties.optimize_lstm:
@@ -654,37 +624,28 @@ def xpu_check_channel_last():
                 + " will use non-fused master weight update for bf16 training on XPU."
             )
 
-    # convert optimizer for training case.
-    params_attr = {}
-    if hasattr(optimized_optimizer, "params_attr"):
-        params_attr = optimized_optimizer.params_attr
-    if dtype == torch.bfloat16 and model.training:
-        (
-            optimized_model,
-            optimized_optimizer,
-            params_attr,
-        ) = utils._weight_cast.weight_dtype_convert_with_ipex(
-            optimized_model,
-            optimized_optimizer,
-            params_attr,
-            opt_properties.split_master_weight_for_bf16,
-            convert_dtype=torch.bfloat16,
-        )
-    if dtype == torch.half and model.training:
-        assert (
-            device_type != "xpu"
-        ), "For now, XPU device does not support model training with half precision."
-        (
-            optimized_model,
-            optimized_optimizer,
-            params_attr,
-        ) = utils._weight_cast.weight_dtype_convert_with_ipex(
-            optimized_model,
-            optimized_optimizer,
-            params_attr,
-            False,
-            convert_dtype=torch.half,
-        )
+    if model.training:
+        if hasattr(optimized_optimizer, "params_attr"):
+            params_attr = optimized_optimizer.params_attr
+        if dtype == torch.float16:
+            assert (
+                device_type != "xpu"
+            ), "For now, XPU device does not support model training with half precision."
+            opt_properties.split_master_weight_for_bf16 = False
+        if dtype in (torch.bfloat16, torch.float16):
+            # convert optimizer for training case.
+            (
+                optimized_model,
+                optimized_optimizer,
+                params_attr,
+            ) = utils._weight_cast.weight_dtype_convert_with_ipex(
+                optimized_model,
+                optimized_optimizer,
+                params_attr,
+                opt_properties.split_master_weight_for_bf16,
+                dtype,
+            )
+
     # Since TorchDynamo cannot handle custom operations yet, for the case of inference graph mode,
     # the weights prepacking here is temporarily cancelled, and it will be completed on the graph.
     if opt_properties.weights_prepack:
@@ -704,7 +665,7 @@ def xpu_check_channel_last():
                 optimized_optimizer,
                 params_attr,
             ) = utils._weight_prepack.weight_prepack_with_ipex(
-                optimized_model, optimized_optimizer, params_attr, inplace, "cpu"
+                optimized_model, optimized_optimizer, params_attr, "cpu"
             )
             torch._dynamo.allow_in_graph(utils._weight_prepack._IPEXConv2d)
             torch._dynamo.allow_in_graph(utils._weight_prepack._IPEXConvTranspose2d)
@@ -719,7 +680,7 @@ def xpu_check_channel_last():
                 optimized_optimizer,
                 params_attr,
             ) = utils._weight_prepack.weight_prepack_with_ipex(
-                optimized_model, optimized_optimizer, params_attr, inplace, "xpu"
+                optimized_model, optimized_optimizer, params_attr, "xpu"
             )
 
     if opt_properties.graph_mode:

intel_extension_for_pytorch/nn/utils/_model_convert.py

@@ -1,10 +1,9 @@
 import torch
 import copy
-import warnings
-import types
-
 from torch.nn.utils.rnn import PackedSequence
-
+from ._parameter_wrapper import get_shared_parameter_status
+import contextlib
+import types
 
 class _LSTM(torch.nn.LSTM):
     # This is a solution to swap the lstm module with the ipex counterpart
@@ -127,91 +126,44 @@ def replace_dropout_with_identity(model):
                 replace_dropout_with_identity(child)
 
 
-def _save_to_state_dict(self, destination, prefix, keep_vars):
-    # convert weights(bias) of module to float while saving check point
-    param_dict = {}
-    for name, para in self.named_parameters():
-        if not hasattr(self, name):
-            continue
-        param_dict.update({name: para})
-        temp_param = torch.nn.Parameter(
-            para.to(torch.float), requires_grad=para.requires_grad
-        )
-        setattr(self, name, temp_param)
-    super(type(self), self)._save_to_state_dict(destination, prefix, keep_vars)
-    for p in param_dict:
-        origin_param = param_dict[p]
-        setattr(self, p, origin_param)
-
-
-def convert_module_data_type(module, dtype):
-    # convert weights(bias) of module to dtype to reduce dtype reorder
+def convert_model_data_type(model, dtype):
+    # convert weights(bias) of model to dtype to reduce dtype reorder
     assert dtype in [
         torch.bfloat16,
         torch.float16,
-    ], "module convert only support bf16 and fp16"
-    module_convert_list_bf16 = [
-        torch.nn.Conv2d,
-        torch.nn.Conv3d,
-        torch.nn.ConvTranspose2d,
-        torch.nn.ConvTranspose3d,
-        torch.nn.Linear,
-        torch.nn.Embedding,
-        torch.nn.LSTM,
-    ]
-
-    module_convert_list_fp16 = [
-        torch.nn.Conv1d,
-        torch.nn.Conv2d,
-        torch.nn.Conv3d,
-        torch.nn.Linear,
-    ]
-
-    module_convert_lists = {
-        torch.bfloat16: module_convert_list_bf16,
-        torch.float16: module_convert_list_fp16,
-    }
-
-    for module_cls in module_convert_lists[dtype]:
-        if isinstance(module, module_cls):
-            setattr(
-                module,
-                "_save_to_state_dict",
-                types.MethodType(_save_to_state_dict, module),
-            )
-            if module_cls is torch.nn.LSTM:
-                for name, param in module.named_parameters():
-                    ori_data = getattr(getattr(module, name), "data")
-                    ori_data_dtype = ori_data.dtype
-                    if (
-                        ori_data_dtype == torch.float
-                        or ori_data_dtype == torch.bfloat16
-                    ):
-                        casted_data = ori_data.detach().clone().to(dtype)
-                        setattr(getattr(module, name), "data", casted_data)
-                    else:
-                        warnings.warn(
-                            f"WARNING: Can't convert model's parameters dtyep from {ori_data_dtype} to {dtype}"
-                        )
-                        break
-            else:
-                ori_data_dtype = module.weight.dtype
-                # Assume weight and bias have same dtype, only need check weight dtype here.
-                if (
-                    ori_data_dtype == torch.float
-                    or ori_data_dtype == torch.bfloat16
-                    or ori_data_dtype == torch.half
-                ):
-                    weight_data = module.weight.detach().clone().to(dtype)
-                    module.weight.data = weight_data
-                    if hasattr(module, "bias") and module.bias is not None:
-                        bias_data = module.bias.detach().clone().to(dtype)
-                        module.bias.data = bias_data
-                else:
-                    warnings.warn(
-                        f"WARNING: Can't convert model's parameters dtype from {ori_data_dtype} to {dtype}"
-                    )
-            break
-    for child in module.children():
-        convert_module_data_type(child, dtype)
-    return module
+    ], "model convert only support bf16 and fp16"
+
+    params_attr = {}
+    get_shared_parameter_status(model, params_attr)
+
+    for _, param in model.named_parameters():
+        if param is None:
+            continue
+        if params_attr[param].can_cast_inference(dtype):
+            params_attr[param].cast_for_inference(dtype)
+
+    def patch_state_dict():
+        def cast_back_state_dict(
+            self, *args, destination=None, prefix="", keep_vars=False
+        ):
+            with torch.no_grad(), contextlib.ExitStack() as stack:
+                for v in params_attr.values():
+                    stack.enter_context(v.inference_cast_save())
+                out = self._original_state_dict(
+                    *args,
+                    destination=destination,
+                    prefix=prefix,
+                    keep_vars=keep_vars
+                )
+            return out
+
+        if not hasattr(model, "_original_state_dict"):
+            setattr(model, "_original_state_dict", model.state_dict)
+        setattr(
+            model,
+            "state_dict",
+            types.MethodType(cast_back_state_dict, model),
+        )
+
+    patch_state_dict()
+    return params_attr, model