mindspore 2.6.0__cp311-none-any.whl → 2.7.0__cp311-none-any.whl

mindspore/nn/cell.py

@@ -15,6 +15,10 @@
 """cell"""
 from __future__ import absolute_import
 
+__all__ = [
+    "register_cell_buffer_registration_hook",
+]
+
 import inspect
 import os
 import time
@@ -24,7 +28,6 @@ from collections import OrderedDict, namedtuple
 from typing import (
     Dict,
     Optional,
-    Set,
     Callable,
     List,
     Tuple,
@@ -34,36 +37,30 @@ from typing import (
     Mapping
 )
 
+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.parameter import PARAMETER_NAME_DEFAULT
-from mindspore.common.hook_handle import HookHandle
-from mindspore.context import ParallelMode
+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
 from mindspore.common import dtype as mstype
 from mindspore.common.api import _cell_graph_executor, _pynative_executor, _get_args_for_run, cells_compile_cache, \
-    _no_grad
-from mindspore.common.api import _convert_python_data, _get_args_for_run_predict
+    _no_grad, _get_mutable_flags
+from mindspore.common.api import _convert_python_data
 from mindspore.common.api import _process_dyn_args, _generate_dyn_compile_args
-from mindspore.common.parameter import _Buffer, Parameter, ParameterTuple
+from mindspore.common.parameter import _Buffer, Parameter, ParameterTuple, _is_parameter_generated
 from mindspore.common.tensor import Tensor
-from mindspore.ops.operations import Cast
 from mindspore.ops.primitive import Primitive
 from mindspore.ops.operations import _inner_ops as inner
 from mindspore.parallel.shard import Shard
 from mindspore.parallel._utils import _init_auto_parallel_context, _clear_auto_parallel_context
 from mindspore._check_jit_forbidden_api import jit_forbidden_register
-from mindspore.common._decorator import deprecated
 from mindspore.common._register_for_recompute import recompute_registry
-
-
-__all__ = [
-    "register_cell_buffer_registration_hook",
-]
+from mindspore.common.jit_config import JitConfig
 
 _global_buffer_registration_hooks: Dict[int, Callable] = OrderedDict()
 _EXTRA_STATE_KEY_SUFFIX = "_extra_state"
@@ -96,13 +93,11 @@ 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
 
 
-
 class Cell(Cell_):
     """
     The basic building block of neural networks in MindSpore. The model or neural network layer should inherit this
@@ -160,51 +155,59 @@ 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', '_is_run', '_param_prefix',
-                   '_attr_synced', 'pynative', 'requires_grad', 'cell_type',
-                   '_parameters_forward_hook', '_parameters_backward_hook']
+                   '_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]]
-    _non_persistent_buffers_set: Set[str]
+    global_cells = weakref.WeakKeyDictionary()
+    _no_auto_lazy_inline = True
+
+    def __new__(class_, *args, **kwargs):
+        # Use class_ to avoid name conflicts with input args and kwargs.
+        this = Cell_.__new__(class_, *args, **kwargs)
+        if Cell._no_auto_lazy_inline:
+            return this
+
+        Cell.global_cells[this] = (class_, args, kwargs)
+        return this
 
     def __init__(self, auto_prefix=True, flags=None):
         Cell_.__init__(self, self._cell_tag)
         Cell.total_instance_count += 1
-        self.instance_count = Cell.total_instance_count
-        self._params = OrderedDict()
-        self._cells = OrderedDict()
+        super().__setattr__("_params", OrderedDict())
+        super().__setattr__("_cells", OrderedDict())
         super().__setattr__("_buffers", {})
-        super().__setattr__("_non_persistent_buffers_set", set())
-        super().__setattr__("_state_dict_hooks", OrderedDict())
-        super().__setattr__("_state_dict_pre_hooks", OrderedDict())
-        super().__setattr__("_load_state_dict_pre_hooks", OrderedDict())
-        super().__setattr__("_load_state_dict_post_hooks", OrderedDict())
-        self._params_list = OrderedDict()
-        self._primitives = OrderedDict()
-        self.training = False
-        self.requires_grad = False
-        self.is_top_cell = False
-        self.pynative = False
-        self._attr_synced = False
-        self._param_prefix = ''
-        self._auto_prefix = auto_prefix
-        self._scope = None
-        self._phase = 'train'
-        self._parameter_layout_dict = {}
-        self._parallel_parameter_name_list = ()
-        self._parallel_parameter_merge_net_dict = {}
-        self._create_time = int(time.time() * 1e9)
-        self.arguments_key = ""
-        self.compile_cache = set()
-        self.phase_cache = dict()
+        super().__setattr__("_params_list", OrderedDict())
+        super().__setattr__("_primitives", OrderedDict())
+
+        super().__setattr__("_lazy_non_persistent_buffers_set", None)
+        super().__setattr__("_lazy_state_dict_hooks", None)
+        super().__setattr__("_lazy_state_dict_pre_hooks", None)
+        super().__setattr__("_lazy_load_state_dict_pre_hooks", None)
+        super().__setattr__("_lazy_load_state_dict_post_hooks", None)
+        super().__setattr__("training", False)
+        super().__setattr__("requires_grad", False)
+        super().__setattr__("is_top_cell", False)
+        super().__setattr__("_param_prefix", '')
+        super().__setattr__("_auto_prefix", auto_prefix)
+        super().__setattr__("_scope", None)
+        super().__setattr__("_phase", 'train')
+        super().__setattr__("_parameter_layout_dict", None)
+        super().__setattr__("_parallel_parameter_name_list", None)
+        super().__setattr__("_parallel_parameter_merge_net_dict", None)
+        super().__setattr__("_create_time", int(time.time() * 1e9))
+        super().__setattr__("arguments_key", "")
+        super().__setattr__("_compile_cache", None)
+        super().__setattr__("_phase_cache", None)
         cells_compile_cache[id(self)] = self.compile_cache
-        self.parameter_broadcast_done = False
-        self._id = 1
-        self._exist_objs = None
-        self._exist_names = None
-        self._recompute_cell = None
-        self.mixed_precision_type = None
-        self.sig = inspect.signature(self.construct)
+        super().__setattr__("_id", 1)
+        super().__setattr__("_exist_objs", None)
+        super().__setattr__("_exist_names", None)
+        super().__setattr__("_recompute_cell", None)
+        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
@@ -214,38 +217,35 @@ class Cell(Cell_):
 
         if flags:
             self.add_flags(**flags)
-        self._bprop_debug = False
+        super().__setattr__("_bprop_debug", False)
 
         # hook
-        self._forward_pre_hook = OrderedDict()
-        self._forward_hook = OrderedDict()
-        self._backward_pre_hook = OrderedDict()
-        self._cell_backward_pre_hook = None
-        self._backward_hook = OrderedDict()
-        self._cell_backward_hook = None
-        self._is_recursion_hook = False
-
-        # parameters hook
-        self._parameters_forward_hook = None
-        self._parameters_backward_hook = None
-
-        self.cell_type = None
-        self.cast = Cast()
-        self._has_config_recompute = False
-        self._user_parameters = []
-        self._dynamic_shape_inputs = None
-        self._compile_args = None
-        self.saved_dynamic_shape = None
-        self._jit_config_dict = dict()
-        self.grad_ops_label = False
-        self.ge_sync_data = False
-        self._is_check_and_refresh = False
-        self._amp_level = ""
-        self._init_flag = False
-        self._shard_fn = None
-        self.has_bprop = False
+        super().__setattr__("_lazy_forward_pre_hook", None)
+        super().__setattr__("_lazy_forward_hook", None)
+        super().__setattr__("_lazy_backward_pre_hook", None)
+        super().__setattr__("_lazy_backward_hook", None)
+        super().__setattr__("_lazy_forward_pre_hook_with_kwargs", None)
+        super().__setattr__("_lazy_forward_hook_with_kwargs", None)
+        super().__setattr__("_cell_backward_pre_hook", None)
+        super().__setattr__("_cell_backward_hook", None)
+        super().__setattr__("_is_recursion_hook", False)
+
+        super().__setattr__("cell_type", None)
+        super().__setattr__("_has_config_recompute", False)
+        super().__setattr__("_lazy_user_parameters", None)
+        super().__setattr__("_dynamic_shape_inputs", None)
+        super().__setattr__("_has_mutable_args_list", None)
+        super().__setattr__("_jit_config_dict", dict())
+        super().__setattr__("grad_ops_label", False)
+        super().__setattr__("_is_check_and_refresh", False)
+        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"):
-            self.has_bprop = True
+            super().__setattr__("has_bprop", True)
 
     def __getstate__(self):
         base = Cell_.__getstate__(self)
@@ -255,7 +255,6 @@ class Cell(Cell_):
         base, dict_ = state
         Cell_.__setstate__(self, base)
         self.__dict__ = dict_
-        self._attr_synced = False
 
     def __bool__(self):
         return True
@@ -269,6 +268,112 @@ class Cell(Cell_):
     def create_time(self):
         return self._create_time
 
+    @property
+    def _non_persistent_buffers_set(self):
+        """_non_persistent_buffers_set"""
+        if self._lazy_non_persistent_buffers_set is None:
+            super().__setattr__("_lazy_non_persistent_buffers_set", set())
+        return self._lazy_non_persistent_buffers_set
+
+    @property
+    def _state_dict_hooks(self):
+        """_state_dict_hooks"""
+        if self._lazy_state_dict_hooks is None:
+            super().__setattr__("_lazy_state_dict_hooks", OrderedDict())
+        return self._lazy_state_dict_hooks
+
+    @property
+    def _state_dict_pre_hooks(self):
+        """_state_dict_pre_hooks"""
+        if self._lazy_state_dict_pre_hooks is None:
+            super().__setattr__("_lazy_state_dict_pre_hooks", OrderedDict())
+        return self._lazy_state_dict_pre_hooks
+
+    @property
+    def _load_state_dict_pre_hooks(self):
+        """_load_state_dict_pre_hooks"""
+        if self._lazy_load_state_dict_pre_hooks is None:
+            super().__setattr__("_lazy_load_state_dict_pre_hooks", OrderedDict())
+        return self._lazy_load_state_dict_pre_hooks
+
+    @property
+    def _load_state_dict_post_hooks(self):
+        """_load_state_dict_post_hooks"""
+        if self._lazy_load_state_dict_post_hooks is None:
+            super().__setattr__("_lazy_load_state_dict_post_hooks", OrderedDict())
+        return self._lazy_load_state_dict_post_hooks
+
+    @property
+    def compile_cache(self):
+        """compile_cache"""
+        if self._compile_cache is None:
+            super().__setattr__("_compile_cache", set())
+        return self._compile_cache
+
+    @property
+    def phase_cache(self):
+        """phase_cache"""
+        if self._phase_cache is None:
+            super().__setattr__("_phase_cache", dict())
+        return self._phase_cache
+
+    @property
+    def _forward_pre_hook(self):
+        """_forward_pre_hook"""
+        if self._lazy_forward_pre_hook is None:
+            super().__setattr__("_lazy_forward_pre_hook", OrderedDict())
+        return self._lazy_forward_pre_hook
+
+    @property
+    def _forward_hook(self):
+        """_forward_hook"""
+        if self._lazy_forward_hook is None:
+            super().__setattr__("_lazy_forward_hook", OrderedDict())
+        return self._lazy_forward_hook
+
+    @property
+    def _backward_pre_hook(self):
+        """_backward_pre_hook"""
+        if self._lazy_backward_pre_hook is None:
+            super().__setattr__("_lazy_backward_pre_hook", OrderedDict())
+        return self._lazy_backward_pre_hook
+
+    @property
+    def _backward_hook(self):
+        """_backward_hook"""
+        if self._lazy_backward_hook is None:
+            super().__setattr__("_lazy_backward_hook", OrderedDict())
+        return self._lazy_backward_hook
+
+    @property
+    def _forward_pre_hook_with_kwargs(self):
+        """_backward_hook"""
+        if self._lazy_forward_pre_hook_with_kwargs is None:
+            super().__setattr__("_lazy_forward_pre_hook_with_kwargs", OrderedDict())
+        return self._lazy_forward_pre_hook_with_kwargs
+
+    @property
+    def _forward_hook_with_kwargs(self):
+        """_backward_hook"""
+        if self._lazy_forward_hook_with_kwargs is None:
+            super().__setattr__("_lazy_forward_hook_with_kwargs", OrderedDict())
+        return self._lazy_forward_hook_with_kwargs
+
+    @property
+    def _user_parameters(self):
+        """_user_parameters"""
+        if self._lazy_user_parameters is None:
+            super().__setattr__("_lazy_user_parameters", [])
+        return self._lazy_user_parameters
+
+    @_user_parameters.setter
+    def _user_parameters(self, value):
+        """_user_parameters"""
+        if not isinstance(value, list):
+            raise TypeError(f"For 'Cell', the property '_user_parameters' must be list type, "
+                            f"but got type {type(value)}.")
+        self._lazy_user_parameters = value
+
     @property
     def cell_init_args(self):
         return self._cell_init_args
@@ -279,15 +384,21 @@ class Cell(Cell_):
         Get exist parameter names adding by tuple or list of parameter.
         """
         if self._exist_names is None:
-            self._exist_names = set("")
+            super().__setattr__("_exist_names", set(""))
         return self._exist_names
 
     @property
     def exist_objs(self):
         if self._exist_objs is None:
-            self._exist_objs = set()
+            super().__setattr__("_exist_objs", set())
         return self._exist_objs
 
+    @property
+    def _construct_sig(self):
+        if self._lazy_construct_sig is None:
+            super().__setattr__("_lazy_construct_sig", inspect.signature(self.construct))
+        return self._lazy_construct_sig
+
     @property
     def param_prefix(self):
         """
@@ -319,6 +430,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):
         """
@@ -381,6 +499,8 @@ class Cell(Cell_):
         `parameter_layout_dict` represents the tensor layout of a parameter, which is inferred by shard strategy and
         distributed operator information.
         """
+        if self._parameter_layout_dict is None:
+            super().__setattr__("_parameter_layout_dict", {})
         return self._parameter_layout_dict
 
     @property
@@ -396,6 +516,8 @@ class Cell(Cell_):
 
     @property
     def parallel_parameter_name_list(self):
+        if self._parallel_parameter_name_list is None:
+            super().__setattr__("_parallel_parameter_name_list", ())
         return self._parallel_parameter_name_list
 
     @parallel_parameter_name_list.setter
@@ -435,10 +557,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)))
@@ -450,6 +585,8 @@ class Cell(Cell_):
 
     @property
     def parallel_parameter_merge_net_dict(self):
+        if self._parallel_parameter_merge_net_dict is None:
+            super().__setattr__("_parallel_parameter_merge_net_dict", {})
         return self._parallel_parameter_merge_net_dict
 
     @parallel_parameter_merge_net_dict.setter
@@ -867,6 +1004,7 @@ class Cell(Cell_):
         if hasattr(self, "compile_cache") and self.compile_cache:
             _cell_graph_executor.del_net_res(self, self.compile_cache)
         Cell.total_instance_count -= 1
+        Cell.global_cells.pop(self, None)
 
     def __delattr__(self, name):
         if name in self._params:
@@ -879,47 +1017,15 @@ class Cell(Cell_):
             del self._params_list[name]
         else:
             object.__delattr__(self, name)
-        self._attr_synced = False
-
-    def _cast_mixed_precision_inputs(self, inputs, dst_type):
-        """Cast input for mixed precision"""
-        res = list()
-        for item in inputs:
-            if isinstance(item, tuple):
-                res.append(self._cast_mixed_precision_inputs(item, dst_type))
-            elif isinstance(item, float):
-                res.append(self.cast(item, dst_type))
-            elif hasattr(item, "dtype") and item.dtype in \
-                    {mstype.float16, mstype.float32, mstype.float64, mstype.bfloat16} and item.dtype != dst_type:
-                res.append(self.cast(item, dst_type))
-            else:
-                res.append(item)
-        return tuple(res)
 
     def cast_inputs(self, inputs, dst_type):
         """
         Cast inputs to specified type.
 
-        Args:
-            inputs (tuple[Tensor]): The cell inputs.
-            dst_type (mindspore.dtype): The specified data type.
-
-        returns:
-            tuple[Tensor], the result with destination data type.
+        .. warning::
+            This interface will be deprecated in future versions.
         """
-        res = list()
-        for item in inputs:
-            if isinstance(item, tuple):
-                res.append(self.cast_inputs(item, dst_type))
-            else:
-                res.append(self.cast(item, dst_type))
-        return tuple(res)
-
-    def _do_parameter_broadcast(self):
-        if context.get_auto_parallel_context("parallel_mode") == ParallelMode.DATA_PARALLEL:
-            if not self.parameter_broadcast_done:
-                _pynative_executor.parameter_broadcast(self, self.phase)
-                self.parameter_broadcast_done = True
+        logger.warning(f"'cast_inputs' will be deprecated in future versions.")
 
     def run_construct(self, cast_inputs, kwargs):
         """
@@ -940,30 +1046,34 @@ class Cell(Cell_):
         output = self._run_construct(cast_inputs, kwargs)
         return output
 
-    def _run_construct(self, *inputs, **kwargs):
+    def _run_construct(self, *args, **kwargs):
         """Run the construct function"""
         if self._forward_pre_hook:
-            inputs = self._run_forward_pre_hook(inputs)
+            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(*inputs, **kwargs)
+            output = self._shard_fn(*args, **kwargs)
         elif _pynative_executor.requires_grad():
-            if self._backward_hook:
-                output = self._backward_hook_construct(*inputs, **kwargs)
-            elif self._recompute_cell is not None:
-                output = self._recompute_cell(*inputs, **kwargs)
+            if self._recompute_cell is not None:
+                output = self._recompute_cell(*args, **kwargs)
             elif self.has_bprop:
-                output = self._call_custom_bprop(*inputs, **kwargs)
+                output = self._call_custom_bprop(*args, **kwargs)
             else:
-                output = self.construct(*inputs, **kwargs)
+                output = self.construct(*args, **kwargs)
         else:
-            output = self.construct(*inputs, **kwargs)
+            output = self.construct(*args, **kwargs)
 
         if self._forward_hook:
-            output = self._run_forward_hook(inputs, output)
+            output = self._run_forward_hook(args, kwargs, output)
+
+        if self._backward_hook:
+            output = self._cell_backward_hook(output)
 
         if self._backward_pre_hook:
-            output = self._run_backward_pre_hook(output)
+            output = self._cell_backward_pre_hook(output)
 
         return output
 
@@ -998,22 +1108,6 @@ class Cell(Cell_):
                             f"{default_args} default argument, total {positional_args + default_args}, "
                             f"but got {len(args)}.")
 
-    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():
@@ -1039,9 +1133,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")
@@ -1120,8 +1211,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:
@@ -1129,11 +1218,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
@@ -1145,11 +1233,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'.
@@ -1157,14 +1248,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
@@ -1182,40 +1265,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
-        shard_fn = Shard()
-        fn = shard_fn(self, in_strategy, out_strategy, parameter_plan, device, level)
-        self._shard_fn = fn
-        return fn
-
-    def auto_cast_inputs(self, inputs):
-        """
-        Auto cast inputs in mixed precision scenarios.
-
-        Args:
-            inputs (tuple): the inputs of construct.
-
-        Returns:
-            Tuple, the inputs after data type cast.
-        """
-        msg = f"'auto_cast_inputs' is deprecated from version 2.0 and will be removed in a future version."
-        logger.warning(msg)
-        cast_inputs = inputs
-        mixed_type = self.get_mixed_precision_type()
-        if mixed_type == MixedPrecisionType.FP16:
-            cast_inputs = self._cast_mixed_precision_inputs(inputs, mstype.float16)
-        if mixed_type == MixedPrecisionType.FP32:
-            cast_inputs = self._cast_mixed_precision_inputs(inputs, mstype.float32)
+            return
 
-        return cast_inputs
+        shard_fn = Shard()
+        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):
@@ -1224,15 +1301,25 @@ 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'''
         if not hasattr(self, "phase"):
             return False, None
         if (self.phase == "prefill" or self.phase == 'increment') and self.phase in self.phase_cache:
-            new_args = _get_args_for_run_predict(self, args, kwargs, self._compile_args)
+            new_args = _get_args_for_run(self, args, kwargs, self._has_mutable_args_list, True)
+            if self.jit_config_dict:
+                jit_config_dict = self.jit_config_dict
+            else:
+                jit_config_dict = JitConfig().jit_config_dict
+            _cell_graph_executor._graph_executor.set_jit_config(jit_config_dict)
             res = _cell_graph_executor._graph_executor(tuple(new_args), self.phase_cache[self.phase])
             res = _convert_python_data(res)
             return True, res
@@ -1241,8 +1328,9 @@ 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.sig.bind(*args, **kwargs)
+                bound_arguments = self._construct_sig.bind(*args, **kwargs)
                 bound_arguments.apply_defaults()
                 args = bound_arguments.args
                 kwargs = bound_arguments.kwargs
@@ -1251,11 +1339,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
 
@@ -1324,37 +1409,12 @@ class Cell(Cell_):
         """
         with _no_grad():
             output = self.construct(*args, **kwargs)
-        _pynative_executor.call_custom_bprop(self, output, *args, **kwargs)
-        return output
+        return _pynative_executor.call_custom_bprop(self, output, *args, **kwargs)
 
     def _add_attr(self, name, value):
         if name and name[:2] != '__' and name not in Cell.IGNORE_LIST:
             super(Cell, self)._add_attr(name, value)
 
-    def _sync_attr_for_compile(self):
-        """Sync the attr to c++ object."""
-        if self._attr_synced:
-            return
-        cells = self.__dict__.get('_cells')
-        for key in cells:
-            cell = cells[key]
-            cell._sync_attr_for_compile()
-            self._add_attr(key, cell)
-        params = self.__dict__.get('_params')
-        for key in params:
-            if '.' in key:
-                continue
-            param = params[key]
-            self._add_attr(key, param)
-        params_list = self.__dict__.get('_params_list')
-        for key in params_list:
-            params_list_item = params_list[key]
-            self._add_attr(key, params_list_item)
-        for key in self.__dict__:
-            value = self.__dict__[key]
-            self._add_attr(key, value)
-        self._attr_synced = True
-
     def _set_attr_for_param_or_param_tuple(self, name, value):
         """Set attr for param and tensor."""
         if isinstance(value, Parameter):
@@ -1369,27 +1429,16 @@ class Cell(Cell_):
                     # If there are multiple identical objects, their names only check once.
                     continue
                 exist_objs.add(item)
-                if item.name == PARAMETER_NAME_DEFAULT:
-                    logger.warning("For 'Cell', the parameter definition is deprecated.\n"
-                                   "Please set a unique name for the parameter in ParameterTuple '{}'.".format(value))
-                    item.name = item.name + "$" + str(self._id)
+                if _is_parameter_generated(item.name):
+                    item.name = "Parameter$" + str(self._id)
                     self._id += 1
-                self.insert_param_to_cell(item.name, item, check_name_contain_dot=False)
                 if item.name in exist_names:
                     raise ValueError("The value {} , its name '{}' already exists. "
                                      "Please set a unique name for the parameter.".format(value, item.name))
                 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."""
@@ -1398,9 +1447,6 @@ class Cell(Cell_):
                 # If there are multiple identical objects, their names only check once.
                 continue
             self.exist_objs.add(item)
-            if item.name == PARAMETER_NAME_DEFAULT:
-                item.name = item.name + "$" + str(self._id)
-                self._id += 1
             if item.name in self.exist_names:
                 raise ValueError(f"The value {value} , its name '{item.name}' already exists. "
                                  "Please set a unique name for the parameter.")
@@ -1513,24 +1559,6 @@ class Cell(Cell_):
         main_str += ")"
         return main_str
 
-    def load_parameter_slice(self, params):
-        """
-        Replace parameters with sliced tensors by parallel strategies.
-
-        Note:
-            This interface is deprecated.
-        """
-        logger.warning("'load_parameter_slice' function is deprecated.")
-
-    def set_parallel_input_with_inputs(self, *inputs):
-        """
-        Slice inputs tensors by parallel strategies.
-
-        Note:
-            This interface is deprecated.
-        """
-        logger.warning("'set_parallel_input_with_inputs' function is deprecated.")
-
     def set_inputs(self, *inputs, **kwargs):
         """
         Save set inputs for computation graph. The number of inputs should be the same with that of the datasets. When
@@ -1589,8 +1617,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)
 
@@ -1665,7 +1691,6 @@ class Cell(Cell_):
             _cell_graph_executor._graph_executor.check_argument_consistency(compile_args, args, "set_inputs")
             self._check_parameter_consistency(compile_args, args)
             Validator.check_symbolic_shape(compile_args, args)
-            self.saved_dynamic_shape = compile_args
             return compile_args
         return args
 
@@ -1678,8 +1703,9 @@ class Cell(Cell_):
             kwargs (dict): Kwargs of the Cell object.
         """
         _init_auto_parallel_context(self)
-        self._compile_args = self._get_compile_args(args)
-        _cell_graph_executor.compile(self, *self._compile_args, phase=self.phase,
+        compile_args = self._get_compile_args(args)
+        self._has_mutable_args_list = _get_mutable_flags(compile_args)
+        _cell_graph_executor.compile(self, *compile_args, phase=self.phase,
                                      jit_config_dict=self._jit_config_dict, **kwargs)
         _clear_auto_parallel_context(self)
 
@@ -1698,25 +1724,14 @@ class Cell(Cell_):
             Object, the result of executing.
         """
         self.compile(*args, **kwargs)
-        self.add_flags(ge_sync_data=False)
-        new_args = _get_args_for_run(self, args, kwargs, self._compile_args)
+        new_args = _get_args_for_run(self, args, kwargs, self._has_mutable_args_list, False)
+        if self.jit_config_dict:
+            jit_config_dict = self.jit_config_dict
+        else:
+            jit_config_dict = JitConfig().jit_config_dict
+        _cell_graph_executor._graph_executor.set_jit_config(jit_config_dict)
         return _cell_graph_executor(self, *new_args, phase=self.phase)
 
-    def auto_parallel_compile_and_run(self):
-        """
-        Whether or not to execute compile and run in 'AUTO_PARALLEL' or 'SEMI_AUTO_PARALLEL' mode.
-
-        Note:
-            This interface is deprecated.
-        """
-        logger.warning("'auto_parallel_compile_and_run' function is deprecated.")
-
-    def exec_checkpoint_graph(self):
-        """Executes GE saving checkpoint graph operation."""
-        logger.warning("'exec_checkpoint_graph' function is deprecated.")
-        self.add_flags(ge_sync_data=True)
-        _cell_graph_executor(self, phase='save')
-
     def insert_param_to_cell(self, param_name, param, check_name_contain_dot=True):
         """
         Adds a parameter to the current cell.
@@ -1762,35 +1777,10 @@ class Cell(Cell_):
         if not isinstance(param, Parameter) and param is not None:
             raise TypeError(f"For 'insert_param_to_cell', the argument 'param' must be 'Parameter' if not None, "
                             f"but got {type(param)}.")
-        if isinstance(param, Parameter) and param.name == PARAMETER_NAME_DEFAULT:
+        if isinstance(param, Parameter) and _is_parameter_generated(param.name):
             param.name = param_name
         self._params[param_name] = param
 
-    def cast_param(self, param):
-        """
-        Cast parameter according to auto mix precision level in pynative mode.
-
-        This interface is currently used in the case of auto mix precision and usually needs not to be used explicitly.
-
-        Args:
-            param (Parameter): Parameters, the type of which should be cast.
-
-        Returns:
-            Parameter, the input parameter with type automatically cast.
-        """
-        msg = f"'cast_param' is deprecated from version 2.0 and will be removed in a future version."
-        logger.warning(msg)
-        mixed_type = self.get_mixed_precision_type()
-        if mixed_type != MixedPrecisionType.NOTSET:
-            if mixed_type == MixedPrecisionType.FP32:
-                param.set_cast_dtype(mstype.float32)
-            elif mixed_type == MixedPrecisionType.FP16:
-                param.set_cast_dtype(mstype.float16)
-        elif hasattr(param, "set_cast_dtype"):
-            # retest dtype
-            param.set_cast_dtype()
-        return param
-
     def insert_child_to_cell(self, child_name, child_cell):
         """
         Adds a child cell to the current cell with a given name.
@@ -1850,27 +1840,10 @@ class Cell(Cell_):
         """
         Remove the redundant parameters.
 
-        This interface usually needs not to be used explicitly.
+        .. warning::
+            This interface will be deprecated in future versions.
         """
-        cells = self.cells_and_names()
-        for _, cell in cells:
-            params = cell._params.items()
-            for param_name, param in list(params):
-                if param.name not in self.parallel_parameter_name_list:
-                    cell._params.pop(param_name)
-                    logger.info("remove the redundant parameter: %s", param.name)
-                    continue
-            cell_dict = cell.__dict__
-            for key in cell_dict:
-                if isinstance(cell_dict[key], ParameterTuple):
-                    param_tuple = cell_dict[key]
-                    new_param_tuple = []
-                    for param in param_tuple:
-                        if param.name not in self.parallel_parameter_name_list:
-                            logger.info("remove the redundant parameter: %s in ParameterTuple", param.name)
-                            continue
-                        new_param_tuple.append(param)
-                    cell.__dict__[key] = ParameterTuple(new_param_tuple)
+        logger.warning(f"'remove_redundant_parameters' will be deprecated in future versions.")
 
     def _get_cell_parallel_mode(self):
         """Determine whether the current cell is in parallel mode."""
@@ -1926,16 +1899,13 @@ class Cell(Cell_):
         # replace all original usage.
         cells = self.cells_and_names()
         is_parallel_mode = self._get_cell_parallel_mode()
-        is_graph_mode = context.get_context('mode') == context.GRAPH_MODE
 
         for _, cell in cells:
             params = cell._params.items()
             for param_name, param in params:
-                not_sliced = not param.sliced
-                judgment = not_sliced
                 if param.param_info.is_pipeline_shared_param:
                     continue
-                if is_graph_mode and is_parallel_mode and judgment:
+                if is_parallel_mode and not param.sliced:
                     continue
                 if not auto_parallel_mode:
                     cell._params[param_name] = _updata(param)
@@ -1948,11 +1918,9 @@ class Cell(Cell_):
                     param_tuple = cell_dict[key]
                     new_param_tuple = []
                     for param in param_tuple:
-                        not_sliced = not param.sliced
-                        judgment = not_sliced
                         if param.param_info.is_pipeline_shared_param:
                             continue
-                        if is_graph_mode and is_parallel_mode and judgment:
+                        if is_parallel_mode and not param.sliced:
                             continue
                         if not auto_parallel_mode:
                             new_param_tuple.append(_updata(param))
@@ -2591,15 +2559,6 @@ class Cell(Cell_):
         self.add_flags_recursive(broadcast_flag=mode)
         return self
 
-    def set_auto_parallel(self):
-        """
-        Set the cell to auto parallel mode.
-
-        Note:
-            This interface is deprecated.
-        """
-        logger.warning("'set_auto_parallel' function is deprecated.")
-
     def set_jit_config(self, jit_config):
         """
         Set jit config for cell.
@@ -2645,25 +2604,38 @@ class Cell(Cell_):
             raise ValueError(f"Negative 'fusion_size' {fusion_size} is invalid.")
         Tensor._flatten_tensors(self.trainable_params(), fusion_size)  # pylint: disable=W0212
 
-    def register_forward_pre_hook(self, hook_fn):
+    @jit_forbidden_register
+    def register_forward_pre_hook(self, hook_fn, with_kwargs=False):
         """
         Register forward pre hook function for Cell object.
 
+        The hook will be called before :func:`mindspore.nn.Cell.construct` is invoked.
+
+        The hook function should be one of the following signatures:
+
+        - `hook_fn(cell, args) -> None or new_args` , when `with_kwargs` is ``Flase`` .
+        - `hook_fn(cell, args, kwargs) -> None or (new_args, new_kwargs)` , when `with_kwargs` is ``True`` .
+
+        where:
+
+        - `cell` (Cell): Cell object on which the hook is registered.
+        - `args` (tuple): Positional arguments passed to the `construct` function.
+        - `kwargs` (dict): Keyword arguments passed to the `construct` function. Only passed to `hook_fn` when
+          `with_kwargs` is ``True`` .
+
         Note:
-            - The `register_forward_pre_hook(hook_fn)` does not work in graph mode or functions decorated with 'jit'.
-            - 'hook_fn' must be defined as the following code.
-              `cell` is the object of registered Cell. `inputs` is the forward
-              input objects passed to the Cell. The 'hook_fn' can modify the forward input objects by returning new
-              forward input objects.
-            - It should have the following signature:
-              hook_fn(cell, inputs) -> new input objects or none.
-            - In order to prevent running failed when switching to graph mode, it is not recommended to write it in the
-              `construct` function of Cell object. In the pynative mode, if the `register_forward_pre_hook` function is
-              called in the `construct` function of the Cell object, a hook function will be added at each run time of
-              Cell object.
+            - 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.
+            - 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.
+            - In the pynative mode, if this method is called inside the `construct` function of the Cell object, a
+              `hook_fn` will be added at each run time of Cell object.
 
         Args:
             hook_fn (function): Python function. Forward pre hook function.
+            with_kwargs (bool, optional): Specifies whether hook_fn will be passed the kwargs given to the `construct`
+                function. Default: ``False`` .
 
         Returns:
             A handle corresponding to the `hook_fn` . The handle can be used to remove the added `hook_fn` by calling
@@ -2702,16 +2674,41 @@ 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)
+        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
 
-    def _run_forward_pre_hook(self, inputs):
+    @jit_forbidden_register
+    def _run_forward_pre_hook(self, args, kwargs):
         """
         Running forward pre hook function registered on Cell object.
+        """
+        for hook_id, hook_fn in self._forward_pre_hook.items():
+            if hook_id in self._forward_pre_hook_with_kwargs:
+                ret = hook_fn(self, args, kwargs)
+                if ret is not None:
+                    if isinstance(ret, tuple) and len(ret) == 2:
+                        args, kwargs = ret
+                    else:
+                        raise RuntimeError(
+                            "forward pre hook with kwargs must return None or a tuple of (new_args, new_kwargs), "
+                            f"but got {ret}"
+                        )
+            else:
+                ret = hook_fn(self, args)
+                if ret is not None:
+                    if not isinstance(ret, tuple):
+                        ret = (ret,)
+                    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.
@@ -2731,34 +2728,43 @@ class Cell(Cell_):
                 else:
                     forward_pre_hook_inputs = ret
 
-        if isinstance(inputs, tuple):
-            if not isinstance(forward_pre_hook_inputs, tuple):
-                forward_pre_hook_inputs = (forward_pre_hook_inputs,)
-            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)))
+        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
 
-    def register_forward_hook(self, hook_fn):
+    @jit_forbidden_register
+    def register_forward_hook(self, hook_fn, with_kwargs=False):
         """
-        Set the Cell forward hook function.
+        Register forward hook function for Cell object.
+
+        This hook will be called after :func:`mindspore.nn.Cell.construct` has computed an output.
+
+        The hook function should be one of the following signatures:
+
+        - `hook_fn(cell, args, output) -> None or new_output` , when `with_kwargs` is ``False`` .
+        - `hook_fn(cell, args, kwargs, output) -> None or new_output` , when `with_kwargs` is ``True`` .
+
+        where:
+
+        - `cell` (Cell): Cell object on which the hook is registered.
+        - `args` (tuple): Positional arguments passed to the `construct` function.
+        - `kwargs` (dict): Keyword arguments passed to the `construct` function. Only passed to `hook_fn` when
+          `with_kwargs` is ``True`` .
+        - `output`: Output generated by the `construct` function.
 
         Note:
-            - The `register_forward_hook(hook_fn)` does not work in graph mode or functions decorated with 'jit'.
-            - 'hook_fn' must be defined as the following code.
-              `cell` is the object of registered Cell. `inputs` is the forward
-              input objects passed to the Cell. `output` is the forward output object of the Cell. The 'hook_fn' can
-              modify the forward output object by returning new forward output object.
-            - It should have the following signature:
-              hook_fn(cell, inputs, output) -> new output object or none.
-            - In order to prevent running failed when switching to graph mode, it is not recommended to write it in the
-              `construct` function of Cell object. In the pynative mode, if the `register_forward_hook` function is
-              called in the `construct` function of the Cell object, a hook function will be added at each run time of
-              Cell object.
+            - 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.
+            - In the pynative mode, if this method is called inside the `construct` function of the Cell object, a
+              `hook_fn` will be added at each run time of Cell object.
 
         Args:
             hook_fn (function): Python function. Forward hook function.
+            with_kwargs (bool, optional): Specifies whether hook_fn will be passed the kwargs given to the `construct`
+                function. Default: ``False`` .
 
         Returns:
             A handle corresponding to the `hook_fn` . The handle can be used to remove the added `hook_fn` by calling
@@ -2801,16 +2807,17 @@ 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)
+        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 _run_forward_hook(self, inputs, output):
+    def _jit_forward_hook(self, inputs, output):
         """
-        Running forward hook function registered on Cell object.
+        Compile forward hook function registered on Cell object.
 
         Args:
             inputs: The input objects of Cell object.
@@ -2837,12 +2844,26 @@ class Cell(Cell_):
                         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.
+        """
+        for hook_id, hook_fn in self._forward_hook.items():
+            if hook_id in self._forward_hook_with_kwargs:
+                ret = hook_fn(self, args, kwargs, output)
+            else:
+                ret = hook_fn(self, args, output)
+            if ret is not None:
+                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:
@@ -2891,44 +2912,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.
 
@@ -2981,9 +2975,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
@@ -3029,9 +3022,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):
@@ -3116,7 +3108,6 @@ class Cell(Cell_):
             OrderedDict([('param_a', Parameter (name=param_a, shape=(3,), dtype=Int64, requires_grad=True)), \
             ('buffer_a', Tensor(shape=[3], dtype=Int64, value= [4, 5, 6]))])
         """
-        # TODO: Remove `args` and the parsing logic when BC allows.
         if args:
             # DeprecationWarning is ignored by default
             warnings.warn(
@@ -3169,7 +3160,7 @@ class Cell(Cell_):
 
         It should have the following signature:
 
-        hook(cell, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None  # noqa: B950
+        hook(cell, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None
 
         Args:
             hook (Callable): The hook function before `load_state_dict` is called.
@@ -3178,9 +3169,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
@@ -3212,9 +3202,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(
@@ -3450,12 +3439,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
@@ -3507,65 +3496,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
@@ -3584,12 +3525,6 @@ class Cell(Cell_):
         for param in params:
             param.set_param_ps(init_in_server)
 
-    @deprecated("1.8", "set_param_fl")
-    def set_param_fl(self, push_to_server=False, pull_from_server=False, requires_aggr=True):
-        params = self.parameters_and_names()
-        for param in params:
-            param[1].set_param_fl(push_to_server, pull_from_server, requires_aggr)
-
     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
@@ -3650,7 +3585,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."
@@ -3693,12 +3628,12 @@ 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:
+        if context._get_mode() == context.PYNATIVE_MODE:
             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))
-        if 'parallel_optimizer_comm_recompute' in kwargs.keys():
+        if 'parallel_optimizer_comm_recompute' in kwargs:
             if kwargs.get('parallel_optimizer_comm_recompute', False):
                 logger.warning("Currently, the communication operator allgathers introduced by optimizer shard "
                                "is replaced with zero3.")
@@ -3711,38 +3646,6 @@ class Cell(Cell_):
                                  "the key kwargs must be 'mp_comm_recompute', "
                                  "'parallel_optimizer_comm_recompute', 'recompute_slice_activation'" % key)
 
-    @deprecated("2.3", "infer_param_pipeline_stage")
-    def infer_param_pipeline_stage(self):
-        """
-        Infer pipeline stages of all parameters in the cell.
-
-        Note:
-            - The interface is deprecated from version 2.3 and will be removed in a future version.
-
-        Returns:
-            The params belong to current stage in pipeline parallel.
-
-        Raises:
-            RuntimeError: If there is a parameter does not belong to any stage.
-        """
-        from mindspore.parallel._utils import _get_global_rank, _get_device_num
-        logger.warning(f"This interface may be deleted in the future.")
-        stage_num = context.get_auto_parallel_context("pipeline_stages")
-        device_num = _get_device_num()
-        rank_id = _get_global_rank()
-        per_stage_devices = device_num // stage_num
-        current_stage = rank_id // per_stage_devices
-        params = []
-        for param in self.trainable_params():
-            if not param._pipeline_stage_list:  # pylint: disable=W0212
-                raise RuntimeError("For 'infer_param_pipeline_stage', the parameter {} does not belong to any stage, "
-                                   "please check whether the cell where the param locates has been set "
-                                   "'pipeline_stage'. Otherwise, the parameter should use 'add_pipeline_stage' "
-                                   "to add its stage information".format(param.name))
-            if current_stage in param._pipeline_stage_list:
-                params.append(param)
-        return params
-
     def place(self, role, rank_id):
         """
         Set the label for all operators in this cell.
@@ -3772,19 +3675,6 @@ class Cell(Cell_):
         for op in all_ops:
             op.place(role, rank_id)
 
-    def _mixed_precision_cast(self, inputs):
-        mixed_type = self.get_mixed_precision_type()
-        if mixed_type == MixedPrecisionType.NOTSET:
-            return inputs
-        if mixed_type == MixedPrecisionType.FP16:
-            cast_type = mstype.float16
-        elif mixed_type == MixedPrecisionType.BF16:
-            cast_type = mstype.bfloat16
-        else:
-            cast_type = mstype.float32
-        cast_inputs = self._cast_mixed_precision_inputs(inputs, cast_type)
-        return cast_inputs
-
     def _get_attr_from_cell(self, network):
         if not isinstance(network, Cell):
             return
@@ -3793,92 +3683,70 @@ class Cell(Cell_):
         if hasattr(network, "_amp_level"):
             self._amp_level = getattr(network, "_amp_level")
 
-    def _register_parameters_hook(self, forward_hook=None, backward_hook=None, all=False):
+    def _set_jit_graph_name(self, key):
+        """
+        Set jit graph name.
         """
-        Register the forward hook for parameters and register the backward hook for the corresponding gradient.
+        self._jit_graph_name = key
 
-        .. warning::
-            This is an experimental prototype that is subject to change and/or deletion.
+    def _jit_backward_pre_hook(self, grad_output):
+        new_grad_output = grad_output
+        if not isinstance(grad_output, tuple):
+            new_grad_output = (grad_output,)
 
-        Note:
-            - The `_register_parameters_hook(forward_hook, backward_hook)` only work in graph mode
-            - The `forward_hook` must be defined as the following code.
-              `parameters`: the tuple of the trainble parameters of the Cell, each element in the tuple shuould be
-               in the format of `(param_name, Parameter)`.
-            - The `forward_hook` should have the following signature:
-              forward_hook(parameters) -> None.
-            - The `backward_hook` must be defined as the following code.
-              `gradients`: the tuple of the gradients corresponding to the trainble parameters of the Cell, each
-               element in the tuple shuould be in the format of `(param_name, gradient)`.
-            - The `backward_hook` should have the following signature:
-              backward_hook(parameters) -> New gradients.
+        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
 
-        Args:
-            forward_hook (function, optional): Python function or ``None``, Forward hook function. Default: ``None``
-            backward_hook (function, optional): Python function or ``None``, Backward hook function. Default ``None``
-            all (bool, optional): bool, whether to set hooks for all sub cells recursively. Default: ``False``
+        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)))
 
-        Returns:
-            None
+        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)))
 
-        Raises:
-            RuntimeError: If the `forward_hook` or `backward_hook ` has unspoorted syntax under GRAPH MODE.
-            TypeError: If the `forward_hook` or `backward_hook` is not defined as required.
+        return new_grad_output
 
-        Supported Platforms:
-        ``Ascend`` ``GPU`` ``CPU``
+    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,)
 
-        Examples:
-            >>> import mindspore as ms
-            >>> from mindspore import Tensor, nn, ops, Parameter
-            >>>
-            >>> ms.set_context(mode=ms.GRAPH_MODE)
-            >>> def parameter_hook(parameters):
-            ...     print("--- enter parameter hook ---")
-            ...     for name, param in parameters:
-            ...         print (name, param)
-            ...     print("--- leave parameter hook ---")
-            ...
-            >>> def gradient_hook(gradients):
-            ...     print("--- enter gradient hook ---")
-            ...     outs = []
-            ...     for name, gradient in gradients:
-            ...         print(name, gradient)
-            ...         outs.append(gradient * 2) # double gradient
-            ...     print("--- leave gradient hook ---")
-            ...     return outs
-            ...
-            >>> class Net(nn.Cell):
-            ...     def __init__(self)
-            ...         super(Net, self).__init__()
-            ...         self.w = Parameter(Tensor(np.array([3.0], np.float32)), name='w')
-            ...     def construct(self, x):
-            ...         return self.w * x
-            ...
-            >>> grad = ops.GradOperation(get_by_list=True)
-            >>> net = Net()
-            >>> net._register_parameters_hook(forward_hook=parameter_hook, backward_hook=gradient_hook)
-            >>> x = Tensor(np.array([4.0]).astype(np.float32))
-            >>> output = grad(net, net.trainable_params())(x)
-            --- enter parameter hook ---
-            w
-            Tensor(shape=[1], dtype=Float32, value=[ 3.00000000e+00])
-            --- leave parameter hook ---
-            --- enter gradient hook ---
-            w
-            Tensor(shape=[1], dtype=Float32, value=[ 4.00000000e+00])
-            --- leave gradient hook ---
-            >>> print("doubled grad: ", output)
-            doubled grad: (Tensor(shape=[1], dtype=Float32, value=[ 8.00000000e+00]),)
-        """
-        if not all:
-            self._parameters_forward_hook = forward_hook
-            self._parameters_backward_hook = backward_hook
-        else:
-            for _, cell in self.cells_and_names():
-                cell._parameters_forward_hook = forward_hook
-                cell._parameters_backward_hook = backward_hook
+        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):
     """