mindspore 2.3.0__cp310-cp310-win_amd64.whl → 2.4.0__cp310-cp310-win_amd64.whl

mindspore/common/api.py

@@ -38,12 +38,13 @@ from mindspore.common.tensor import Tensor as PythonTensor
 from mindspore.common.sparse_tensor import CSRTensor as PythonCSRTensor
 from mindspore.common.sparse_tensor import COOTensor as PythonCOOTensor
 from mindspore.common.sparse_tensor import RowTensor as PythonRowTensor
+from mindspore._c_expression.amp import get_curr_amp_strategy
 from mindspore._c_expression import GraphExecutor_, Tensor, CSRTensor, RowTensor, COOTensor, \
     PyNativeExecutor_, verify_inputs_signature, init_exec_dataset, _set_dataset_mode_config, init_pipeline, \
-    _ms_memory_recycle, _bind_device_ctx, jit_mode_pi_enable, jit_mode_pi_compile
+    _ms_memory_recycle, _bind_device_ctx
 from mindspore.parallel._ps_context import _is_role_sched
 from mindspore.parallel._utils import _check_full_batch, _get_parameter_broadcast, _is_pynative_parallel, \
-    _is_in_auto_parallel_mode
+    _is_in_auto_parallel_mode, _is_parallel_mode
 from mindspore import _checkparam as Validator
 from mindspore._checkparam import is_stub_tensor
 from mindspore.common._utils import is_shape_unknown
@@ -51,6 +52,8 @@ from mindspore.common.mutable import mutable
 from mindspore.common._register_for_adapter import ms_adapter_registry
 from mindspore.common.auto_dynamic_shape import get_auto_dynamic_shape_args, update_auto_dynamic_shape_phase, \
     get_auto_dynamic_shape_args_with_check_input_signature, update_auto_dynamic_shape_phase_with_check_input_signature
+from mindspore.common._pijit_context import PIJitCaptureContext
+from mindspore.common.parameter import Parameter
 
 # Store ms_function class compiled pipeline cache.
 ms_compile_cache = set()
@@ -513,6 +516,19 @@ def _generate_dyn_compile_args(compile_args, dyn_args):
     return tuple(new_compile_args)
 
 
+def _get_parameter_ids(args, kwargs):
+    """Get the ids of parameters."""
+    parameter_ids = ""
+    for arg in args:
+        if isinstance(arg, Parameter):
+            parameter_ids += str(id(arg))
+    for _, value in kwargs.items():
+        # The type of key is usually String type.
+        if isinstance(value, Parameter):
+            parameter_ids += str(id(value))
+    return parameter_ids
+
+
 class _MindsporeFunctionExecutor:
     """
     Represents a function compiled by graph compiler.
@@ -625,6 +641,10 @@ class _MindsporeFunctionExecutor:
 
         self._graph_executor.set_enable_tuple_broaden(self.enable_tuple_broaden)
         key = self._graph_executor.generate_arguments_key(self.fn, compile_args, kwargs, self.enable_tuple_broaden)
+
+        parameter_ids = _get_parameter_ids(args, kwargs)
+        if parameter_ids != "":
+            key = str(key) + '.' + parameter_ids
         phase = generate_name + '.' + str(key)
 
         update_auto_dynamic_shape_phase_with_check_input_signature(compile_args, key_id, phase, self.input_signature)
@@ -783,31 +803,28 @@ def _get_jit_hash(hash_input):
     return _get_obj_id(hash_input)
 
 
-def _update_graph_executor_config(jit_config):
-    """Update GraphExecutor jit_config"""
-    if isinstance(jit_config, JitConfig):
-        jit_config = jit_config.jit_config_dict
-    if not isinstance(jit_config, dict):
-        return
-    valid_config = dict()
-    for k, v in jit_config.items():
-        valid_config[str(k)] = str(v)
-    GraphExecutor_.get_instance().set_jit_config(JitConfig(**valid_config).jit_config_dict)
-
-
 def jit(fn=None, mode="PSJit", input_signature=None, hash_args=None, jit_config=None, compile_once=False):
     """
     Create a callable MindSpore graph from a Python function.
 
     This allows the MindSpore runtime to apply optimizations based on graph.
 
+    Note:
+        - If `input_signature` is specified, each input of `fn` must be a Tensor. And the input arguments for `fn`
+          will not accept `**kwargs`.
+        - It is not supported to run a function with decoration @jit(mode=“PIJit”)
+          in static graph mode, in which case the decoration @jit(mode=“PIJit”) is considered invalid.
+        - Calls to functions with decorated @jit(mode=“PIJit”) inside functions
+          decorated with @jit(mode=“PIJit”) are not supported,
+          and the decoration @jit(mode=“PIJit”) is considered invalid.
+
     Args:
         fn (Function): The Python function that will be run as a graph. Default: ``None`` .
         mode (str): The type of jit used, the value of mode should be ``PIJit`` or ``PSJit``. Default: ``PSJit`` .
 
-            - `PSJit <https://www.mindspore.cn/docs/en/master/note/static_graph_syntax_support.html>`_ :
+            - `PSJit <https://www.mindspore.cn/docs/en/master/model_train/program_form/static_graph.html>`_ :
               Parse python ast to build graph.
-            - `PIJit <https://www.mindspore.cn/docs/en/master/design/dynamic_graph_and_static_graph.html>`_ :
+            - `PIJit <https://www.mindspore.cn/docs/en/master/model_train/program_form/pynative.html#pijit>`_ :
               Parse python bytecode to build graph at runtime.
 
         input_signature (Union[Tuple, List, Dict, Tensor]): The Tensor which describes the input arguments. The
@@ -831,10 +848,6 @@ def jit(fn=None, mode="PSJit", input_signature=None, hash_args=None, jit_config=
             it was created again.
             Default: ``False`` .
 
-    Note:
-        If `input_signature` is specified, each input of `fn` must be a Tensor. And the input arguments for `fn`
-        will not accept `**kwargs`.
-
     Returns:
         Function, if `fn` is not None, returns a callable function that will execute the compiled function; If `fn` is
         None, returns a decorator and when this decorator invokes with a single `fn` argument, the callable function is
@@ -938,45 +951,20 @@ def jit(fn=None, mode="PSJit", input_signature=None, hash_args=None, jit_config=
             # only the function or cell instance wrapped by shard will fall into this branch
             if _is_pynative_parallel() and func.__name__ == _PYNATIVE_PARALLEL_FUNC_NAME:
                 process_obj = hash_args
+            # Handle auto mixed precision strategy.
+            if not hasattr(func, "amp_strategy"):
+                if isinstance(func, types.MethodType):
+                    setattr(func.__func__, "amp_strategy", get_curr_amp_strategy())
+                else:
+                    setattr(func, "amp_strategy", get_curr_amp_strategy())
             out = _MindsporeFunctionExecutor(func, hash_obj, dyn_args, process_obj, jit_config)(*args, **kwargs)
             return out
 
         return staging_specialize
 
-    def pi_wrap_mindspore(decorated):
-        func = decorated
-        if isinstance(func, ms.nn.Cell):
-            func = func.construct
-        if isinstance(func, type) and issubclass(func, ms.nn.Cell):
-            func = func.construct
-        if isinstance(func, types.MethodType):
-            func = func.__func__
-        if not isinstance(func, types.FunctionType):
-            logger.warning("only support function and mindspore.nn.Cell instance")
-            return decorated
-
-        # generator, coroutine, awaitable and a function that return them is unsupported
-        UNSUPPORTED_CODE_TYPE = (inspect.CO_GENERATOR | inspect.CO_COROUTINE |
-                                 inspect.CO_ASYNC_GENERATOR | inspect.CO_ITERABLE_COROUTINE)
-        if func.__code__.co_flags & UNSUPPORTED_CODE_TYPE:
-            return decorated
-
-        _update_graph_executor_config(jit_config)
-        config = dict()
-        if isinstance(jit_config, JitConfig):
-            config.update(jit_config.jit_config_dict)
-        elif jit_config is not None:
-            config.update(jit_config)
-        jit_mode_pi_enable()
-
-        if jit_mode_pi_compile(func, config, input_signature) is False:
-            logger.warning('add fn {} to compile failed '.format(func))
-
-        return decorated
-
     wrap_func = wrap_mindspore
     if mode == "PIJit":
-        wrap_func = pi_wrap_mindspore
+        wrap_func = PIJitCaptureContext(jit_config, input_signature)
 
     if fn is not None:
         return wrap_func(fn)
@@ -1272,7 +1260,7 @@ def jit_class(cls):
     if not inspect.isclass(cls):
         raise TypeError(f'Decorator jit_class can only be used for class type, but got {cls}.')
     # Check if cls is nn.Cell.
-    if issubclass(cls, nn.Cell):
+    if issubclass(cls, nn.cell.Cell):
         raise TypeError(f"Decorator jit_class is used for user-defined classes and cannot be used for nn.Cell: {cls}.")
     setattr(cls, '__ms_class__', True)
     return cls
@@ -1463,23 +1451,22 @@ class _PyNativeExecutor:
         """
         self._executor.end_graph(obj, output, *args, *(kwargs.values()))
 
-    def check_run(self, grad, obj, weights, grad_hash_id, *args, **kwargs):
+    def check_run(self, grad, obj, weights, grad_hash_id, *args):
         """
         Whether the forward graph need to construct.
 
         Args:
             grad (GradOperation): The gradoperation object.
             obj (Function/Cell): The function or cell instance.
-            grad_hash_id (tuple): The id of objects which contribute to cache of compiled graph in pynative mode.
+            grad_hash_id (tuple): The id of objects, which contributes to cache of compiled graph in pynative mode.
             args (tuple): Function or cell input arguments.
-            kwargs (dict): keyword arguments.
 
         Return:
-            bool, specifies whether the forward graph need to construct.
+            bool, specifies whether the forward graph needs to construct.
         """
-        return self._executor.check_run(grad, obj, weights, grad_hash_id, *args, *(kwargs.values()))
+        return self._executor.check_run(grad, obj, weights, grad_hash_id, *args)
 
-    def grad(self, obj, grad, weights, grad_position, *args, **kwargs):
+    def grad(self, obj, grad, weights, grad_position, *args):
         """
         Get grad graph.
 
@@ -1490,12 +1477,11 @@ class _PyNativeExecutor:
             grad_position (Union(int, tuple[int])): If int, get the gradient with respect to single input.
               If tuple, get the gradients with respect to selected inputs. 'grad_position' begins with 0. Default: 0.
             args (tuple): Function or cell input arguments.
-            kwargs (dict): keyword arguments.
 
         Return:
             None.
         """
-        return self._executor.grad(grad, obj, weights, grad_position, *args, *(kwargs.values()))
+        return self._executor.grad(grad, obj, weights, grad_position, *args)
 
     def clear_res(self):
         """
@@ -1528,9 +1514,23 @@ class _PyNativeExecutor:
         """
         return self._executor.grad_jit(output, *args)
 
+    def call_custom_bprop(self, obj, output, *args, **kwargs):
+        """
+        Call custom bprop to build variable for cell bprop.
+        Args:
+            obj (Cell): The function or cell instance.
+            output (Tensor/tuple/list): Function or cell output object.
+            args (tuple): Function or cell input arguments.
+            kwargs (dict): keyword arguments.
+
+        Return:
+            None.
+        """
+        return self._executor.call_custom_bprop(obj, output, *args, *(kwargs.values()))
+
     def grad_flag(self):
         """
-        The flag of building grad graph.
+        The flag of whether the net building grad graph.
 
         Return:
             bool, whether building grad graph.
@@ -1563,7 +1563,7 @@ class _PyNativeExecutor:
 
     def enable_grad(self):
         """
-        The global flag whether needing to calculate gradient.
+        The global flag that whether need to calculate gradient use in no_grad.
 
         Return:
             bool, whether needing to calculate gradient.
@@ -1582,6 +1582,18 @@ class _PyNativeExecutor:
         """
         self._executor.set_enable_grad(flag)
 
+    def requires_grad(self):
+        """
+        When both enable_grad is true and grad_flag is true, that the flag requires_grad will be true.
+
+        Args:
+            flag (bool): Specifying whether calculating gradient.
+
+        Return:
+            None.
+        """
+        return self._executor.requires_grad()
+
     def set_jit_compile_status(self, status, phase):
         """
         Set jit is compiling
@@ -1605,6 +1617,18 @@ class _PyNativeExecutor:
         """
         self._executor.set_is_run_recompute(status)
 
+    def set_cell_use_dynamic_shape_process(self, flag):
+        """
+        Set the dynamic shape flag of eval process.
+
+        Args:
+            flag (bool): Specifying whether using a dynamic process.
+
+        Return:
+            None.
+        """
+        self._executor.set_cell_use_dynamic_shape_process(flag)
+
     def set_dynamic_input(self, obj, *args):
         """
         Set dynamic shape tensor of input arguments.
@@ -1630,27 +1654,19 @@ class _PyNativeExecutor:
         """
         return self._executor.get_dynamic_input(*actual_args)
 
-    def is_first_cell(self):
-        """
-        The flag of first cell instance.
-
-        Return:
-            bool, specifies whether is the first cell.
+    def set_mixed_precision_type(self, mixed_precision_type, is_push=True):
         """
-
-        return self._executor.is_first_cell()
-
-    def set_hook_changed(self, cell):
-        """
-        The flag of registering or removing a hook function on Cell instance.
+        The value of mixed precision type.
 
         Args:
-            cell (Cell): The cell instance.
+            type(MixedPrecisionType): Mix precision type.
+            is_push(bool): If called by __enter__, is push will be True
 
         Return:
             None.
         """
-        self._executor.set_hook_changed(cell)
+
+        return self._executor.set_mixed_precision_type(mixed_precision_type, is_push)
 
     def constant_folding(self, *args):
         """
@@ -1687,6 +1703,7 @@ class _CellGraphExecutor:
         self._graph_executor = GraphExecutor_.get_instance()
         self._graph_executor.set_py_exe_path(sys.executable)
         self._graph_executor.set_kernel_build_server_dir(os.path.split(kernel_build_server.__file__)[0] + os.sep)
+        self._pid = os.getpid()
 
     def init_dataset(self, queue_name, dataset_size, batch_size, dataset_types, dataset_shapes,
                      input_indexs, phase='dataset', need_run=True):
@@ -1789,6 +1806,10 @@ class _CellGraphExecutor:
         self._graph_executor.set_enable_tuple_broaden(self.enable_tuple_broaden)
         key = self._graph_executor.generate_arguments_key(obj, args, kwargs, self.enable_tuple_broaden)
         obj.arguments_key = str(key)
+        # When exist parameter in the top graph inputs, need check if the parameter object has changed.
+        parameter_ids = _get_parameter_ids(args, kwargs)
+        if parameter_ids != "":
+            obj.arguments_key = obj.arguments_key + '.' + parameter_ids
         raw_phase = phase
         phase = phase + '.' + str(obj.create_time) + '.' + str(id(obj)) + '.' + obj.arguments_key
         obj.phase_cache[raw_phase] = phase
@@ -1825,7 +1846,7 @@ class _CellGraphExecutor:
         if graph is None:
             raise RuntimeError("Compile graph failed for phase {}.".format(phase))
 
-        auto_parallel_mode = _is_in_auto_parallel_mode()
+        auto_parallel_mode = _is_in_auto_parallel_mode() or _is_parallel_mode()
         if not auto_parallel_mode:
             replace = obj.init_parameters_data(auto_parallel_mode=auto_parallel_mode)
             self._update_param_node_default_input(phase, replace)
@@ -1913,15 +1934,9 @@ class _CellGraphExecutor:
 
     def del_net_res(self, obj, net_id):
         """Clear the memory resource of a network."""
-        self._graph_executor.del_net_res(obj, net_id)
-
-    def inc_graph_cell_count(self):
-        """Increase the count of GraphCell instance."""
-        self._graph_executor.inc_graph_cell_count()
-
-    def dec_graph_cell_count(self):
-        """Decrease the count of GraphCell instance."""
-        self._graph_executor.dec_graph_cell_count()
+        # no need to del net res by gc in independent dataset process which is a subprocess forked by main process
+        if self._pid == os.getpid():
+            self._graph_executor.del_net_res(obj, net_id)
 
     def _get_branch_control_input(self):
         if ('obf_ratio' not in self.obfuscate_config.keys()) or (

mindspore/common/dump.py

@@ -27,18 +27,17 @@ def set_dump(target, enabled=True):
     `target` should be an instance of :class:`mindspore.nn.Cell` or :class:`mindspore.ops.Primitive` .
     Please note that this API takes effect only when Synchronous Dump is enabled and the `dump_mode`
     field in dump config file is ``"2"`` . See the `dump document
-    <https://www.mindspore.cn/tutorials/experts/en/master/debug/dump.html>`_ for details.
+    <https://www.mindspore.cn/docs/en/master/model_train/debug/dump.html>`_ for details.
     The default enabled status for
     a :class:`mindspore.nn.Cell` or :class:`mindspore.ops.Primitive` is False.
 
     Note:
-        1. This API is only effective for GRAPH_MODE whose graph compilation level is O0/O1 with Ascend backend.
-        2. This API only supports being called before training starts.
+        1. This API only supports being called before training starts.
            If you call this API during training, it may not be effective.
-        3. After using `set_dump(Cell, True)` , operators in forward and backward
+        2. After using `set_dump(Cell, True)` , operators in forward and backward
            computation  (computation generated by the grad operations) of the
            cell will be dumped.
-        4. For :class:`mindspore.nn.SoftmaxCrossEntropyWithLogits` layer, the forward
+        3. For :class:`mindspore.nn.SoftmaxCrossEntropyWithLogits` layer, the forward
            computation and backward computation use the same set of
            operators. So you can only see dump data from backward computation.
            Please note that :class:`mindspore.nn.SoftmaxCrossEntropyWithLogits` layer will also use
@@ -58,7 +57,7 @@ def set_dump(target, enabled=True):
         .. note::
             Please set environment variable `MINDSPORE_DUMP_CONFIG` to the dump config file and set `dump_mode` field
             in dump config file to 2 before running this example.
-            See `dump document <https://www.mindspore.cn/tutorials/experts/en/master/debug/dump.html>`_ for details.
+            See `dump document <https://www.mindspore.cn/docs/en/master/model_train/debug/dump.html>`_ for details.
 
         >>> import numpy as np
         >>> import mindspore as ms

mindspore/common/generator.py

@@ -56,12 +56,6 @@ class Generator:
     A generator that manages the state of random numbers and provides seed and offset for random functions.
     When the seed and offset are fixed, the random function generates the same random sequence.
 
-    Inputs:
-        - **step** (int) - Set the step size for offset update.
-
-    Outputs:
-        Tuple consisting of the seed and offset of generator.
-
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
 
@@ -199,7 +193,7 @@ def manual_seed(seed):  # pylint: disable=redefined-outer-name
         >>> print(initial_seed())
         13
     """
-    default_generator.manual_seed(seed)
+    return default_generator.manual_seed(seed)
 
 
 def initial_seed():

mindspore/common/hook_handle.py

@@ -77,27 +77,19 @@ class HookHandle:
         It is only supported in pynative mode and works when registering or removing hook function for Cell object.
 
     Args:
-        hook_cell (Cell): The Cell object with hook function registered on. Default value: None.
-        hook_key (int): The key of cell hook function in dict. It is generated during cell hook function registration.
-                        Default value: -1.
-        hook_type (str): The type of cell hook function: '_forward_pre_hook', '_forward_hook' or '_cell_backward_hook'.
-                         Default value: "".
+        hook_dict (Dict): The hook object with hook function registered on. Default value: None.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
     """
-    def __init__(self, hook_cell=None, hook_key=-1, hook_type=""):
-        if hook_cell is not None:
-            self._hook_cell = weakref.ref(hook_cell)
-        else:
-            self._hook_cell = hook_cell
-        self._hook_key = hook_key
-        self._hook_type = hook_type
-
-    def __del__(self):
-        self._hook_cell = None
-        self._hook_key = None
-        self._hook_type = None
+    unique_id = 0
+
+    def __init__(self, hook_dict=None):
+        self.hook_dict_ref = None
+        if hook_dict is not None:
+            self.hook_dict_ref = weakref.ref(hook_dict)
+            self.handle_id = HookHandle.unique_id
+            HookHandle.unique_id += 1
 
     def remove(self):
         """
@@ -121,7 +113,7 @@ class HookHandle:
             >>> from mindspore import Tensor
             >>> from mindspore.ops import GradOperation
             >>> ms.set_context(mode=ms.PYNATIVE_MODE)
-            >>> def forward_pre_hook_fn(cell_id, inputs):
+            >>> def forward_pre_hook_fn(cell, inputs):
             ...     print("forward inputs: ", inputs)
             ...
             >>> class Net(nn.Cell):
@@ -145,11 +137,7 @@ class HookHandle:
             (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], dtype=Float32,
             value= [ 2.00000000e+00]))
         """
-        if self._hook_cell is not None:
-            hook_cell = self._hook_cell()
-            if self._hook_type == "_forward_pre_hook" and self._hook_key in hook_cell._forward_pre_hook:
-                del hook_cell._forward_pre_hook[self._hook_key]
-            elif self._hook_type == "_forward_hook" and self._hook_key in hook_cell._forward_hook:
-                del hook_cell._forward_hook[self._hook_key]
-            elif self._hook_type == "_cell_backward_hook":
-                hook_cell._cell_backward_hook.remove_backward_hook(self._hook_key)
+        if self.hook_dict_ref is not None:
+            hook_dict = self.hook_dict_ref()
+            if hook_dict is not None and self.handle_id in hook_dict:
+                del hook_dict[self.handle_id]

mindspore/common/mindir_util.py

@@ -90,9 +90,9 @@ def save_mindir(model, file_name):
     if not file_name.endswith('.mindir'):
         file_name += ".mindir"
 
-    current_path = os.path.abspath(file_name)
+    current_path = os.path.realpath(file_name)
     dirname = os.path.dirname(current_path)
-    os.makedirs(dirname, exist_ok=True)
+    os.makedirs(dirname, mode=0o700, exist_ok=True)
     if os.path.exists(file_name):
         os.chmod(file_name, stat.S_IWUSR)
 

mindspore/common/parameter.py

@@ -41,6 +41,8 @@ from mindspore.parallel._ps_context import _is_role_worker, _is_role_pserver, _i
                                            _is_ps_mode
 from mindspore.parallel._ps_context import _reinsert_hash_table_size, _insert_accumu_init_info, _cache_enable
 from mindspore.common._decorator import deprecated
+from mindspore.communication._comm_helper import _is_initialized
+from mindspore.communication import get_group_size
 import mindspore.common._monad as monad
 
 __all__ = ['Parameter', 'ParameterTuple']
@@ -52,11 +54,22 @@ PARAMETER_NAME_PREFIX_MAX_LEN = 1024
 _GLOBAL_PARAMETER_KEY = -1
 
 
-def _is_in_parallel_mode():
+def _is_in_auto_parallel_mode():
     """Get parallel mode."""
     return auto_parallel_context().get_parallel_mode() in ["semi_auto_parallel", "auto_parallel"]
 
 
+def _is_parallel_mode():
+    """ Whether is parallel mode """
+    if not _is_initialized() or context.get_context('mode') == context.PYNATIVE_MODE:
+        return False
+    if os.getenv("RUN_MODE") != "predict":
+        return False
+    if get_group_size() > 1 and _get_parallel_mode() == "stand_alone":
+        return True
+    return False
+
+
 def init_to_value(init):
     """
     Get value of initializer.
@@ -91,6 +104,15 @@ def _get_unique_parameter_key():
     return _GLOBAL_PARAMETER_KEY
 
 
+def _gen_offload_file_path(offload_dir):
+    offload_dir = os.path.relpath(offload_dir)
+    if not os.path.exists(offload_dir):
+        os.makedirs(offload_dir, mode=0o700, exist_ok=True)
+    offload_file_path = offload_dir + "/" + str(_get_global_rank()) + "_" + str(
+        _get_unique_parameter_key()) + "_" + str(time.time()) + ".data"
+    return offload_file_path
+
+
 def _offload_if_config(data):
     """
     Offload parameter(data size > 512) to file when enable memory offload and offload parameter to disk.
@@ -111,11 +133,7 @@ def _offload_if_config(data):
     offload_file_path = data.offload_file_path()
     if offload_file_path is None or offload_file_path == "":
         offload_dir = offload_context.get("offload_path", "./offload")
-        offload_dir = os.path.relpath(offload_dir)
-        if not os.path.exists(offload_dir):
-            os.makedirs(offload_dir)
-        offload_file_path = offload_dir + "/" + str(_get_global_rank()) + "_" + str(
-            _get_unique_parameter_key()) + "_" + str(time.time()) + ".data"
+        offload_file_path = _gen_offload_file_path(offload_dir)
     data.offload(offload_file_path)
 
 
@@ -191,6 +209,12 @@ class Parameter(Tensor_):
         storage_format (str): Only Ascend device target is supported. It is used to specify the format of the weight
             loaded to the device. By default, the format is not changed. The optional values are ``"FRACTAL_NZ"`` ,
             ``"NC1HWC0"`` , ``"FRACTAL_Z"`` , etc. Default: ``""`` .
+        device(str): Only Ascend device target is supported. It is used to specify the device which the parameter is
+            stored. By default, the parameter will be stored on NPU while computing. When the device is specified as
+            ``"CPU"``, the parameter will be loaded into the device when it needs to be used, and unloaded to the CPU
+            after use. It takes effext only when `memory_offload` is ``"ON"``, `jit_level` is not ``"O2"`` and
+            `memory_optimize_level` is ``O0`` in `mindspore.set_context()`. Less device memory is needed when device is
+            specified as ``"CPU"``.
 
     Examples:
         >>> import numpy as np
@@ -244,7 +268,7 @@ class Parameter(Tensor_):
             Parameter, (data, self.name, self.requires_grad, self.layerwise_parallel))
 
     def __init__(self, default_input, name=None, requires_grad=True, layerwise_parallel=False, parallel_optimizer=True,
-                 storage_format=""):
+                 storage_format="", device=None):
         self.param_info = ParamInfo()
         self.init_in_server = False
         self.name = name
@@ -263,7 +287,7 @@ class Parameter(Tensor_):
         self.requires_aggr = True
         self._cast_type = None
         self._unique = False
-        self.is_in_parallel = _is_in_parallel_mode()
+        self.is_in_parallel = _is_in_auto_parallel_mode()
         self.is_in_shard = False
         self._pipeline_stage_list = []
         self.slice_num = 1
@@ -296,6 +320,10 @@ class Parameter(Tensor_):
                             f" 'numpy.ndarray', 'list']. But got type {type(default_input)}.")
         self.param_info.parameter_shape = self.shape
         self.param_info.storage_format = storage_format
+        if device is not None:
+            if device != "CPU":
+                raise ValueError(f"Only 'CPU' is supported for device, but got ${device}.")
+            self._set_user_data("parameter_device", device)
 
         import mindspore.ops.operations.other_ops as other_ops
         self.load = other_ops.Load()
@@ -342,7 +370,8 @@ class Parameter(Tensor_):
                     return (Tensor, data.asnumpy(), mstype.qint4x2)
                 return (Tensor, data.asnumpy())
 
-            not_init_data = _is_role_sched() or (_is_role_pserver() and _cache_enable()) or _is_in_parallel_mode()
+            not_init_data = _is_role_sched() or (_is_role_pserver() and _cache_enable()
+                                                 ) or _is_in_auto_parallel_mode() or _is_parallel_mode()
             if not_init_data:
                 # do not init data while in auto parallel.
                 return (Tensor, None, data.dtype, get_slice_shape(data.dtype, data.shape), data.init)
@@ -368,7 +397,7 @@ class Parameter(Tensor_):
 
         Tutorial Examples:
             - `Parameter Server Mode
-              <https://www.mindspore.cn/tutorials/experts/en/master/parallel/parameter_server_training.html>`_
+              <https://www.mindspore.cn/docs/en/master/model_train/parallel/parameter_server_training.html>`_
         """
         if not _is_ps_mode() or not (_is_role_worker() or _is_role_pserver() or _is_role_sched()):
             raise RuntimeError("Must complete following two steps before calling set_param_ps: \n"
@@ -616,6 +645,9 @@ class Parameter(Tensor_):
             shape = self.shape if self.slice_num == 1 else self.param_info.origin_shape
             dtype = self.dtype
             x.set_data(initializer(init, shape=shape, dtype=dtype))
+        device = self._get_user_data("parameter_device")
+        if device is not None:
+            x._set_user_data("parameter_device", device)
         return x
 
     @property
@@ -942,7 +974,7 @@ class Parameter(Tensor_):
             >>> x = Parameter(Tensor(np.array([[1, 2], [3, 4]], dtype=np.float32)), name="param")
             >>> x.init_data()
         """
-        if self.is_default_input_init and self.is_in_parallel != _is_in_parallel_mode():
+        if self.is_default_input_init and self.is_in_parallel != _is_in_auto_parallel_mode():
             raise RuntimeError("Must set or change parallel mode before any initializer Tensor created.")
         if self.init_mode is None:
             return self
@@ -1026,8 +1058,9 @@ class ParameterTuple(tuple):
             Tuple, the new Parameter tuple.
 
         Tutorial Examples:
-            - `Cell and Parameter - Parameter Tuple
-              <https://mindspore.cn/tutorials/en/master/advanced/modules/layer.html#parameter-tuple>`_
+            - `Tensor and Parameter - Parameter Tuple
+              <https://mindspore.cn/docs/en/master/model_train/model_building/tensor_and_parameter.html
+              #parameter-tuple>`_
         """
         Validator.check_str_by_regular(prefix)
         new = []