mindspore 2.7.0rc1__cp311-cp311-win_amd64.whl → 2.7.1__cp311-cp311-win_amd64.whl

mindspore/nn/cell.py

@@ -39,10 +39,11 @@ from typing import (
 
 import weakref
 import mindspore as ms
+import mindspore.ops as ops
 from mindspore._checkparam import args_type_check, check_hook_fn
-from mindspore.common._auto_dynamic import is_auto_dynamic, convert_inputs_to_dynamic
+from mindspore.common.dynamic_shape._auto_dynamic import is_auto_dynamic, convert_inputs_to_dynamic
 from mindspore import log as logger
-from mindspore.common.hook_handle import HookHandle
+from mindspore.common.hook_handle import HookHandle, _update_hook_version
 from mindspore import context
 from mindspore._c_expression import init_pipeline, update_func_graph_hyper_params, Cell_, FuncGraph, MixedPrecisionType
 from mindspore import _checkparam as Validator
@@ -92,9 +93,8 @@ def register_cell_buffer_registration_hook(hook: Callable[..., None],):
         A handle that can be used to remove the added hook by calling
         `handle.remove()`.
     """
-    from mindspore.utils.hooks import _RemovableHandle
-    handle = _RemovableHandle(_global_buffer_registration_hooks)
-    _global_buffer_registration_hooks[handle.id] = hook
+    handle = HookHandle(_global_buffer_registration_hooks)
+    _global_buffer_registration_hooks[handle.handle_id] = hook
     return handle
 
 
@@ -155,7 +155,8 @@ class Cell(Cell_):
     IGNORE_LIST = ['_scope', '_cell_init_args', '_auto_prefix', '_cells', '_params', '_create_time',
                    '_func_graph_flags', '_parameter_layout_dict', '_params_list', '_phase', '_bprop_debug',
                    '_forward_pre_hook', '_forward_hook', '_backward_pre_hook', '_backward_hook',
-                   '_cell_backward_pre_hook', '_cell_backward_hook', '_param_prefix', 'requires_grad', 'cell_type']
+                   '_cell_backward_pre_hook', '_cell_backward_hook', '_param_prefix',
+                   'requires_grad', 'cell_type', '_in_strategy', '_out_strategy']
     total_instance_count = 0
     _buffers: Dict[str, Optional[Tensor]]
     global_cells = weakref.WeakKeyDictionary()
@@ -191,6 +192,7 @@ class Cell(Cell_):
         super().__setattr__("_auto_prefix", auto_prefix)
         super().__setattr__("_scope", None)
         super().__setattr__("_phase", 'train')
+        super().__setattr__("_compile_phase", None)
         super().__setattr__("_parameter_layout_dict", None)
         super().__setattr__("_parallel_parameter_name_list", None)
         super().__setattr__("_parallel_parameter_merge_net_dict", None)
@@ -206,6 +208,7 @@ class Cell(Cell_):
         super().__setattr__("mixed_precision_type", None)
         super().__setattr__("_lazy_construct_sig", None)
         super().__setattr__("_jit_graph_name", '')
+        super().__setattr__("_compiled", False)
         init_pipeline()
 
         # call gc to release GE session resources used by non-used cell objects
@@ -239,6 +242,8 @@ class Cell(Cell_):
         super().__setattr__("_amp_level", "")
         super().__setattr__("_init_flag", False)
         super().__setattr__("_shard_fn", None)
+        super().__setattr__("_in_strategy", None)
+        super().__setattr__("_out_strategy", None)
         super().__setattr__("has_bprop", False)
         if hasattr(self, "bprop"):
             super().__setattr__("has_bprop", True)
@@ -426,6 +431,13 @@ class Cell(Cell_):
         """
         return self._bprop_debug
 
+    @property
+    def compiled(self):
+        """
+        Get whether `Cell` is compiled in graph mode.
+        """
+        return self._compiled
+
     @bprop_debug.setter
     def bprop_debug(self, value):
         """
@@ -482,6 +494,19 @@ class Cell(Cell_):
             raise TypeError(f"For 'Cell', the property 'phase' must be string type, but got type {type(value)}.")
         self._phase = value
 
+    @property
+    def compile_phase(self):
+        return self._compile_phase
+
+    @compile_phase.setter
+    def compile_phase(self, value):
+        if not isinstance(value, str):
+            raise TypeError(f"For 'Cell', 'compile_phase' must be string type, but got type {type(value)}.")
+        self._compile_phase = value
+        for cell in self._cells.values():
+            if cell is not None:
+                cell.compile_phase = value
+
     @property
     def parameter_layout_dict(self):
         """
@@ -546,10 +571,23 @@ class Cell(Cell_):
 
     @property
     def pipeline_segment(self):
+        """
+        `pipeline_segment` represents the pipeline segment of current Cell.
+        """
         return self._pipeline_segment
 
     @pipeline_segment.setter
     def pipeline_segment(self, value):
+        """
+        Set the `pipeline_segment` of a Cell. Only effective in zero_bubble_v scheduler.
+
+        Args:
+            value (int): The pipeline segment of a parameter.
+
+        Raises:
+            TypeError: If `value` is not int type or is a bool type.
+            ValueError: If `value` is not a positive integer.
+        """
         if not isinstance(value, int) or isinstance(value, bool):
             raise TypeError("For 'context.set_auto_parallel_context', the argument 'pipeline_stages' "
                             "must be int type, but got type : {}".format(type(value)))
@@ -1027,12 +1065,13 @@ class Cell(Cell_):
         if self._forward_pre_hook:
             args, kwargs = self._run_forward_pre_hook(args, kwargs)
 
+        if self._backward_hook:
+            args = self._cell_backward_hook(args)
+
         if self._shard_fn is not None:
             output = self._shard_fn(*args, **kwargs)
         elif _pynative_executor.requires_grad():
-            if self._backward_hook:
-                output = self._backward_hook_construct(*args, **kwargs)
-            elif self._recompute_cell is not None:
+            if self._recompute_cell is not None:
                 output = self._recompute_cell(*args, **kwargs)
             elif self.has_bprop:
                 output = self._call_custom_bprop(*args, **kwargs)
@@ -1044,8 +1083,11 @@ class Cell(Cell_):
         if self._forward_hook:
             output = self._run_forward_hook(args, kwargs, output)
 
-        if self._backward_pre_hook and _pynative_executor.requires_grad():
-            output = self._run_backward_pre_hook(output)
+        if self._backward_hook:
+            output = self._cell_backward_hook(output)
+
+        if self._backward_pre_hook:
+            output = self._cell_backward_pre_hook(output)
 
         return output
 
@@ -1080,23 +1122,6 @@ class Cell(Cell_):
                             f"{default_args} default argument, total {positional_args + default_args}, "
                             f"but got {len(args)}.")
 
-    # pylint: disable=E0203
-    def _hook_fn_registered(self):
-        '''Hook function in graph mode'''
-        # Check super().__init__() in graph mode.
-        try:
-            if self._forward_pre_hook or self._forward_hook or self._backward_pre_hook or self._backward_hook:
-                return True
-        except AttributeError as e:
-            raise AttributeError(f"The '{type(self).__name__}' object does not inherit attribute from 'cell'. "
-                                 f"Please use 'super().__init__()'.") from e
-        if not self._is_recursion_hook:
-            self._is_recursion_hook = True
-            for cell in self.cells():
-                if cell._hook_fn_registered():
-                    return True
-        return False
-
     def _get_prims_recursively(self):
         all_prims = list()
         for _, value in self._primitives.items():
@@ -1122,9 +1147,6 @@ class Cell(Cell_):
             >>> net = nn.Dense(3, 4)
             >>> net.set_data_parallel()
         """
-        if context._get_mode() == context.PYNATIVE_MODE:
-            raise ValueError("set_data_parallel: does not support PyNative mode.")
-
         all_prims = self._get_prims_recursively()
         for prim in all_prims:
             prim.add_prim_attr("strategy_gen_mode", "data_parallel")
@@ -1203,8 +1225,6 @@ class Cell(Cell_):
             ...             out = self.blocks[i](out)
             ...         return out
         """
-        if context._get_mode() == context.PYNATIVE_MODE:
-            raise ValueError("The Cell offload does not support PyNative mode now.")
         if isinstance(backward_prefetch, str):
             Validator.check_string(backward_prefetch, ['Auto'], 'backward_prefetch', self.cls_name)
         else:
@@ -1212,11 +1232,10 @@ class Cell(Cell_):
         for prim in self._get_prims_recursively():
             prim._offload(backward_prefetch=backward_prefetch)
 
-    def shard(self, in_strategy, out_strategy=None, parameter_plan=None, device="Ascend", level=0):
+    def shard(self, in_strategy, out_strategy=None, parameter_plan=None):
         """
         Defining the input and output layouts of this cell and the parallel strategies of remaining ops will be
-        generated by sharding propagation. In PyNative mode, use this method to specify a Cell for distributed
-        execution in graph mode. In Graph mode, use this method to specify distribution strategy for a Cell,
+        generated by sharding propagation. In Graph mode, use this method to specify distribution strategy for a Cell,
         strategy for others will be set by sharding propagation.
         in_strategy and out_strategy define the input and output layout respectively.
         in_strategy/out_strategy should be a tuple, each element of which corresponds to the desired layout of
@@ -1228,11 +1247,14 @@ class Cell(Cell_):
               In other parallel modes, strategies set here will be ignored.
             - If the input contain Parameter, its strategy should be set in `in_strategy`.
 
+        .. warning::
+            The method is currently not supported in PyNative mode.
+
         Args:
             in_strategy (tuple): Define the layout of inputs, each element of the tuple should be a tuple. Tuple
                                  defines the layout of the corresponding input.
             out_strategy (Union[None, tuple]): Define the layout of outputs similar with in_strategy.
-                                               It is not in use right now. Default: ``None`` .
+                                               Default: ``None`` .
             parameter_plan (Union[dict, None]): Define the layout for the specified parameters. Each element in dict
                                                 defines the layout of the parameter like "param_name: layout".
                                                 The key is a parameter name of type 'str'.
@@ -1240,14 +1262,6 @@ class Cell(Cell_):
                                                 If the parameter name is incorrect or the corresponding parameter
                                                 has been set, the parameter setting will be ignored.
                                                 Default: ``None`` .
-            device (str): Select a certain device target. It is not in use right now.
-                             Support [ ``"CPU"`` , ``"GPU"`` , ``"Ascend"`` ]. Default: ``"Ascend"`` .
-            level (int): Option for parallel strategy infer algorithm, namely the object function, maximize computation
-                         over communication ratio, maximize speed performance, minimize memory usage etc. It is not in
-                         use right now. Support [ ``"0"`` , ``"1"`` , ``"2"`` ]. Default: ``0`` .
-
-        Returns:
-            Function, return the cell construct function that will be executed under auto parallel process.
 
         Examples:
             >>> import mindspore.nn as nn
@@ -1265,19 +1279,34 @@ class Cell(Cell_):
             ...   def __init__(self):
             ...     self.block1 = Block()
             ...     self.block2 = Block()
-            ...     self.block2_shard = self.block2.shard(in_strategy=((2, 1),),
-            ...                                           parameter_plan={'self.block2.shard.dense1.weight': (4, 1)})
+            ...     self.block2.shard(in_strategy=((2, 1),), parameter_plan={'self.block2.dense1.weight': (4, 1)})
             ...   def construct(self, x):
             ...     x = self.block1(x)
-            ...     x = self.block2_shard(x)
+            ...     x = self.block2(x)
             ...     return x
         """
         if ms.communication.management.get_group_size() == 1:
-            return self
+            return
+
         shard_fn = Shard()
-        fn = shard_fn(self, in_strategy, out_strategy, parameter_plan, device, level)
-        self._shard_fn = fn
-        return fn
+        self._shard_fn = shard_fn(self, in_strategy, out_strategy, parameter_plan)
+
+        if self._in_strategy is not None:  # pylint: disable=E0203
+            msg = (
+                "For '%s', 'Shard' has been configured more than once. "
+                "The existing in_strategy is %s and the existing out_strategy is %s. "
+                "The new in_strategy %s and out_strategy %s may not take effect. "
+                "It is recommended to configure 'Shard' only once."
+            ) % (
+                self._cell_tag,
+                self._in_strategy,  # pylint: disable=E0203
+                self._out_strategy,  # pylint: disable=E0203
+                shard_fn.in_strategy,
+                shard_fn.out_strategy,
+            )
+            logger.warning(msg)
+        self._in_strategy = shard_fn.in_strategy
+        self._out_strategy = shard_fn.out_strategy
 
     def _init_check(self):
         for param in self.get_parameters(expand=False):
@@ -1286,9 +1315,13 @@ class Cell(Cell_):
         self._init_flag = True
 
     def _self_check(self):
-        if not self._is_check_and_refresh:
-            self.check_names_and_refresh_name()
-            self._is_check_and_refresh = True
+        try:
+            if not self._is_check_and_refresh:  # pylint: disable=E0203
+                self.check_names_and_refresh_name()
+                self._is_check_and_refresh = True
+        except AttributeError as e:
+            raise AttributeError(f"The '{type(self).__name__}' object does not inherit attribute from 'cell'. "
+                                 f"Please use 'super().__init__()'.") from e
 
     def _predict(self, *args, **kwargs):
         '''Graph executor for predict'''
@@ -1309,6 +1342,7 @@ class Cell(Cell_):
     def __call__(self, *args, **kwargs):
         # Run in Graph mode.
         if context._get_mode() == context.GRAPH_MODE and os.getenv("MS_JIT") != '0':
+            self._compiled = True
             if kwargs:
                 bound_arguments = self._construct_sig.bind(*args, **kwargs)
                 bound_arguments.apply_defaults()
@@ -1319,11 +1353,8 @@ class Cell(Cell_):
             if predict_compiled:
                 return res
             self._check_construct_args(*args)
-
-            if self._hook_fn_registered():
-                logger.warning(f"For 'Cell', it's not support hook function in graph mode. If you want to use hook "
-                               f"function, please use context.set_context to set pynative mode.")
             self._self_check()
+            self.__compile_cell_hook__ = True
             out = self.compile_and_run(*args, **kwargs)
             return out
 
@@ -1421,16 +1452,7 @@ class Cell(Cell_):
                 exist_names.add(item.name)
                 self.insert_param_to_cell(item.name, item, check_name_contain_dot=False)
 
-            if context._get_mode() == context.PYNATIVE_MODE:
-                if name in self.__dict__:
-                    del self.__dict__[name]
-                params = self.__dict__.get('_params')
-                if name in params:
-                    del params[name]
-                params_list = self.__dict__.get('_params_list')
-                params_list[name] = value
-            else:
-                object.__setattr__(self, name, value)
+            object.__setattr__(self, name, value)
 
     def _set_attr_for_parameter_in_list_or_tuple(self, name, value):
         """Set attr for parameter in list or tuple."""
@@ -1609,8 +1631,6 @@ class Cell(Cell_):
                 _pynative_executor.set_dynamic_input(self, *self._dynamic_shape_inputs)
             else:
                 self._check_construct_args(*inputs)
-                # TODO(tronzhang): It may error for no actually args here. So just set in fullmode,
-                #                  which means that incremental mode is lacking dynamic input.
         else:
             self._dynamic_shape_inputs = _process_dyn_args(self.construct, kwargs)
 
@@ -1699,6 +1719,7 @@ class Cell(Cell_):
         _init_auto_parallel_context(self)
         compile_args = self._get_compile_args(args)
         self._has_mutable_args_list = _get_mutable_flags(compile_args)
+        _cell_graph_executor.set_real_args(args, kwargs)
         _cell_graph_executor.compile(self, *compile_args, phase=self.phase,
                                      jit_config_dict=self._jit_config_dict, **kwargs)
         _clear_auto_parallel_context(self)
@@ -2581,23 +2602,7 @@ class Cell(Cell_):
         else:
             self._jit_config_dict = jit_config.jit_config_dict
 
-    def flatten_weights(self, fusion_size=0):
-        """
-        Reset data for weight parameters so that they are using contiguous memory chunks grouped by data type.
-
-        Note:
-            By default, parameters with same data type will using a single contiguous memory chunk. but for
-            some models with huge number of parameters, splitting a large memory chunk into several smaller
-            memory chunks has the potential for performance gains, if this is the case, we can use 'fusion_size'
-            to limit the maximum memory chunk size.
-
-        Args:
-            fusion_size (int): Maximum memory chunk size in bytes, ``0`` for unlimited. Default: ``0`` .
-        """
-        if fusion_size < 0:
-            raise ValueError(f"Negative 'fusion_size' {fusion_size} is invalid.")
-        Tensor._flatten_tensors(self.trainable_params(), fusion_size)  # pylint: disable=W0212
-
+    @jit_forbidden_register
     def register_forward_pre_hook(self, hook_fn, with_kwargs=False):
         """
         Register forward pre hook function for Cell object.
@@ -2617,7 +2622,6 @@ class Cell(Cell_):
           `with_kwargs` is ``True`` .
 
         Note:
-            - The feature does not take effect in graph mode or in PyNative mode with functions decorated by jit.
             - The `hook_fn` can modify the forward inputs by returning new inputs. If `with_kwargs` is ``Flase`` , a
               single value (whick will be wrapped into a tuple unless already a tuple) or a tuple of args should be
               returned. If `with_kwargs` is ``True`` , both `args` and `kwargs` should be returned.
@@ -2668,15 +2672,15 @@ class Cell(Cell_):
             (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], dtype=Float32,
             value= [ 2.00000000e+00]))
         """
-        if context._get_mode() == context.GRAPH_MODE:
-            return HookHandle()
         check_hook_fn(hook_fn)
         handle = HookHandle(self._forward_pre_hook, extra_dict=self._forward_pre_hook_with_kwargs)
         self._forward_pre_hook[handle.handle_id] = hook_fn
         if with_kwargs:
             self._forward_pre_hook_with_kwargs[handle.handle_id] = True
+        _update_hook_version()
         return handle
 
+    @jit_forbidden_register
     def _run_forward_pre_hook(self, args, kwargs):
         """
         Running forward pre hook function registered on Cell object.
@@ -2700,6 +2704,35 @@ class Cell(Cell_):
                     args = ret
         return args, kwargs
 
+    def _jit_forward_pre_hook(self, inputs):
+        """
+        Compile forward pre hook function registered on Cell object.
+
+        Args:
+            inputs: The input objects of cell object.
+
+        Returns:
+            - **outputs** - New input objects or none.
+
+        Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+        """
+        forward_pre_hook_inputs = inputs
+        for fn in self._forward_pre_hook.values():
+            ret = fn(self, forward_pre_hook_inputs)
+            if ret is not None:
+                if not isinstance(ret, tuple):
+                    forward_pre_hook_inputs = (ret,)
+                else:
+                    forward_pre_hook_inputs = ret
+
+        if len(forward_pre_hook_inputs) != len(inputs):
+            raise TypeError(
+                "The forward pre hook return value size is {} not equal to input size {}".format(
+                    len(forward_pre_hook_inputs), len(inputs)))
+        return forward_pre_hook_inputs
+
+    @jit_forbidden_register
     def register_forward_hook(self, hook_fn, with_kwargs=False):
         """
         Register forward hook function for Cell object.
@@ -2720,7 +2753,6 @@ class Cell(Cell_):
         - `output`: Output generated by the `construct` function.
 
         Note:
-            - The feature does not take effect in graph mode or in PyNative mode with functions decorated by jit.
             - The `hook_fn` can modify the forward outputs by returning new outputs.
             - In order to prevent running failed when switching to graph mode, it is not recommended to call it in the
               `construct` function of Cell object.
@@ -2773,15 +2805,44 @@ class Cell(Cell_):
         """
         if self.has_bprop:
             return HookHandle()
-        if context._get_mode() == context.GRAPH_MODE:
-            return HookHandle()
         check_hook_fn(hook_fn)
         handle = HookHandle(self._forward_hook, extra_dict=self._forward_hook_with_kwargs)
         self._forward_hook[handle.handle_id] = hook_fn
         if with_kwargs:
             self._forward_hook_with_kwargs[handle.handle_id] = True
+        _update_hook_version()
         return handle
 
+    def _jit_forward_hook(self, inputs, output):
+        """
+        Compile forward hook function registered on Cell object.
+
+        Args:
+            inputs: The input objects of Cell object.
+            output: The output object of Cell object.
+
+        Returns:
+            - **output** - New output object or none.
+
+        Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+        """
+        forward_hook_output = output
+        for fn in self._forward_hook.values():
+            ret = fn(self, inputs, forward_hook_output)
+            if ret is not None:
+                forward_hook_output = ret
+
+        if isinstance(output, tuple):
+            if not isinstance(forward_hook_output, tuple):
+                forward_hook_output = (forward_hook_output,)
+            if len(forward_hook_output) != len(output):
+                raise TypeError(
+                    "The forward hook return value size is {} not equal to output size {}".format(
+                        len(forward_hook_output), len(output)))
+        return forward_hook_output
+
+    @jit_forbidden_register
     def _run_forward_hook(self, args, kwargs, output):
         """
         Running forward hook function registered on Cell object.
@@ -2795,12 +2856,12 @@ class Cell(Cell_):
                 output = ret
         return output
 
+    @jit_forbidden_register
     def register_backward_pre_hook(self, hook_fn):
         """
         Register the backward pre hook function.
 
         Note:
-            - The `register_backward_pre_hook(hook_fn)` does not work in graph mode or functions decorated with 'jit'.
             - The 'hook_fn' must be defined as the following code.
               `cell` is the Cell object. `grad_output` is the gradient passed to the Cell.
             - The 'hook_fn' should have the following signature:
@@ -2849,44 +2910,17 @@ class Cell(Cell_):
             >>> print(output)
             (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]),)
         """
-        if context._get_mode() == context.GRAPH_MODE:
-            return HookHandle()
         check_hook_fn(hook_fn)
-        handle = HookHandle(self._backward_pre_hook)
+        handle = HookHandle(self._backward_pre_hook, extra_dict=None)
         self._backward_pre_hook[handle.handle_id] = hook_fn
-        if self._cell_backward_pre_hook is None:
+        if self._cell_backward_pre_hook is None:  # pylint: disable=E0203
             # Generate a CellBackwardHook prim, and add function for it
             self._cell_backward_pre_hook = inner.CellBackwardHook(self.cls_name + "(" + str(id(self)) + ")",
                                                                   self, self._backward_pre_hook)
             self._cell_backward_pre_hook.register_backward_pre_hook()
+        _update_hook_version()
         return handle
 
-    def _run_backward_pre_hook(self, outputs):
-        """
-        Running backward pre hook function registered on Cell object.
-
-        Args:
-            outputs: The output objects of cell object.
-
-        Returns:
-            - **outputs** - New backward gradient or None.
-
-        Supported Platforms:
-        ``Ascend`` ``GPU`` ``CPU``
-        """
-        if isinstance(outputs, tuple):
-            ret = self._cell_backward_pre_hook(*outputs)
-        else:
-            ret = self._cell_backward_pre_hook(outputs)
-        if isinstance(outputs, tuple):
-            if len(outputs) == 1:
-                ret = (ret,)
-            if len(ret) != len(outputs):
-                raise TypeError(
-                    "The backward pre hook return value size is {} not equal to output size {}".format(
-                        len(ret), len(outputs)))
-        return ret
-
     def get_extra_state(self) -> Any:
         """Return any extra state to include in the cell's state_dict.
 
@@ -2939,9 +2973,8 @@ class Cell(Cell_):
             A handle that can be used to remove the added hook by calling
             `handle.remove()`.
         """
-        from mindspore.utils.hooks import _RemovableHandle
-        handle = _RemovableHandle(self._state_dict_hooks)
-        self._state_dict_hooks[handle.id] = hook
+        handle = HookHandle(self._state_dict_hooks)
+        self._state_dict_hooks[handle.handle_id] = hook
         return handle
 
     @jit_forbidden_register
@@ -2987,9 +3020,8 @@ class Cell(Cell_):
             >>> print("extra_param" in net_state_dict)
             True
         """
-        from mindspore.utils.hooks import _RemovableHandle
-        handle = _RemovableHandle(self._state_dict_pre_hooks)
-        self._state_dict_pre_hooks[handle.id] = hook
+        handle = HookHandle(self._state_dict_pre_hooks)
+        self._state_dict_pre_hooks[handle.handle_id] = hook
         return handle
 
     def _save_to_state_dict(self, destination, prefix, keep_vars):
@@ -3135,9 +3167,8 @@ class Cell(Cell_):
             A handle that can be used to remove the added hook by calling
             `handle.remove()`.
         """
-        from mindspore.utils.hooks import _RemovableHandle
-        handle = _RemovableHandle(self._load_state_dict_pre_hooks)
-        self._load_state_dict_pre_hooks[handle.id] = hook
+        handle = HookHandle(self._load_state_dict_pre_hooks)
+        self._load_state_dict_pre_hooks[handle.handle_id] = hook
         return handle
 
     @jit_forbidden_register
@@ -3169,9 +3200,8 @@ class Cell(Cell_):
             A handle that can be used to remove the added hook by calling
             `handle.remove()`.
         """
-        from mindspore.utils.hooks import _RemovableHandle
-        handle = _RemovableHandle(self._load_state_dict_post_hooks)
-        self._load_state_dict_post_hooks[handle.id] = hook
+        handle = HookHandle(self._load_state_dict_post_hooks)
+        self._load_state_dict_post_hooks[handle.handle_id] = hook
         return handle
 
     def _load_from_state_dict(
@@ -3407,12 +3437,12 @@ class Cell(Cell_):
             )
         return _IncompatibleKeys(missing_keys, unexpected_keys)
 
+    @jit_forbidden_register
     def register_backward_hook(self, hook_fn):
         """
         Register the backward hook function.
 
         Note:
-            - The `register_backward_hook(hook_fn)` does not work in graph mode or functions decorated with 'jit'.
             - The 'hook_fn' must be defined as the following code.
               `cell` is the registered Cell object. `grad_input` is the gradient computed and passed to
               the next Cell or primitive, which can be return a new gradient or None. `grad_output` is the gradient
@@ -3464,83 +3494,17 @@ class Cell(Cell_):
             >>> print(output)
             (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]),)
         """
-        if context._get_mode() == context.GRAPH_MODE:
-            return HookHandle()
         check_hook_fn(hook_fn)
-        handle = HookHandle(self._backward_hook)
+        handle = HookHandle(self._backward_hook, extra_dict=None)
         self._backward_hook[handle.handle_id] = hook_fn
-        if self._cell_backward_hook is None:
+        if self._cell_backward_hook is None:  # pylint: disable=E0203
             # Generate a CellBackwardHook prim, and add function for it
             self._cell_backward_hook = inner.CellBackwardHook(self.cls_name + "(" + str(id(self)) + ")",
                                                               self, self._backward_hook)
             self._cell_backward_hook.register_backward_hook()
+        _update_hook_version()
         return handle
 
-    def _backward_hook_construct(self, *inputs, **kwargs):
-        """
-        Backward hook construct method to replace original construct method.
-
-        Args:
-            inputs: The input objects of Cell object.
-            kwargs (dict): Dictionary of variable keyword parameters.
-
-        Returns:
-            - **outputs** - The output objects of Cell object.
-
-        Supported Platforms:
-        ``Ascend`` ``GPU`` ``CPU``
-        """
-        # cell_backward_hook has CellBackwardHook op, so keep input args as they are.
-        outputs = self._cell_backward_hook(*inputs)
-        # If the inputs have more than two args, the outputs will also have more than two args and will be wrapped into
-        # a tuple, so need to do unwrapping. If inputs is empty, we also need to unwrap it.
-        # Because when output of runop method is one, it will not wrap a tuple, we need not unwrap it.
-        is_need_unwrap = False
-        if isinstance(outputs, tuple) and len(inputs) != 1:
-            is_need_unwrap = True
-
-        if self._recompute_cell is not None:
-            if is_need_unwrap:
-                outputs = self._recompute_cell(*outputs, **kwargs)
-            else:
-                outputs = self._recompute_cell(outputs, **kwargs)
-        elif self.has_bprop:
-            if is_need_unwrap:
-                outputs = self._call_custom_bprop(*outputs, **kwargs)
-            else:
-                outputs = self._call_custom_bprop(outputs, **kwargs)
-        else:
-            if is_need_unwrap:
-                outputs = self.construct(*outputs, **kwargs)
-            else:
-                outputs = self.construct(outputs, **kwargs)
-        if isinstance(outputs, tuple):
-            new_outputs = self._cell_backward_hook(*outputs)
-        else:
-            new_outputs = self._cell_backward_hook(outputs)
-        # if outputs is (X,) and new_outpus is X
-        if isinstance(outputs, tuple) and len(outputs) == 1:
-            new_outputs = (new_outputs,)
-        return new_outputs
-
-    def set_param_ps(self, recurse=True, init_in_server=False):
-        """
-        Set whether the trainable parameters are updated by parameter server and whether the
-        trainable parameters are initialized on server.
-
-        Note:
-            It only works when a running task is in the parameter server mode.
-            It is only supported in graph mode.
-
-        Args:
-            recurse (bool): Whether sets the trainable parameters of subcells. Default: ``True`` .
-            init_in_server (bool): Whether trainable parameters updated by parameter server are
-                initialized on server. Default: ``False`` .
-        """
-        params = self.trainable_params(recurse)
-        for param in params:
-            param.set_param_ps(init_in_server)
-
     def set_comm_fusion(self, fusion_type, recurse=True):
         """
         Set `comm_fusion` for all the parameters in this cell. Please refer to the description of
@@ -3601,7 +3565,7 @@ class Cell(Cell_):
         """
         Validator.check_bool(mode)
         Validator.check_bool(output_recompute)
-        if not self._has_config_recompute:
+        if not self._has_config_recompute:  # pylint: disable=E0203
             self._has_config_recompute = True
         else:
             logger.info("The recompute interface can be configured only once."
@@ -3644,8 +3608,7 @@ class Cell(Cell_):
                 introduced by optimizer shard are recomputed in auto parallel or semi auto parallel mode.
                 Default: ``False`` .
         """
-        if context.get_context("mode") == context.PYNATIVE_MODE:
-            self._recompute_cell = recompute_registry.get()(self.construct)
+        self._recompute_cell = recompute_registry.get()(self.construct)
         self._recompute()
         if 'mp_comm_recompute' in kwargs.keys():
             self._mp_comm_recompute(kwargs.get('mp_comm_recompute', False))
@@ -3662,35 +3625,6 @@ class Cell(Cell_):
                                  "the key kwargs must be 'mp_comm_recompute', "
                                  "'parallel_optimizer_comm_recompute', 'recompute_slice_activation'" % key)
 
-    def place(self, role, rank_id):
-        """
-        Set the label for all operators in this cell.
-        This label tells MindSpore compiler on which process this cell should be launched.
-        And each process's identical label consists of input `role` and `rank_id`.
-        So by setting different cells with different labels, which will be launched on different processes,
-        users can launch a distributed training or predicting job.
-
-        Note:
-            - This method is effective only after
-              `mindspore.communication.init()` is called for dynamic cluster building.
-
-        Args:
-            role (str): The role of the process on which this cell will be launched.
-                        Only 'MS_WORKER' is supported for now.
-            rank_id (int): The rank id of the process on which this cell will be launched.
-                           The rank is unique in processes with the same role.
-
-        Examples:
-            >>> from mindspore import context
-            >>> import mindspore.nn as nn
-            >>> context.set_context(mode=context.GRAPH_MODE)
-            >>> fc = nn.Dense(2, 3)
-            >>> fc.place('MS_WORKER', 0)
-        """
-        all_ops = self._get_prims_recursively()
-        for op in all_ops:
-            op.place(role, rank_id)
-
     def _get_attr_from_cell(self, network):
         if not isinstance(network, Cell):
             return
@@ -3705,6 +3639,64 @@ class Cell(Cell_):
         """
         self._jit_graph_name = key
 
+    def _jit_backward_pre_hook(self, grad_output):
+        new_grad_output = grad_output
+        if not isinstance(grad_output, tuple):
+            new_grad_output = (grad_output,)
+
+        for fn in self._backward_pre_hook.values():
+            ret = fn(self, new_grad_output)
+            if ret is not None:
+                if not isinstance(ret, tuple):
+                    output = (ret,)
+                else:
+                    output = ret
+            else:
+                output = ops.Depend()(new_grad_output, ret)
+            new_grad_output = output
+
+        if not isinstance(grad_output, tuple):
+            if len(new_grad_output) == 1:
+                return new_grad_output[0]
+            raise TypeError(
+                "The backward pre hook return value size is {} not equal to input size 1".format(
+                    len(new_grad_output)))
+
+        if len(new_grad_output) != len(grad_output):
+            raise TypeError(
+                "The backward pre hook return value size is {} not equal to input size {}".format(
+                    len(new_grad_output), len(grad_output)))
+
+        return new_grad_output
+
+    def _jit_backward_hook(self, grad_input, grad_output):
+        backward_hook_input = grad_input
+        backward_hook_output = grad_output
+        if not isinstance(grad_input, tuple):
+            backward_hook_input = (grad_input,)
+        if not isinstance(grad_output, tuple):
+            backward_hook_output = (grad_output,)
+
+        for fn in self._backward_hook.values():
+            ret = fn(self, backward_hook_input, backward_hook_output)
+            if ret is not None:
+                if not isinstance(ret, tuple):
+                    output = (ret,)
+                else:
+                    output = ret
+            else:
+                output = ops.Depend()(backward_hook_input, ret)
+
+            backward_hook_input = output
+
+        if not isinstance(grad_input, tuple):
+            return backward_hook_input[0]
+
+        if len(backward_hook_input) != len(grad_input):
+            raise TypeError(
+                "The backward hook return value size is {} not equal to input size {}".format(
+                    len(backward_hook_input), len(grad_input)))
+        return backward_hook_input
 
 class GraphCell(Cell):
     """