mindspore 2.5.0__cp311-cp311-win_amd64.whl → 2.6.0__cp311-cp311-win_amd64.whl

mindspore/nn/cell.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2024 Huawei Technologies Co., Ltd
+# Copyright 2020-2025 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,13 +15,26 @@
 """cell"""
 from __future__ import absolute_import
 
-import gc
 import inspect
 import os
 import time
-from collections import OrderedDict
-import numpy
-
+import warnings
+import itertools
+from collections import OrderedDict, namedtuple
+from typing import (
+    Dict,
+    Optional,
+    Set,
+    Callable,
+    List,
+    Tuple,
+    Iterator,
+    Any,
+    TypeVar,
+    Mapping
+)
+
+import mindspore as ms
 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 import log as logger
@@ -34,19 +47,62 @@ 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 _generate_branch_control_input, _convert_python_data, _get_args_for_run_predict
+from mindspore.common.api import _convert_python_data, _get_args_for_run_predict
 from mindspore.common.api import _process_dyn_args, _generate_dyn_compile_args
-from mindspore.common.parameter import Parameter, ParameterTuple
+from mindspore.common.parameter import _Buffer, Parameter, ParameterTuple
 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",
+]
+
+_global_buffer_registration_hooks: Dict[int, Callable] = OrderedDict()
+_EXTRA_STATE_KEY_SUFFIX = "_extra_state"
+
+
+class _IncompatibleKeys(namedtuple("IncompatibleKeys", ["missing_keys", "unexpected_keys"]),):
+    def __repr__(self):
+        if not self.missing_keys and not self.unexpected_keys:
+            return "<All keys matched successfully>"
+        return super().__repr__()
+
+    __str__ = __repr__
+
+
+def register_cell_buffer_registration_hook(hook: Callable[..., None],):
+    r"""Register a buffer registration hook common to all cells.
+
+    .. warning ::
+
+        This adds global state to the `nn.Cell` cell
+
+    The hook will be called every time :func:`register_buffer` is invoked.
+    It should have the following signature::
+
+        hook(cell, name, buffer) -> None or new buffer
+
+    The hook can modify the input or return a single modified value in the hook.
+
+    Returns:
+        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
+    return handle
+
+
+
 class Cell(Cell_):
     """
     The basic building block of neural networks in MindSpore. The model or neural network layer should inherit this
@@ -108,6 +164,8 @@ class Cell(Cell_):
                    '_attr_synced', 'pynative', 'requires_grad', 'cell_type',
                    '_parameters_forward_hook', '_parameters_backward_hook']
     total_instance_count = 0
+    _buffers: Dict[str, Optional[Tensor]]
+    _non_persistent_buffers_set: Set[str]
 
     def __init__(self, auto_prefix=True, flags=None):
         Cell_.__init__(self, self._cell_tag)
@@ -115,10 +173,17 @@ class Cell(Cell_):
         self.instance_count = Cell.total_instance_count
         self._params = OrderedDict()
         self._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 = ''
@@ -135,8 +200,8 @@ class Cell(Cell_):
         cells_compile_cache[id(self)] = self.compile_cache
         self.parameter_broadcast_done = False
         self._id = 1
-        self.exist_names = set("")
-        self.exist_objs = set()
+        self._exist_objs = None
+        self._exist_names = None
         self._recompute_cell = None
         self.mixed_precision_type = None
         self.sig = inspect.signature(self.construct)
@@ -146,7 +211,6 @@ class Cell(Cell_):
         if os.getenv('GC_COLLECT_IN_CELL') == '1':
             logger.warning("The convenient environment 'GC_COLLECT_IN_CELL' is deprecated from version 2.5 "
                            "and will be removed in a future version.")
-            gc.collect()
 
         if flags:
             self.add_flags(**flags)
@@ -209,6 +273,21 @@ class Cell(Cell_):
     def cell_init_args(self):
         return self._cell_init_args
 
+    @property
+    def exist_names(self):
+        """
+        Get exist parameter names adding by tuple or list of parameter.
+        """
+        if self._exist_names is None:
+            self._exist_names = set("")
+        return self._exist_names
+
+    @property
+    def exist_objs(self):
+        if self._exist_objs is None:
+            self._exist_objs = set()
+        return self._exist_objs
+
     @property
     def param_prefix(self):
         """
@@ -237,11 +316,6 @@ class Cell(Cell_):
     def bprop_debug(self):
         """
         Get whether cell custom bprop debug is enabled.
-
-        Tutorial Examples:
-            - `Custom Neural Network Layers - Custom Cell Reverse
-              <https://mindspore.cn/docs/en/master/model_train/custom_program/network_custom.html
-              #custom-cell-reverse>`_
         """
         return self._bprop_debug
 
@@ -358,8 +432,6 @@ class Cell(Cell_):
             raise ValueError("For 'Cell', the property 'pipeline_stage' "
                              "can not be less than 0, but got {}".format(value))
         self._pipeline_stage = value
-        for item in self.trainable_params():
-            item.add_pipeline_stage(value)
 
     @property
     def pipeline_segment(self):
@@ -395,6 +467,374 @@ class Cell(Cell_):
     def enable_backward_hook(self):
         return self._enable_backward_hook
 
+    @jit_forbidden_register
+    def register_buffer(
+            self, name: str, tensor: Optional[Tensor], persistent: bool = True
+    ) -> None:
+        r"""Add a buffer to the cell.
+
+        This is typically used to register a buffer that should not to be
+        considered a model parameter. For example, BatchNorm's `running_mean`
+        is not a parameter, but is part of the cell's state. Buffers, by
+        default, are persistent and will be saved alongside parameters. This
+        behavior can be changed by setting `persistent` to ``False`` . The
+        only difference between a persistent buffer and a non-persistent buffer
+        is that the latter will not be a part of this cell's :attr:`state_dict` .
+
+        Buffers can be accessed as attributes using given names.
+
+        Args:
+            name (str): name of the buffer. The buffer can be accessed
+                from this cell using the given name.
+            tensor (Tensor): Buffer to be registered. If ``None`` ,
+                the buffer is not included in the cell's :attr:`state_dict` .
+            persistent (bool, optional): Whether the buffer is part of this cell's :attr:`state_dict`. Default ``True``.
+
+        Examples:
+            >>> import mindspore
+            ...
+            >>> class Net(mindspore.nn.Cell):
+            ...    def __init__(self):
+            ...        super().__init__()
+            ...        self.register_buffer("buffer0", mindspore.tensor([1, 2, 3]))
+            ...
+            ...    def construct(self, x):
+            ...        return x + self.net_buffer
+            ...
+            >>> net = Net()
+            >>> net.register_buffer("buffer0", mindspore.tensor([4, 5, 6]))
+            >>> print(net.buffer0)
+            [4 5 6]
+        """
+
+        if "_buffers" not in self.__dict__:
+            raise AttributeError("cannot assign buffer before Cell.__init__() call")
+        if not isinstance(name, str):
+            raise TypeError(
+                f"buffer name should be a string.But got this type: {type(name)}"
+            )
+        if "." in name:
+            raise KeyError('buffer name can\'t contain "."')
+        if name == "":
+            raise KeyError('buffer name can\'t be empty string ""')
+        if hasattr(self, name) and name not in self._buffers:
+            raise KeyError(f"attribute '{name}' already exists")
+        if tensor is not None and not isinstance(tensor, Tensor):
+            raise TypeError(
+                f"cannot assign '{type(tensor)}' object to buffer '{name}' "
+                "(mindspore Tensor or None required)"
+            )
+        for hook in _global_buffer_registration_hooks.values():
+            output = hook(self, name, tensor)
+            if output is not None:
+                tensor = output
+        if tensor is not None:
+            tensor._is_buffer = True
+        self._buffers[name] = tensor
+        if persistent:
+            self._non_persistent_buffers_set.discard(name)
+        else:
+            self._non_persistent_buffers_set.add(name)
+
+    @jit_forbidden_register
+    def get_buffer(self, target: str) -> "Tensor":
+        """Return the buffer given by `target` if it exists, otherwise throw an error.
+
+        See the docstring for `get_sub_cell` for a more detailed
+        explanation of this method's functionality as well as how to
+        correctly specify `target` .
+
+        Args:
+            target (str): The fully-qualified string name of the buffer
+                to look for. (See `get_sub_cell` for how to specify a
+                fully-qualified string.)
+
+        Returns:
+            Tensor
+
+        Examples:
+            >>> import mindspore
+            ...
+            ...
+            >>> class NetC(mindspore.nn.Cell):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.register_buffer("buffer_c", mindspore.tensor([0, 0, 0]))
+            ...
+            ...     def construct(self, x):
+            ...         return x + self.buffer_c
+            ...
+            ...
+            >>> class NetB(mindspore.nn.Cell):
+            ...     def __init__(self, net_c):
+            ...         super().__init__()
+            ...         self.net_c = net_c
+            ...         self.register_buffer("buffer_b", mindspore.tensor([1, 2, 3]))
+            ...
+            ...     def construct(self, x):
+            ...         return self.net_c(x) + self.buffer_b
+            ...
+            ...
+            >>> class NetA(mindspore.nn.Cell):
+            ...     def __init__(self, net_b):
+            ...         super().__init__()
+            ...         self.net_b = net_b
+            ...         self.register_buffer("buffer_a", mindspore.tensor([4, 5, 6]))
+            ...
+            ...     def construct(self, x):
+            ...         return self.net_b(x) + self.buffer_a
+            ...
+            ...
+            >>> net_c = NetC()
+            >>> net_b = NetB(net_c)
+            >>> net_a = NetA(net_b)
+            >>> buffer_c = net_a.get_buffer("net_b.net_c.buffer_c")
+            >>> print(f'buffer_c is {buffer_c}')
+            buffer_c is [0 0 0]
+
+        """
+        cell_path, _, buffer_name = target.rpartition(".")
+
+        cell = self.get_sub_cell(cell_path)
+
+        if not hasattr(cell, buffer_name):
+            raise AttributeError(
+                cell._get_name() + " has no attribute `" + buffer_name + "`"
+            )
+
+        buffer = getattr(cell, buffer_name)
+
+        if buffer_name not in cell._buffers:
+            raise AttributeError("`" + buffer_name + "` is not a buffer")
+
+        return buffer
+
+    @jit_forbidden_register
+    def named_buffers(
+            self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
+    ) -> Iterator[Tuple[str, Tensor]]:
+        r"""Return an iterator over cell buffers, yielding both the name of the buffer as well as the buffer itself.
+
+        Args:
+            prefix (str, optional): prefix to prepend to all buffer names. Default ``""``.
+            recurse (bool, optional): if ``True`` , then yields buffers of this cell
+                and all sub cells. Otherwise, yields only buffers that
+                are direct members of this cell. Default ``True``.
+            remove_duplicate (bool, optional): Whether to remove the duplicated buffers in the result. Default ``True``.
+
+        Returns:
+            Iterator[Tuple[str, Tensor]], an iterator of tuple containing the name and buffer.
+
+        Examples:
+            >>> import mindspore
+            ...
+            ...
+            >>> class NetB(mindspore.nn.Cell):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.register_buffer("buffer_b", mindspore.tensor([1, 2, 3]))
+            ...
+            ...     def construct(self, x):
+            ...         return x + self.buffer_b
+            ...
+            ...
+            >>> class NetA(mindspore.nn.Cell):
+            ...     def __init__(self, net_b):
+            ...         super().__init__()
+            ...         self.net_b = net_b
+            ...         self.register_buffer("buffer_a", mindspore.tensor([4, 5, 6]))
+            ...
+            ...     def construct(self, x):
+            ...         return self.net_b(x) + self.buffer_a
+            ...
+            ...
+            >>> net_b = NetB()
+            >>> net_a = NetA(net_b)
+            >>>
+            >>> for name, buffer in net_a.named_buffers():
+            ...     print(f'buffer name is {name}, buffer is {buffer}')
+            buffer name is buffer_a, buffer is [4 5 6]
+            buffer name is net_b.buffer_b, buffer is [1 2 3]
+
+        """
+        gen = self._named_members(
+            lambda cell: cell._buffers.items(),
+            prefix=prefix,
+            recurse=recurse,
+            remove_duplicate=remove_duplicate,
+        )
+        yield from gen
+
+    @jit_forbidden_register
+    def buffers(self, recurse: bool = True) -> Iterator[Tensor]:
+        r"""Return an iterator over cell buffers.
+
+        Args:
+            recurse (bool, optional): If ``True`` , then yields buffers of this cell
+                and all sub cells. Otherwise, yields only buffers that
+                are direct members of this cell. Default ``True``.
+
+        Returns:
+            Iterator[Tensor], an iterator of buffer.
+
+        Examples:
+            >>> import mindspore
+            ...
+            ...
+            >>> class NetB(mindspore.nn.Cell):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.register_buffer("buffer_b", mindspore.tensor([1, 2, 3]))
+            ...
+            ...     def construct(self, x):
+            ...         return x + self.buffer_b
+            ...
+            ...
+            >>> class NetA(mindspore.nn.Cell):
+            ...     def __init__(self, net_b):
+            ...         super().__init__()
+            ...         self.net_b = net_b
+            ...         self.register_buffer("buffer_a", mindspore.tensor([4, 5, 6]))
+            ...
+            ...     def construct(self, x):
+            ...         return self.net_b(x) + self.buffer_a
+            ...
+            ...
+            >>> net_b = NetB()
+            >>> net_a = NetA(net_b)
+            >>>
+            >>> for buffer in net_a.buffers():
+            ...     print(f'buffer is {buffer}')
+            buffer is [4 5 6]
+            buffer is [1 2 3]
+
+        """
+        for _, buf in self.named_buffers(recurse=recurse):
+            yield buf
+
+    def _named_members(self, get_members_fn, prefix="", recurse=True, remove_duplicate: bool = True):
+        r"""Help yield various names + members of cells."""
+        memo = set()
+        cells = (
+            self.cells_and_names(name_prefix=prefix)
+            if recurse
+            else [(prefix, self)]
+        )
+        for cell_prefix, cell in cells:
+            members = get_members_fn(cell)
+            for k, v in members:
+                if v is None or v in memo:
+                    continue
+                if remove_duplicate:
+                    memo.add(v)
+                name = cell_prefix + ("." if cell_prefix else "") + k
+                yield name, v
+
+    @jit_forbidden_register
+    def get_sub_cell(self, target: str) -> "Cell":
+        """Return the sub cell given by `target` if it exists, otherwise throw an error.
+
+        For example, let's say you have an ``nn.Cell`` ``A`` that
+        looks like this:
+
+        .. code-block:: text
+
+            A(
+                (net_b): NetB(
+                    (net_c): NetC(
+                        (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
+                    )
+                    (dense): Dense(in_features=100, out_features=200, bias=True)
+                )
+            )
+
+        (The diagram shows an ``nn.Cell`` ``A``. ``A`` has a nested
+        sub cell ``net_b``, which itself has two sub cells ``net_c``
+        and ``dense``. ``net_c`` then has a sub cell ``conv``.)
+
+        To check whether we have the ``dense`` sub cell, we
+        would call `get_sub_cell("net_b.dense")`. To check whether
+        we have the ``conv`` sub cell, we would call
+        `get_sub_cell("net_b.net_c.conv")`.
+
+        The runtime of ``get_sub_cell`` is bounded by the degree
+        of cell nesting in `target`. A query against
+        `name_cells` achieves the same result, but it is O(N) in
+        the number of transitive cells. So, for a simple check to see
+        if some sub cells exist, ``get_sub_cell`` should always be
+        used.
+
+        Args:
+            target (str): The fully-qualified string name of the sub cell
+                to look for. (See above example for how to specify a
+                fully-qualified string.)
+
+        Returns:
+            Cell
+
+        Examples:
+            >>> import mindspore
+            ...
+            ...
+            >>> class NetC(mindspore.nn.Cell):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.register_buffer("buffer_c", mindspore.tensor([0, 0, 0]))
+            ...         self.dense_c = mindspore.nn.Dense(5, 3)
+            ...
+            ...     def construct(self, x):
+            ...         return self.dense_c(x) + self.buffer_c
+            ...
+            ...
+            >>> class NetB(mindspore.nn.Cell):
+            ...     def __init__(self, net_c):
+            ...         super().__init__()
+            ...         self.net_c = net_c
+            ...         self.register_buffer("buffer_b", mindspore.tensor([1, 2, 3]))
+            ...
+            ...     def construct(self, x):
+            ...         return self.net_c(x) + self.buffer_b
+            ...
+            ...
+            >>> class NetA(mindspore.nn.Cell):
+            ...     def __init__(self, net_b):
+            ...         super().__init__()
+            ...         self.net_b = net_b
+            ...         self.register_buffer("buffer_a", mindspore.tensor([4, 5, 6]))
+            ...
+            ...     def construct(self, x):
+            ...         return self.net_b(x) + self.buffer_a
+            ...
+            ...
+            >>> net_c = NetC()
+            >>> net_b = NetB(net_c)
+            >>> net_a = NetA(net_b)
+            >>> net_c = net_a.get_sub_cell("net_b.net_c")
+            >>> print(f'net_c is {net_c}')
+            net_c is NetC(
+                (dense_c): Dense(input_channels=5, output_channels=3, has_bias=True)
+            )
+
+        """
+        if target == "":
+            return self
+
+        atoms: List[str] = target.split(".")
+        cell = self
+
+        for item in atoms:
+            if not hasattr(cell, item):
+                raise AttributeError(
+                    cell._get_name() + " has no " "attribute `" + item + "`"
+                )
+
+            cell = getattr(cell, item)
+
+            if not isinstance(cell, Cell):
+                raise AttributeError("`" + item + "` is not " "an nn.Cell")
+
+        return cell
+
     def get_func_graph_proto(self):
         """Return graph binary proto."""
         exec_id = ".".join([self.phase, str(self.create_time), str(id(self))])
@@ -405,6 +845,10 @@ class Cell(Cell_):
             params = self.__dict__['_params']
             if name in params:
                 return params[name]
+        if '_buffers' in self.__dict__:
+            buffers = self.__dict__['_buffers']
+            if name in buffers:
+                return buffers[name]
         if '_cells' in self.__dict__:
             cells = self.__dict__['_cells']
             if name in cells:
@@ -427,6 +871,8 @@ class Cell(Cell_):
     def __delattr__(self, name):
         if name in self._params:
             del self._params[name]
+        elif name in self._buffers:
+            del self._buffers[name]
         elif name in self._cells:
             del self._cells[name]
         elif '_params_list' in self.__dict__ and name in self._params_list:
@@ -600,6 +1046,89 @@ class Cell(Cell_):
         for prim in all_prims:
             prim.add_prim_attr("strategy_gen_mode", "data_parallel")
 
+    def offload(self, backward_prefetch="Auto"):
+        """
+        Set the cell offload. All primitive ops in the cell will be set offload. For the intermediate
+        activations calculated by these primitive ops, we will not save them in the forward pass, but
+        offload them and onload them in the backward pass.
+
+        Note:
+            - If Cell.offload is called, the mode should be set to "GRAPH_MODE".
+            - If Cell.offload is called, lazyinline should be enabled.
+
+        Args:
+            backward_prefetch(Union[str, int], optional): The timing for prefetching activations in advance in backward
+                                                          pass. Default: ``"Auto"``. If set it to ``"Auto"``, framework
+                                                          will start to prefetch activations one operator in advance.
+                                                          If set it to a positive int value, framework will start to
+                                                          prefetch activations ``backward_prefetch`` operators in
+                                                          advance, such as 1, 20, 100.
+        Examples:
+            >>> import mindspore.nn as nn
+            >>> from mindspore import ops
+            >>> from mindspore.common import Tensor, Parameter
+            >>> from mindspore.common.lazy_inline import lazy_inline
+            >>>
+            >>> class Block(nn.Cell):
+            ...     def __init__(self):
+            ...         super(Block, self).__init__()
+            ...         self.transpose1 = ops.Transpose()
+            ...         self.transpose2 = ops.Transpose()
+            ...         self.transpose3 = ops.Transpose()
+            ...         self.transpose4 = ops.Transpose()
+            ...         self.real_div1 = ops.RealDiv()
+            ...         self.real_div2 = ops.RealDiv()
+            ...         self.batch_matmul1 = ops.BatchMatMul()
+            ...         self.batch_matmul2 = ops.BatchMatMul()
+            ...         self.softmax = ops.Softmax(-1)
+            ...         self.expand_dims = ops.ExpandDims()
+            ...         self.sub = ops.Sub()
+            ...         self.y = Parameter(Tensor(np.ones((1024, 128, 128)).astype(np.float32)))
+            ...     def construct(self, x):
+            ...         transpose1 = self.transpose1(x, (0, 2, 1, 3))
+            ...         real_div1 = self.real_div1(transpose1, Tensor(2.37891))
+            ...         transpose2 = self.transpose2(x, (0, 2, 3, 1))
+            ...         real_div2 = self.real_div2(transpose2, Tensor(2.37891))
+            ...         batch_matmul1 = self.batch_matmul1(real_div1, real_div2)
+            ...         expand_dims = self.expand_dims(self.y, 1)
+            ...         sub = self.sub(Tensor([1.0]), expand_dims)
+            ...         soft_max = self.softmax(sub)
+            ...         transpose3 = self.transpose3(x, (0, 2, 1, 3))
+            ...         batch_matmul2 = self.batch_matmul2(soft_max[0], transpose3)
+            ...         transpose4 = self.transpose4(batch_matmul2, (0, 2, 1, 3))
+            ...         return transpose4
+            >>>
+            >>> class OuterBlock(nn.Cell):
+            ...     @lazy_inline
+            ...     def __init__(self):
+            ...         super(OuterBlock, self).__init__()
+            ...         self.block = Block()
+            ...     def construct(self, x):
+            ...         return self.block(x)
+            >>>
+            >>> class Nets(nn.Cell):
+            ...     def __init__(self):
+            ...         super(Nets, self).__init__()
+            ...         self.blocks = nn.CellList()
+            ...         for _ in range(3):
+            ...             b = OuterBlock()
+            ...             b.offload()
+            ...             self.blocks.append(b)
+            ...     def construct(self, x):
+            ...         out = x
+            ...         for i in range(3):
+            ...             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:
+            Validator.check_non_negative_int(backward_prefetch)
+        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):
         """
         Defining the input and output layouts of this cell and the parallel strategies of remaining ops will be
@@ -612,9 +1141,9 @@ class Cell(Cell_):
         The parallel strategies of remaining operators are derived from the strategy specified by the input and output.
 
         Note:
-            If Cell.shard is called, the parallel mode in `set_auto_parallel_context` (parallel_mode) will be set to
-            "auto_parallel" and the search mode (search_mode) to "sharding_propagation".
-            If the input contain Parameter, its strategy should be set in `in_strategy`.
+            - It is valid only in semi auto parallel or auto parallel mode.
+              In other parallel modes, strategies set here will be ignored.
+            - If the input contain Parameter, its strategy should be set in `in_strategy`.
 
         Args:
             in_strategy (tuple): Define the layout of inputs, each element of the tuple should be a tuple. Tuple
@@ -628,7 +1157,7 @@ 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 (string): Select a certain device target. It is not in use right now.
+            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
@@ -660,10 +1189,8 @@ class Cell(Cell_):
             ...     x = self.block2_shard(x)
             ...     return x
         """
-        if context.get_auto_parallel_context("parallel_mode") not in ["auto_parallel", "semi_auto_parallel"]:
-            raise AssertionError(f"Cell shard only supports auto parallel or semi_auto_parallel "
-                                 f"Please check the parallel mode in parallel context.")
-
+        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
@@ -766,7 +1293,8 @@ class Cell(Cell_):
         """
         Process cell info before call construct
         """
-        if self.requires_grad:
+        if self.requires_grad and (not _pynative_executor.grad_flag() or _pynative_executor.high_order()):
+            self.is_top_cell = True
             _pynative_executor.set_grad_flag(True)
             _pynative_executor.new_graph(self, *args, **kwargs)
         elif self._dynamic_shape_inputs is not None:
@@ -780,8 +1308,9 @@ class Cell(Cell_):
         """
         Process cell info after call construct
         """
-        if self.requires_grad:
+        if self.requires_grad and self.is_top_cell:
             _pynative_executor.end_graph(self, output, *args, **kwargs)
+            self.is_top_cell = False
         elif self._dynamic_shape_inputs is not None:
             _pynative_executor.set_cell_use_dynamic_shape_process(False)
 
@@ -826,52 +1355,41 @@ class Cell(Cell_):
             self._add_attr(key, value)
         self._attr_synced = True
 
-    def _set_attr_for_parameter(self, name, value):
-        """Set attr for parameter."""
-        cells = self.__dict__.get('_cells')
-        params = self.__dict__.get('_params')
-        if params is None:
-            raise AttributeError("For 'Cell', can not assign params before Cell.__init__() is called.")
-        if name in self.__dict__:
-            if self.__dict__[name] is not None:
-                raise TypeError(f"For 'Cell', the {name} should not be Parameter.")
-            del self.__dict__[name]
-        if cells and name in cells:
-            raise TypeError(f"For 'Cell', the {name} must be Cell, but got Parameter.")
-        self.insert_param_to_cell(name, value)
-
-    def _set_attr_for_parameter_tuple(self, name, value):
-        """Set attr for parameter in ParameterTuple."""
-        params = self.__dict__.get('_params')
-        params_list = self.__dict__.get('_params_list')
-        if params is None:
-            raise AttributeError("For 'Cell', can not assign params before Cell.__init__() is called.")
-        exist_names = set("")
-        exist_objs = set()
-        for item in value:
-            if item in exist_objs:
-                # 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)
-                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)
-
-        if context._get_mode() == context.PYNATIVE_MODE:
+    def _set_attr_for_param_or_param_tuple(self, name, value):
+        """Set attr for param and tensor."""
+        if isinstance(value, Parameter):
             if name in self.__dict__:
                 del self.__dict__[name]
-            if name in params:
-                del params[name]
-            params_list[name] = value
-        else:
-            object.__setattr__(self, name, value)
+            self.insert_param_to_cell(name, value)
+        elif isinstance(value, ParameterTuple):
+            exist_names = set("")
+            exist_objs = set()
+            for item in value:
+                if item in exist_objs:
+                    # 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)
+                    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)
+
+            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)
 
     def _set_attr_for_parameter_in_list_or_tuple(self, name, value):
         """Set attr for parameter in list or tuple."""
@@ -884,24 +1402,18 @@ class Cell(Cell_):
                 item.name = item.name + "$" + str(self._id)
                 self._id += 1
             if item.name in self.exist_names:
-                raise ValueError("The value {} , its name '{}' already exists. "
-                                 "Please set a unique name for the parameter.".format(value, item.name))
+                raise ValueError(f"The value {value} , its name '{item.name}' already exists. "
+                                 "Please set a unique name for the parameter.")
             self.exist_names.add(item.name)
         object.__setattr__(self, name, value)
 
     def _set_attr_for_cell(self, name, value):
         """Set attr for cell."""
-        cells = self.__dict__.get('_cells')
-        params = self.__dict__.get('_params')
-        if cells is None:
-            raise AttributeError("For 'Cell', can not assign cells before Cell.__init__() is called.")
         if name in self.__dict__:
             del self.__dict__[name]
-        if params and name in params:
-            raise TypeError(f"For 'Cell', the {name} must be Parameter, but got Cell.")
         if self._auto_prefix:
             value.update_parameters_name(name + '.')
-        cells[name] = value
+        self.insert_child_to_cell(name, value)
         if hasattr(self, '_cell_init_args'):
             self.cell_init_args += str({name: value})
 
@@ -914,30 +1426,57 @@ class Cell(Cell_):
         else:
             self.insert_param_to_cell(name, None)
 
-    def __setattr__(self, name, value):
-        cells = self.__dict__.get('_cells')
+    def _set_attr_for_object(self, name, value):
+        """Set attr for py object."""
         params = self.__dict__.get('_params')
-        if isinstance(value, Parameter):
-            self._set_attr_for_parameter(name, value)
-        elif isinstance(value, ParameterTuple):
-            self._set_attr_for_parameter_tuple(name, value)
-        elif isinstance(value, (list, tuple)) and value and _check_param_list_tuple(value):
+        if params is not None and name in params:
+            if value is not None:
+                if isinstance(value, Tensor):
+                    params[name].set_data(value)
+                    return
+                raise TypeError(
+                    f"Parameter '{name}' already exists in network, "
+                    f"can not assign this type: '{type(value)}' as a parameter.")
+            params[name] = None
+            return
+        cells = self.__dict__.get('_cells')
+        if cells is not None and name in cells:
+            if value is not None:
+                raise TypeError(
+                    f"Sub cell '{name}' already exists in network, "
+                    f"can not assign this type: '{type(value)}' as a cell.")
+            cells[name] = None
+            return
+        buffers = self.__dict__.get('_buffers')
+        if buffers is not None and name in buffers:
+            if value is not None:
+                raise TypeError(
+                    f"Buffer '{name}' already exists in network, "
+                    f"can not assign this type: '{type(value)}' as a buffer.")
+            buffers[name] = None
+            return
+        object.__setattr__(self, name, value)
+
+    def __setattr__(self, name, value):
+        if isinstance(value, (Parameter, ParameterTuple)):
+            self._set_attr_for_param_or_param_tuple(name, value)
+        elif _is_parameter_list_or_tuple(value):
             self._set_attr_for_parameter_in_list_or_tuple(name, value)
         elif isinstance(value, Cell):
             self._set_attr_for_cell(name, value)
-        elif params and name in params:
-            self._set_attr_for_params(name, value)
-        elif cells and name in cells:
-            if value is not None:
-                raise TypeError(f"For 'Cell', the type of {name} must be cell, but got {type(value).__name__}.")
-            self._cells[name] = None
-        else:
-            if isinstance(value, Primitive):
-                value.set_prim_instance_name(name)
-                self._primitives[name] = value
+        elif isinstance(value, _Buffer):
+            if name in self.__dict__:
+                del self.__dict__[name]
+            self.register_buffer(name, value)
+        elif isinstance(value, Primitive):
+            value.set_prim_instance_name(name)
+            self._primitives[name] = value
             object.__setattr__(self, name, value)
-        if name not in Cell.IGNORE_LIST:
-            self._attr_synced = False
+        else:
+            self._set_attr_for_object(name, value)
+
+    def _get_name(self):
+        return self.__class__.__name__
 
     def extend_repr(self):
         """
@@ -951,19 +1490,28 @@ class Cell(Cell_):
         return self.__repr__()
 
     def __repr__(self):
-        extra_str = self.extend_repr()
-        info_str = self.__class__.__name__ + '<'
-        if self._cells:
-            sub_str = '\n'
-            if extra_str:
-                sub_str += '{}\n'.format(self.extend_repr())
-            for key, value in self._cells.items():
-                sub_str += '({}): {}\n'.format(key, repr(value))
-            sub_str = sub_str.replace('\n', '\n  ') + '>'
-            info_str += sub_str
-        else:
-            info_str += extra_str + '>'
-        return info_str
+        extra_lines = []
+        extend_repr = self.extend_repr()
+        # empty string will be split into list ['']
+        if extend_repr:
+            extra_lines = extend_repr.split("\n")
+        child_lines = []
+        for key, cell in self._cells.items():
+            cell_str = repr(cell)
+            cell_str = _addindent(cell_str, 2)
+            child_lines.append("(" + key + "): " + cell_str)
+        lines = extra_lines + child_lines
+
+        main_str = self._get_name() + "("
+        if lines:
+            # simple one-liner info, which most builtin Modules will use
+            if len(extra_lines) == 1 and not child_lines:
+                main_str += extra_lines[0]
+            else:
+                main_str += "\n  " + "\n  ".join(lines) + "\n"
+
+        main_str += ")"
+        return main_str
 
     def load_parameter_slice(self, params):
         """
@@ -1129,9 +1677,11 @@ class Cell(Cell_):
             args (tuple): Args of the Cell object.
             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,
                                      jit_config_dict=self._jit_config_dict, **kwargs)
+        _clear_auto_parallel_context(self)
 
     def compile_and_run(self, *args, **kwargs):
         """
@@ -1262,9 +1812,9 @@ class Cell(Cell_):
             >>> net2 = nn.Dense(2, 2)
             >>> net1.insert_child_to_cell("child", net2)
             >>> print(net1)
-            ReLU<
-              (child): Dense<input_channels=2, output_channels=2, has_bias=True>
-              >
+            ReLU(
+              (child): Dense(input_channels=2, output_channels=2, has_bias=True)
+            )
         """
         if not isinstance(child_name, str):
             raise TypeError(f"For 'insert_child_to_cell', the type of parameter 'child_name' must be str, "
@@ -1322,13 +1872,22 @@ class Cell(Cell_):
                         new_param_tuple.append(param)
                     cell.__dict__[key] = ParameterTuple(new_param_tuple)
 
+    def _get_cell_parallel_mode(self):
+        """Determine whether the current cell is in parallel mode."""
+        is_parallel_mode = False
+        for _, param in self.parameters_and_names():
+            if param.param_info.is_param_init:
+                is_parallel_mode = True
+                break
+        return is_parallel_mode
+
     def init_parameters_data(self, auto_parallel_mode=False):
         """
         Initialize all parameters and replace the original saved parameters in cell.
 
         Note:
             trainable_params() and other similar interfaces may return different parameter instance after
-            `init_parameters_data`, do not save these results.
+            `init_parameters_data`. It is not recommended to save these results.
 
         Args:
             auto_parallel_mode (bool): If running in auto_parallel_mode. Default: ``False`` .
@@ -1366,9 +1925,18 @@ 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:
+                    continue
                 if not auto_parallel_mode:
                     cell._params[param_name] = _updata(param)
                     continue
@@ -1380,6 +1948,12 @@ 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:
+                            continue
                         if not auto_parallel_mode:
                             new_param_tuple.append(_updata(param))
                             continue
@@ -1687,7 +2261,7 @@ class Cell(Cell_):
             ...         return x
             >>> net = Net()
             >>> print(net.cells())
-            odict_values([Dense<input_channels=2, output_channels=2, has_bias=True>])
+            odict_values([Dense(input_channels=2, output_channels=2, has_bias=True)])
         """
         return self.name_cells().values()
 
@@ -1748,7 +2322,7 @@ class Cell(Cell_):
             ...         return x
             >>> net = Net()
             >>> print(net.name_cells())
-            OrderedDict([('dense', Dense<input_channels=2, output_channels=2, has_bias=True>)])
+            OrderedDict([('dense', Dense(input_channels=2, output_channels=2, has_bias=True))])
         """
         value_set = set()
         cells = OrderedDict()
@@ -1789,10 +2363,10 @@ class Cell(Cell_):
             ...     if isinstance(cell, nn.Dense):
             ...         cell.weight.set_data(initializer(One(), cell.weight.shape, cell.weight.dtype))
             >>> net.apply(func)
-            SequentialCell<
-              (0): Dense<input_channels=2, output_channels=2, has_bias=True>
-              (1): Dense<input_channels=2, output_channels=2, has_bias=True>
-              >
+            SequentialCell(
+              (0): Dense(input_channels=2, output_channels=2, has_bias=True)
+              (1): Dense(input_channels=2, output_channels=2, has_bias=True)
+            )
             >>> print(net[0].weight.asnumpy())
             [[1. 1.]
              [1. 1.]]
@@ -1832,9 +2406,6 @@ class Cell(Cell_):
         if not hasattr(self, "_func_graph_flags"):
             self._func_graph_flags = {}
         self._func_graph_flags.update({**flags})
-        if context._get_mode() == context.PYNATIVE_MODE and self._func_graph_flags.get("output_no_recompute"):
-            raise TypeError("Recompute is not supported in PyNative mode currently, you can use "
-                            "'context.set_context(mode=context.GRAPH_MODE)' or @jit to set graph mode.")
         self.__dict__.update({**flags})
         self._add_mixed_precision_flag(**flags)
         return self
@@ -1927,8 +2498,8 @@ class Cell(Cell_):
             >>>
             >>> net = nn.Conv2d(120, 240, 4, has_bias=False, weight_init='normal')
             >>> net.to_float(mstype.float16)
-            Conv2d<input_channels=120, output_channels=240, kernel_size=(4, 4), stride=(1, 1), pad_mode=same,
-            padding=0, dilation=(1, 1), group=1, has_bias=False, weight_init=normal, bias_init=None, format=NCHW>
+            Conv2d(input_channels=120, output_channels=240, kernel_size=(4, 4), stride=(1, 1), pad_mode=same,
+            padding=0, dilation=(1, 1), group=1, has_bias=False, weight_init=normal, bias_init=None, format=NCHW)
         """
         if dst_type not in (mstype.float16, mstype.float32, mstype.bfloat16):
             raise ValueError("For 'to_float', the argument 'dst_type' must be mstype.float32, mstype.float16 or "
@@ -2133,8 +2704,7 @@ class Cell(Cell_):
         """
         if context._get_mode() == context.GRAPH_MODE:
             return HookHandle()
-        if not check_hook_fn("register_forward_pre_hook", hook_fn):
-            return HookHandle()
+        check_hook_fn(hook_fn)
         handle = HookHandle(self._forward_pre_hook)
         self._forward_pre_hook[handle.handle_id] = hook_fn
         return handle
@@ -2233,8 +2803,7 @@ class Cell(Cell_):
             return HookHandle()
         if context._get_mode() == context.GRAPH_MODE:
             return HookHandle()
-        if not check_hook_fn("register_forward_hook", hook_fn):
-            return HookHandle()
+        check_hook_fn(hook_fn)
         handle = HookHandle(self._forward_hook)
         self._forward_hook[handle.handle_id] = hook_fn
         return handle
@@ -2324,8 +2893,7 @@ class Cell(Cell_):
         """
         if context._get_mode() == context.GRAPH_MODE:
             return HookHandle()
-        if not check_hook_fn("register_backward_pre_hook", hook_fn):
-            return HookHandle()
+        check_hook_fn(hook_fn)
         handle = HookHandle(self._backward_pre_hook)
         self._backward_pre_hook[handle.handle_id] = hook_fn
         if self._cell_backward_pre_hook is None:
@@ -2361,6 +2929,527 @@ class Cell(Cell_):
                         len(ret), len(outputs)))
         return ret
 
+    def get_extra_state(self) -> Any:
+        """Return any extra state to include in the cell's state_dict.
+
+        This function is called from ``state_dict``.
+        Implement this and a corresponding ``set_extra_state`` for your cell
+        if you need to store extra state.
+
+        Note that extra state should be picklable to ensure working serialization
+        of the state_dict. Only provide backwards compatibility guarantees
+        for serializing tensors; other objects may break backwards compatibility if
+        their serialized pickled form changes.
+
+        Returns:
+            object, any extra state to store in the cell's state_dict.
+        """
+        raise RuntimeError(
+            "Reached a code path in Cell.get_extra_state() that should never be called."
+
+        )
+
+    def set_extra_state(self, state: Any) -> None:
+        """Set extra state contained in the loaded `state_dict`.
+
+        This function is called from `load_state_dict` to handle any extra state
+        found within the `state_dict`. Implement this function and a corresponding
+        `get_extra_state` for your cell if you need to store extra state within its
+        `state_dict`.
+
+        Args:
+            state (dict): Extra state from the `state_dict`.
+        """
+        raise RuntimeError(
+            "Reached a code path in Cell.set_extra_state() that should never be called."
+        )
+
+    @jit_forbidden_register
+    def register_state_dict_post_hook(self, hook):
+        r"""Register a post-hook for the :func:`mindspore.nn.Cell.state_dict` method.
+
+        It should have the following signature:
+
+        hook(cell, state_dict, prefix, local_metadata) -> None
+
+        The registered hooks can modify the ``state_dict`` inplace.
+
+        Args:
+            hook (Callable): The hook function after `state_dict` is called.
+
+        Returns:
+            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
+        return handle
+
+    @jit_forbidden_register
+    def register_state_dict_pre_hook(self, hook):
+        r"""Register a pre-hook for the :func:`mindspore.nn.Cell.state_dict` method.
+
+        It should have the following signature:
+
+        hook(cell, prefix, keep_vars) -> None
+
+        The registered hooks can be used to perform pre-processing before the `state_dict`
+        call is made.
+
+        Args:
+            hook (Callable): The hook function before `state_dict` is called.
+
+        Returns:
+            A handle that can be used to remove the added hook by calling
+            `handle.remove()`.
+
+        Examples:
+            >>> import mindspore
+            ...
+            ...
+            >>> class NetA(mindspore.nn.Cell):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.register_buffer("buffer_a", mindspore.tensor([1, 2, 3]))
+            ...         self.param_a = mindspore.Parameter(mindspore.tensor([1, 2, 3]))
+            ...
+            ...     def construct(self, x):
+            ...         return x + self.buffer_a + self.param_a
+            ...
+            ...
+            >>> def _add_extra_param(cell, prefix, keep_vars):
+            ...     cell._params["extra_param"] = mindspore.Parameter(mindspore.tensor([4, 5, 6]))
+            ...
+            ...
+            >>> net = NetA()
+            >>> handle = net.register_state_dict_pre_hook(_add_extra_param)
+            >>> net_state_dict = net.state_dict()
+            >>> handle.remove()
+            >>> 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
+        return handle
+
+    def _save_to_state_dict(self, destination, prefix, keep_vars):
+        r"""Save cell state to the `destination` dictionary.
+
+        The `destination` dictionary will contain the state
+        of the cell, but not its descendants. This is called on every
+        sub cell in :func:`mindspore.nn.Cell.state_dict`.
+
+        In rare cases, subclasses can achieve class-specific behavior by
+        overriding this method with custom logic.
+
+        Args:
+            destination (dict): a dict where state will be stored
+            prefix (str): the prefix for parameters and buffers used in this
+                cell
+        """
+        for name, param in self._params.items():
+            if param is not None:
+                destination[prefix + name] = param
+        for name, buf in self._buffers.items():
+            if buf is not None and name not in self._non_persistent_buffers_set:
+                destination[prefix + name] = buf
+        extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX
+        if (
+                getattr(self.__class__, "get_extra_state", Cell.get_extra_state)
+                is not Cell.get_extra_state
+        ):
+            destination[extra_state_key] = self.get_extra_state()
+
+    # The user can pass an optional arbitrary mappable object to `state_dict`, in which case `state_dict` returns
+    # back that same object. But if they pass nothing, an `OrderedDict` is created and returned.
+    T_destination = TypeVar("T_destination", bound=Dict[str, Any])
+
+    @jit_forbidden_register
+    def state_dict(self, *args, destination=None, prefix="", keep_vars=False):
+        r"""Return a dictionary containing references to the whole state of the cell.
+
+        Both parameters and persistent buffers (e.g. running averages) are
+        included. Keys are corresponding parameter and buffer names.
+        Parameters and buffers set to ``None`` are not included.
+
+        .. note::
+            The returned object is a shallow copy. It contains references
+            to the cell's parameters and buffers.
+
+        .. warning::
+            - Currently ``state_dict()`` also accepts positional arguments for
+              ``destination``, ``prefix`` and ``keep_vars`` in order. However,
+              this is being deprecated and keyword arguments will be enforced in
+              future releases.
+
+            - Please avoid the use of argument ``destination`` as it is not
+              designed for end-users.
+
+        Args:
+            destination (dict, optional): If provided, the state of cell will
+                be updated into the dict and the same object is returned.
+                Otherwise, an ``OrderedDict`` will be created and returned.
+                Default: ``None``.
+            prefix (str, optional): A prefix added to parameter and buffer
+                names to compose the keys in state_dict. Default: ``''``.
+            keep_vars (bool, optional): Whether the state_dict returns a copy. Default: ``False`` , returns a reference.
+
+        Returns:
+            Dict, a dictionary containing a whole state of the cell.
+
+        Examples:
+            >>> import mindspore
+            >>> class Model(mindspore.nn.Cell):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.register_buffer("buffer_a", mindspore.tensor([4, 5, 6]))
+            ...         self.param_a = mindspore.Parameter(mindspore.tensor([1, 2, 3]))
+            ...
+            ...     def construct(self, x):
+            ...         return x + self.buffer_a + self.param_a
+            ...
+            ...
+            >>> model = Model()
+            >>> print(model.state_dict())
+            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(
+                "Positional args are being deprecated, use kwargs instead. Refer to "
+                "https://www.mindspore.cn/docs/zh-CN/master/api_python/nn/mindspore.nn.Cell.html"
+                " for details.",
+                FutureWarning,
+                stacklevel=2,
+            )
+            if destination is None:
+                destination = args[0]
+            if len(args) > 1 and prefix == "":
+                prefix = args[1]
+            if len(args) > 2 and keep_vars is False:
+                keep_vars = args[2]
+        if destination is not None and not isinstance(destination, dict):
+            raise TypeError(f"The type of destination must be OrderedDict, but got {type(destination)}")
+        if not isinstance(prefix, str):
+            raise TypeError(f"The type of prefix must be string, but got {type(prefix)}")
+        if not isinstance(keep_vars, bool):
+            raise TypeError(f"The type of keep_vars must be bool, but got {type(keep_vars)}")
+
+        if destination is None:
+            destination = OrderedDict()
+            destination._metadata = OrderedDict()
+
+        local_metadata = {}
+        if hasattr(destination, "_metadata"):
+            destination._metadata[prefix[:-1]] = local_metadata
+
+        for hook in self._state_dict_pre_hooks.values():
+            hook(self, prefix, keep_vars)
+        self._save_to_state_dict(destination, prefix, keep_vars)
+        for name, cell in self._cells.items():
+            if cell is not None:
+                cell.state_dict(
+                    destination=destination,
+                    prefix=prefix + name + ".",
+                    keep_vars=keep_vars,
+                )
+        for hook in self._state_dict_hooks.values():
+            hook_result = hook(self, destination, prefix, local_metadata)
+            if hook_result is not None:
+                raise RuntimeError("state_dict post-hook must return None")
+        return destination
+
+    @jit_forbidden_register
+    def register_load_state_dict_pre_hook(self, hook):
+        r"""Register a pre-hook to be run before cell's :func:`mindspore.nn.Cell.load_state_dict` is called.
+
+        It should have the following signature:
+
+        hook(cell, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None  # noqa: B950
+
+        Args:
+            hook (Callable): The hook function before `load_state_dict` is called.
+
+        Returns:
+            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
+        return handle
+
+    @jit_forbidden_register
+    def register_load_state_dict_post_hook(self, hook):
+        r"""Register a post-hook to be run after cell's :func:`mindspore.nn.Cell.load_state_dict` is called.
+
+        It should have the following signature:
+
+        hook(cell, incompatible_keys) -> None
+
+        The ``cell`` argument is the current cell that this hook is registered
+        on, and the ``incompatible_keys`` argument is a ``NamedTuple`` consisting
+        of attributes ``missing_keys`` and ``unexpected_keys``. ``missing_keys``
+        is a ``list`` of ``str`` containing the missing keys and
+        ``unexpected_keys`` is a ``list`` of ``str`` containing the unexpected keys.
+
+        The given incompatible_keys can be modified inplace if needed.
+
+        Note that the checks performed when calling :func:`load_state_dict` with
+        ``strict=True`` are affected by modifications the hook makes to
+        ``missing_keys`` or ``unexpected_keys``, as expected. Additions to either
+        set of keys will result in an error being thrown when ``strict=True``, and
+        clearing out both missing and unexpected keys will avoid an error.
+
+        Args:
+            hook (Callable): The hook function after `load_state_dict` is called.
+
+        Returns:
+            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
+        return handle
+
+    def _load_from_state_dict(
+            self,
+            state_dict,
+            prefix,
+            local_metadata,
+            strict,
+            missing_keys,
+            unexpected_keys,
+            error_msgs,
+    ):
+        r"""Copy parameters and buffers from :attr:`state_dict` into only this cell, but not its descendants.
+
+        This is called on every sub cell
+        in :func:`mindspore.nn.Cell.load_state_dict`. Metadata saved for this
+        cell in input :attr:`state_dict` is provided as :attr:`local_metadata`.
+        For state dicts without metadata, :attr:`local_metadata` is empty.
+        Subclasses can achieve class-specific backward compatible loading using
+        the version number at `local_metadata.get("version", None)`.
+
+        .. note::
+            :attr:`state_dict` is not the same object as the input
+            :attr:`state_dict` to :func:`mindspore.nn.Cell.load_state_dict`. So
+            it can be modified.
+
+        Args:
+            state_dict (dict): a dict containing parameters and
+                persistent buffers.
+            prefix (str): the prefix for parameters and buffers used in this
+                cell
+            local_metadata (dict): a dict containing the metadata for this cell.
+                See
+            strict (bool): whether to strictly enforce that the keys in
+                :attr:`state_dict` with :attr:`prefix` match the names of
+                parameters and buffers in this cell
+            missing_keys (list of str): if ``strict=True``, add missing keys to
+                this list
+            unexpected_keys (list of str): if ``strict=True``, add unexpected
+                keys to this list
+            error_msgs (list of str): error messages should be added to this
+                list, and will be reported together in
+                :func:`mindspore.nn.Cell.load_state_dict`
+        """
+        for hook in self._load_state_dict_pre_hooks.values():
+            hook(
+                self,
+                state_dict,
+                prefix,
+                local_metadata,
+                strict,
+                missing_keys,
+                unexpected_keys,
+                error_msgs,
+            )
+
+        persistent_buffers = {
+            k: v
+            for k, v in self._buffers.items()
+            if k not in self._non_persistent_buffers_set
+        }
+        local_name_params = itertools.chain(
+            self._params.items(), persistent_buffers.items()
+        )
+        local_state = {k: v for k, v in local_name_params if v is not None}
+
+        for name, param in local_state.items():
+            key = prefix + name
+            if key in state_dict:
+                input_param = state_dict[key]
+                if not isinstance(input_param, Tensor):
+                    error_msgs.append(
+                        f'While copying the parameter named "{key}", '
+                        "expected Tensor or Tensor-like object from checkpoint but "
+                        f"received {type(input_param)}"
+                    )
+                    continue
+
+                if input_param.shape != param.shape:
+                    # local shape should match the one in checkpoint
+                    error_msgs.append(
+                        f"size mismatch for {key}: copying a param with shape {input_param.shape} from checkpoint, "
+                        f"the shape in current model is {param.shape}."
+                    )
+                    continue
+                try:
+                    param.assign_value(Tensor(input_param.asnumpy(), dtype=param.dtype))
+                except Exception as ex:  # pylint: disable=W0703
+                    error_msgs.append(
+                        f'While copy the parameter named "{key}", '
+                        f"whose shape in the model are {param.shape} and "
+                        f"whose shape in the checkpoint are {input_param.shape}, "
+                        f"an exception occurred : {ex.args}."
+                    )
+            elif strict:
+                missing_keys.append(key)
+
+        extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX
+        if getattr(self.__class__, "set_extra_state", Cell.set_extra_state) is not Cell.set_extra_state:
+            if extra_state_key in state_dict:
+                self.set_extra_state(state_dict[extra_state_key])
+            elif strict:
+                missing_keys.append(extra_state_key)
+        elif strict and (extra_state_key in state_dict):
+            unexpected_keys.append(extra_state_key)
+
+        if strict:
+            for key in state_dict.keys():
+                if key.startswith(prefix) and key != extra_state_key:
+                    input_name = key[len(prefix):].split(".", 1)
+                    # Must be cell if it have attributes
+                    if len(input_name) > 1:
+                        if input_name[0] not in self._cells:
+                            unexpected_keys.append(key)
+                    elif input_name[0] not in local_state:
+                        unexpected_keys.append(key)
+
+    @jit_forbidden_register
+    def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True):
+        r"""Copy parameters and buffers from :attr:`state_dict` into this cell and its descendants.
+
+        If :attr:`strict` is ``True``, then
+        the keys of :attr:`state_dict` must exactly match the keys returned
+        by this cell's :func:`mindspore.nn.Cell.state_dict` function.
+
+        Args:
+            state_dict (dict): A dict containing parameters and
+                persistent buffers.
+            strict (bool, optional): Whether to strictly enforce that the keys
+                in input `state_dict` match the keys returned by this cell's
+                :func:`mindspore.nn.Cell.state_dict` function. Default ``True`` .
+
+        Returns:
+            A namedtuple with ``missing_keys`` and ``unexpected_keys`` fields,
+
+            - `missing_keys` is a list of str containing any keys that are expected
+              by this cell but missing from the provided ``state_dict``.
+
+            - `unexpected_keys` is a list of str containing the keys that are not
+              expected by this cell but present in the provided ``state_dict``.
+
+        Note:
+            If `strict` is ``True`` and a parameter or buffer is registered as ``None``, but its corresponding key
+            exists in :attr:`state_dict`, and :func:`mindspore.nn.Cell.load_state_dict` will raise a ``RuntimeError``.
+
+        Examples:
+            >>> import mindspore
+            >>> import os
+            >>> class Model(mindspore.nn.Cell):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.register_buffer("buffer_a", mindspore.tensor([4, 5, 6]))
+            ...         self.param_a = mindspore.Parameter(mindspore.tensor([1, 2, 3]))
+            ...
+            ...     def construct(self, x):
+            ...         return x + self.buffer_a + self.param_a
+            ...
+            ...
+            >>> model = Model()
+            >>> print(model.state_dict())
+            >>> mindspore.save_checkpoint(model.state_dict(), './model_state_dict_ckpt')
+            >>> new_model = Model()
+            >>> new_model.load_state_dict(mindspore.load_checkpoint('./model_state_dict_ckpt'))
+            >>> print(new_model.state_dict())
+            >>> os.remove('./model_state_dict_ckpt')
+            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]))])
+            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]))])
+        """
+        if not isinstance(state_dict, Mapping):
+            raise TypeError(
+                f"Expected state_dict to be dict-like, got {type(state_dict)}."
+            )
+
+        missing_keys: List[str] = []
+        unexpected_keys: List[str] = []
+        error_msgs: List[str] = []
+
+        # copy state_dict so _load_from_state_dict can modify it
+        metadata = getattr(state_dict, "_metadata", None)
+        state_dict = OrderedDict(state_dict)
+        if metadata is not None:
+            # mypy isn't aware that "_metadata" exists in state_dict
+            state_dict._metadata = metadata  # type: ignore[attr-defined]
+
+        def load(cell, local_state_dict, prefix=""):
+            local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
+            cell._load_from_state_dict(
+                local_state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs,
+            )
+            for name, child in cell._cells.items():
+                if child is not None:
+                    child_prefix = prefix + name + "."
+                    child_state_dict = {k: v for k, v in local_state_dict.items() if k.startswith(child_prefix)}
+                    load(child, child_state_dict, child_prefix)  # noqa: F821
+
+            # Note that the hook can modify missing_keys and unexpected_keys.
+            incompatible_keys = _IncompatibleKeys(missing_keys, unexpected_keys)
+            for hook in cell._load_state_dict_post_hooks.values():
+                out = hook(cell, incompatible_keys)
+                if out is not None:
+                    raise RuntimeError(
+                        "Hooks registered with ``register_load_state_dict_post_hook`` are not"
+                        "expected to return new values, if incompatible_keys need to be modified,"
+                        "it should be done inplace."
+                    )
+
+        load(self, state_dict)
+        del load
+
+        if strict:
+            if unexpected_keys:
+                error_msgs.insert(
+                    0,
+                    "Unexpected key(s) in state_dict: {}. ".format(
+                        ", ".join(f'"{k}"' for k in unexpected_keys)
+                    ),
+                )
+            if missing_keys:
+                error_msgs.insert(
+                    0,
+                    "Missing key(s) in state_dict: {}. ".format(
+                        ", ".join(f'"{k}"' for k in missing_keys)
+                    ),
+                )
+
+        if error_msgs:
+            raise RuntimeError(
+                "Error(s) in loading state_dict for {}:\n\t{}".format(
+                    self.__class__.__name__, "\n\t".join(error_msgs)
+                )
+            )
+        return _IncompatibleKeys(missing_keys, unexpected_keys)
+
     def register_backward_hook(self, hook_fn):
         """
         Register the backward hook function.
@@ -2420,8 +3509,7 @@ class Cell(Cell_):
         """
         if context._get_mode() == context.GRAPH_MODE:
             return HookHandle()
-        if not check_hook_fn("register_backward_hook", hook_fn):
-            return HookHandle()
+        check_hook_fn(hook_fn)
         handle = HookHandle(self._backward_hook)
         self._backward_hook[handle.handle_id] = hook_fn
         if self._cell_backward_hook is None:
@@ -2565,8 +3653,9 @@ class Cell(Cell_):
         if not self._has_config_recompute:
             self._has_config_recompute = True
         else:
-            raise RuntimeError("The recompute interface can be configured only once."
-                               " When the parent cell is configured, the child cell should not be configured")
+            logger.info("The recompute interface can be configured only once."
+                        " When the parent cell is configured, the child cell should not be configured")
+            return
         self._set_recompute_scope(mode)
         if mode and not output_recompute:
             self.add_flags(output_no_recompute=True)
@@ -2606,18 +3695,13 @@ class Cell(Cell_):
         """
         if context.get_context("mode") == context.PYNATIVE_MODE:
             self._recompute_cell = recompute_registry.get()(self.construct)
-            self._add_recompute_flag()
-            return
         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 (kwargs.get('parallel_optimizer_comm_recompute', False) and
-                    context.get_auto_parallel_context("pipeline_stages") > 1):
+            if kwargs.get('parallel_optimizer_comm_recompute', False):
                 logger.warning("Currently, the communication operator allgathers introduced by optimizer shard "
-                               "are not support recomputation in pipeline parallel.")
-            elif context.get_auto_parallel_context("pipeline_stages") == 1:
-                self._parallel_optimizer_comm_recompute(kwargs.get('parallel_optimizer_comm_recompute', False))
+                               "is replaced with zero3.")
         if 'recompute_slice_activation' in kwargs:
             self._recompute_slice_activation(kwargs.get('recompute_slice_activation', False))
 
@@ -2709,18 +3793,6 @@ class Cell(Cell_):
         if hasattr(network, "_amp_level"):
             self._amp_level = getattr(network, "_amp_level")
 
-    def _add_recompute_flag(self):
-        """
-        Set pynative cell recomputed.
-        """
-        if not self._has_config_recompute:
-            self._has_config_recompute = True
-        else:
-            logger.info("The recompute interface can be configured only once."
-                        " If the parent cell is configured, the child cell should not be configured")
-        for cell in self.cells():
-            cell._add_recompute_flag()
-
     def _register_parameters_hook(self, forward_hook=None, backward_hook=None, all=False):
         """
         Register the forward hook for parameters and register the backward hook for the corresponding gradient.
@@ -2807,6 +3879,7 @@ class Cell(Cell_):
                 cell._parameters_forward_hook = forward_hook
                 cell._parameters_backward_hook = backward_hook
 
+
 class GraphCell(Cell):
     """
     Base class for running the graph loaded from a MindIR.
@@ -2820,12 +3893,10 @@ class GraphCell(Cell):
             The key is the parameter name whose type is str, and the value is a Tensor or Parameter.
             If the parameter exists in the graph according to the name, update it's value.
             If the parameter does not exist, ignore it. Default: ``None`` .
-        obf_random_seed (Union[int, None]): The random seed used for dynamic obfuscation. "dynamic obfuscation" is
-            used for model protection, which can refer to :func:`mindspore.obfuscate_model`. If the input `graph` is
-            a func_graph loaded from a mindir file obfuscated with `obf_random_seed` , then `obf_random_seed` should be
-            provided. `obf_random_seed` should be in (0, 9223372036854775807]. default: ``None`` .
+        obf_random_seed (Union[int, None]): The random seed used for dynamic obfuscation, which is not supported now.
 
     Raises:
+        NotImplementedError: Dynamic structure obfuscation is not supported now.
         TypeError: If the `graph` is not a FuncGraph.
         TypeError: If the `params_init` is not a dict.
         TypeError: If the key of the `params_init` is not a str.
@@ -2855,20 +3926,12 @@ class GraphCell(Cell):
 
     def __init__(self, graph, params_init=None, obf_random_seed=None):
         super(GraphCell, self).__init__(auto_prefix=True)
+        if obf_random_seed is not None:
+            raise NotImplementedError("Dynamic structure obfuscation is not supported now.")
         if not isinstance(graph, FuncGraph):
             raise TypeError(f"For 'GraphCell', the argument 'graph' must be a FuncGraph loaded from MindIR, "
                             f"but got type {type(graph)}.")
         self.graph = graph
-        self.obf_random_seed = obf_random_seed
-        if obf_random_seed is not None:
-            if not isinstance(obf_random_seed, int):
-                raise TypeError("'obf_random_seed' must be int, but got {}.".format(type(obf_random_seed)))
-            int_64_max = 9223372036854775807
-            if obf_random_seed <= 0 or obf_random_seed > int_64_max:
-                raise ValueError(
-                    "'obf_random_seed' must be larger than 0, and less or equal than int64 ({}),"
-                    "but got {}.".format(int_64_max, obf_random_seed))
-            self._branch_control_input = _generate_branch_control_input(self.obf_random_seed)
         params_init = {} if params_init is None else params_init
         if not isinstance(params_init, dict):
             raise TypeError(f"For 'GraphCell', the argument 'params_init' must be a dict, but got {type(params_init)}.")
@@ -2888,19 +3951,30 @@ class GraphCell(Cell):
     def __call__(self, *args, **kwargs):
         self.phase = "graph_load_from_mindir"
         self._add_attr("graph_load_from_mindir", self.graph)
-        if not self.obf_random_seed:
-            return self.compile_and_run(*args, **kwargs)
-        append_input = Tensor((numpy.ones((1,)) * self._branch_control_input).astype(numpy.int32))
-        return self.compile_and_run(*args, append_input, **kwargs)
+        return self.compile_and_run(*args, **kwargs)
 
 
-def _check_param_list_tuple(value):
+def _is_parameter_list_or_tuple(value):
     """
     Check the type of input in list or tuple is Parameter.
     :param value: list or tuple.
     :return: The types of all inputs are parameter.
     """
-    for item in value:
-        if not isinstance(item, Parameter):
-            return False
-    return True
+    if isinstance(value, (list, tuple)) and value:
+        for item in value:
+            if not isinstance(item, Parameter):
+                return False
+        return True
+    return False
+
+
+def _addindent(s_, num_spaces):
+    s = s_.split("\n")
+    # don't do anything for single-line stuff
+    if len(s) == 1:
+        return s_
+    first = s.pop(0)
+    s = [(num_spaces * " ") + line for line in s]
+    s = "\n".join(s)
+    s = first + "\n" + s
+    return s