mindspore 1.10.0__cp39-cp39-win_amd64.whl → 2.0.0rc1__cp39-cp39-win_amd64.whl

mindspore/context.py

@@ -1,4 +1,4 @@
- # Copyright 2020-2022 Huawei Technologies Co., Ltd
+# Copyright 2020-2022 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.
@@ -27,22 +27,23 @@ from types import FunctionType
 
 from mindspore import log as logger
 from mindspore._c_expression import MSContext, ms_ctx_param
-from mindspore._checkparam import args_type_check, Validator, args_unreset_check
+from mindspore import _checkparam as Validator
+from mindspore._checkparam import args_type_check
 from mindspore.parallel._auto_parallel_context import _set_auto_parallel_context, _get_auto_parallel_context, \
     _reset_auto_parallel_context
 from mindspore.parallel._ps_context import _set_ps_context, _get_ps_context, _reset_ps_context, \
     _need_reset_device_target_for_ps
-from mindspore.default_config import __device_target__, __package_name__
+from mindspore.parallel._offload_context import _set_offload_context, _get_offload_context
 
 __all__ = ['GRAPH_MODE', 'PYNATIVE_MODE', 'set_context', 'get_context', 'set_auto_parallel_context',
            'get_auto_parallel_context', 'reset_auto_parallel_context', 'ParallelMode', 'set_ps_context',
-           'get_ps_context', 'reset_ps_context']
+           'get_ps_context', 'reset_ps_context', 'set_offload_context', 'get_offload_context']
 
 GRAPH_MODE = 0
 PYNATIVE_MODE = 1
 _DEVICE_APP_MEMORY_SIZE = 31  # The max memory size of graph plus variable.
 _re_pattern = r'[1-9][0-9]*(\.)?[0-9]*GB|0\.[0-9]*GB'
-_k_context = None
+K_CONTEXT = None
 
 
 def _make_directory(path):
@@ -184,11 +185,12 @@ class _Context:
             if self.enable_debug_runtime:
                 self.set_backend_policy("vm")
             parallel_mode = _get_auto_parallel_context("parallel_mode")
-            if parallel_mode not in (ParallelMode.DATA_PARALLEL, ParallelMode.STAND_ALONE):
-                raise ValueError(f"Got {parallel_mode}, when the user enabled SEMI_AUTO_PARALELL or AUTO_PARALLEL, "
-                                 f"pynative mode dose not support, you should set "
-                                 f"context.set_auto_parallel_context(parallel_mode='data_parallel') "
-                                 f"or context.set_auto_parallel_context(parallel_mode='stand_alone').")
+            if parallel_mode not in (ParallelMode.DATA_PARALLEL, ParallelMode.STAND_ALONE, ParallelMode.AUTO_PARALLEL):
+                raise ValueError(f"Got {parallel_mode}, when the user enabled SEMI_AUTO_PARALELL, "
+                                 f"pynative mode dose not support, you should set either "
+                                 f"context.set_auto_parallel_context(parallel_mode='data_parallel'), "
+                                 f"context.set_auto_parallel_context(parallel_mode='stand_alone') "
+                                 f"or context.set_auto_parallel_context(parallel_mode='auto_parallel').")
             self._context_switches.push(True, None)
         elif mode == GRAPH_MODE:
             if self.enable_debug_runtime:
@@ -199,6 +201,83 @@ class _Context:
                              f"or context.PYNATIVE_MODE (1), but got {mode}.")
         self.set_param(ms_ctx_param.mode, mode)
 
+    def set_memory_optimize_level(self, memory_optimize_level):
+        """
+        The memory optimize level, support "O0", "O1".
+
+        Args:
+            target (str): "O0", "O1"
+        """
+        memory_optimize_levels = ["O0", "O1"]
+        if memory_optimize_level not in memory_optimize_levels:
+            raise ValueError(f"For 'context.set_context', the argument 'memory_optimize_level' must be one of "
+                             f"{memory_optimize_levels}, but got {memory_optimize_level}.")
+        if memory_optimize_level == "O0":
+            self.set_param(ms_ctx_param.memory_optimize_level, 0)
+        else:
+            self.set_param(ms_ctx_param.memory_optimize_level, 1)
+
+    def set_memory_offload(self, memory_offload):
+        """
+        Enable memory offload or not, support "ON", "OFF".
+
+        Args:
+            memory_offload (str): "ON", "OFF"
+        """
+        memory_offload_options = ["ON", "OFF"]
+        if memory_offload not in memory_offload_options:
+            raise ValueError(f"For 'context.set_context', the argument 'memory_offload' must be one of "
+                             f"{memory_offload_options}, but got {memory_offload}.")
+        if memory_offload == "ON":
+            self.set_param(ms_ctx_param.memory_offload, True)
+        else:
+            self.set_param(ms_ctx_param.memory_offload, False)
+
+    def set_deterministic(self, deterministic):
+        """
+        Enable model run in deterministic, and support the values "ON" and "OFF".
+
+        Args:
+            deterministic (str): "ON", "OFF"
+        """
+        deterministic_options = ["ON", "OFF"]
+        if deterministic not in deterministic_options:
+            raise ValueError(f"For 'context.set_context', the argument 'deterministic' must be one of "
+                             f"{deterministic_options}, but got {deterministic}.")
+        self.set_param(ms_ctx_param.deterministic, deterministic)
+
+    def set_ascend_config(self, ascend_config):
+        """
+        Enable ascend config.
+
+        Args:
+            ascend_config (dict): 'precision_mode'
+                - precision_mode (str): "force_fp16", "allow_fp32_to_fp16", "allow_mix_precision",
+                            "must_keep_origin_dtype", "force_fp32", "force_lowerprecision", "allow_fp32_to_bf16",
+                            "allow_fp32_to_lowprecision", "allow_mix_precision_fp16" and "allow_mix_precision_bf16".
+        """
+
+        ascend_cfgs = {'precision_mode': ["force_fp16", "allow_fp32_to_fp16", "allow_mix_precision",
+                                          "must_keep_origin_dtype", "force_fp32", "force_lowerprecision",
+                                          "allow_fp32_to_bf16", "allow_fp32_to_lowprecision",
+                                          "allow_mix_precision_fp16", "allow_mix_precision_bf16"],
+                       'jit_compile': [True, False]}
+        for ascend_key in ascend_config:
+            if ascend_key not in ascend_cfgs:
+                raise ValueError(f"For 'context.set_context', the key of argument 'ascend_config' must be one of "
+                                 f"{ascend_cfgs}, but got {ascend_key}.")
+            supported_modes = ascend_cfgs.get(ascend_key)
+            if ascend_config[ascend_key] not in supported_modes:
+                raise ValueError(f"For 'ascend_config', the value of argument {ascend_key} must be one of "
+                                 f"{supported_modes}, but got {ascend_config[ascend_key]}.")
+            if ascend_key == 'precision_mode':
+                self.set_param(ms_ctx_param.precision_mode, ascend_config[ascend_key])
+            if ascend_key == 'jit_compile':
+                if ascend_config[ascend_key] is True:
+                    self.set_param(ms_ctx_param.jit_compile, "1")
+                else:
+                    self.set_param(ms_ctx_param.jit_compile, "0")
+
     def set_backend_policy(self, policy):
         success = self._context_handle.set_backend_policy(policy)
         if not success:
@@ -236,7 +315,7 @@ class _Context:
     def set_auto_tune_mode(self, tune_mode):
         candidate = ["NO_TUNE", "RL", "GA", "RL,GA", "GA,RL"]
         if tune_mode in candidate:
-            self.set_param(ms_ctx_param.tune_mode, tune_mode)
+            self.set_param(ms_ctx_param.auto_tune_mode, tune_mode)
         else:
             raise ValueError(f"For 'context.set_context', the argument 'auto_tune_mode' must be in "
                              f"['NO_TUNE', 'RL', 'GA', 'RL,GA', 'GA,RL'], but got {tune_mode}.")
@@ -336,14 +415,14 @@ class _Context:
                 json.load(f)
         except (TypeError, ValueError) as exo:
             raise ValueError(str(exo) + "\nFor 'context.set_context', open or load the 'env_config_path' file {} "
-                             "failed, please check whether 'env_config_path' is json file and correct, or may not "
-                             "have permission to read it.".format(env_config_path))
+                                        "failed, please check whether 'env_config_path' is json file and correct, "
+                                        "or may not have permission to read it.".format(env_config_path))
         self.set_param(ms_ctx_param.env_config_path, env_config_path)
 
     def set_runtime_num_threads(self, runtime_num_threads):
         """Check and set runtime_num_threads."""
-        if runtime_num_threads <= 0:
-            raise ValueError("The num of thread must bigger than 0.")
+        if runtime_num_threads < 0:
+            raise ValueError("The num of thread must bigger than or equal to 0.")
         self.set_param(ms_ctx_param.runtime_num_threads, runtime_num_threads)
 
     def set_op_timeout(self, op_timeout):
@@ -352,6 +431,12 @@ class _Context:
             raise ValueError("The num of op exe timeout must bigger than or equal to 0.")
         self.set_param(ms_ctx_param.op_timeout, op_timeout)
 
+    def set_inter_op_parallel_num(self, inter_op_parallel_num):
+        """Check and set inter_op_parallel_num."""
+        if inter_op_parallel_num < 0:
+            raise ValueError("The num of parallel thread must bigger than or equal to 0.")
+        self.set_param(ms_ctx_param.inter_op_parallel_num, inter_op_parallel_num)
+
     setters = {
         'mode': set_mode,
         'save_graphs_path': set_save_graphs_path,
@@ -365,8 +450,13 @@ class _Context:
         'mempool_block_size': set_mempool_block_size,
         'print_file_path': set_print_file_path,
         'env_config_path': set_env_config_path,
+        'inter_op_parallel_num': set_inter_op_parallel_num,
         'runtime_num_threads': set_runtime_num_threads,
-        'op_timeout': set_op_timeout
+        'memory_optimize_level': set_memory_optimize_level,
+        'op_timeout': set_op_timeout,
+        'memory_offload': set_memory_offload,
+        'deterministic': set_deterministic,
+        'ascend_config': set_ascend_config
     }
 
     @property
@@ -407,7 +497,6 @@ class _Context:
         self._support_binary = support
 
 
-
 def _context():
     """
     Get the global _context, if context is not created, create a new one.
@@ -415,28 +504,28 @@ def _context():
     Returns:
         _Context, the global context in PyNative mode.
     """
-    global _k_context
-    if _k_context is None:
+    global K_CONTEXT
+    if K_CONTEXT is None:
         default_backend = 'debug'
         try:
             from mindspore import default_config
             default_backend = default_config.__backend__
         except ImportError:
             logger.error("import default config fail")
-        _k_context = _Context()
-        _k_context.enable_debug_runtime = False
+        K_CONTEXT = _Context()
+        K_CONTEXT.enable_debug_runtime = False
         if default_backend == 'debug':
-            _k_context.enable_debug_runtime = True
+            K_CONTEXT.enable_debug_runtime = True
             default_backend = 'vm'
-        _k_context.set_backend_policy(default_backend)
-    return _k_context
+            K_CONTEXT.set_backend_policy(default_backend)
+    return K_CONTEXT
 
 
 @args_type_check(device_num=int, global_rank=int, gradients_mean=bool, gradient_fp32_sync=bool, parallel_mode=str,
                  auto_parallel_search_mode=str, search_mode=str, parameter_broadcast=bool, strategy_ckpt_load_file=str,
                  strategy_ckpt_save_file=str, full_batch=bool, enable_parallel_optimizer=bool, enable_alltoall=bool,
                  all_reduce_fusion_config=list, pipeline_stages=int, grad_accumulation_step=int,
-                 parallel_optimizer_config=dict, comm_fusion=dict)
+                 parallel_optimizer_config=dict, comm_fusion=dict, strategy_ckpt_config=dict)
 def set_auto_parallel_context(**kwargs):
     r"""
     Set auto parallel context, only data parallel supported on CPU.
@@ -464,6 +553,7 @@ def set_auto_parallel_context(**kwargs):
     enable_alltoall              grad_accumulation_step
                \                 auto_parallel_search_mode
                \                 comm_fusion
+               \                 strategy_ckpt_config
     ===========================  ===========================
 
     Args:
@@ -502,15 +592,18 @@ def set_auto_parallel_context(**kwargs):
                      data_parallel mode, all parameters are broadcast except for the parameter whose attribute
                      layerwise_parallel is True. Hybrid_parallel, semi_auto_parallel and auto_parallel mode, the
                      segmented parameters do not participate in broadcasting. Default: False.
-        strategy_ckpt_load_file (str): The path to load parallel strategy checkpoint. Default: ''
-        strategy_ckpt_save_file (str): The path to save parallel strategy checkpoint. Default: ''
+        strategy_ckpt_load_file (str): The path to load parallel strategy checkpoint. The parameter is not to be
+                       recommended currently, it is better using 'strategy_ckpt_config' to replace it. Default: ''
+        strategy_ckpt_save_file (str): The path to save parallel strategy checkpoint. The parameter is not to be
+                       recommended currently, it is better using 'strategy_ckpt_config' to replace it. Default: ''
         full_batch (bool): If you load whole batch datasets in auto_parallel mode, this parameter
                        should be set as True. Default: False. The interface is not to be recommended currently,
                        it is better using 'dataset_strategy' to replace it.
         dataset_strategy (Union[str, tuple]): Dataset sharding strategy. Default: "data_parallel".
                        dataset_strategy="data_parallel" is equal to full_batch=False, dataset_strategy="full_batch" is
-                       equal to full_batch=True. For dataset load into net by model parallel strategy likes
-                       ds_stra ((1, 8), (1, 8)), it requires using set_auto_parallel_context(dataset_strategy=ds_stra).
+                       equal to full_batch=True. For execution mode is 'GRAPH_MODE' and dataset load into net by model
+                       parallel strategy likes ds_stra ((1, 8), (1, 8)), it requires using
+                       set_auto_parallel_context(dataset_strategy=ds_stra).
         enable_parallel_optimizer (bool): This is a developing feature, which shards the weight update computation for
                        data parallel training in the benefit of time and memory saving. Currently, auto and semi auto
                        parallel mode support all optimizers in both Ascend and GPU. Data parallel mode only supports
@@ -551,6 +644,9 @@ def set_auto_parallel_context(**kwargs):
                         communication fusion config has two keys: "mode" and "config".
                         It supports following communication fusion types and configurations:
 
+                        - openstate: Whether turn on the communication fusion or not. If `openstate` is `True`, turn on
+                          the communication fusion, otherwise, turn off the communication fusion. Default: `True`.
+
                         - allreduce: If communication fusion type is `allreduce`. The `mode` contains: `auto`, `size`
                           and `index`. In `auto` mode, AllReduce fusion is configured by gradients size and the default
                           fusion threshold is `64` MB. In 'size' mode, AllReduce fusion is configured by gradients size
@@ -565,6 +661,24 @@ def set_auto_parallel_context(**kwargs):
                         - reducescatter: If communication fusion type is `reducescatter`. The `mode` contains: `auto`
                           and `size`. Config is same as `allgather`.
 
+        strategy_ckpt_config (dict): A dict contains the configurations for setting the parallel strategy file. This
+                        interface contains the functions of parameter `strategy_ckpt_load_file` and
+                        `strategy_ckpt_save_file`, it is recommonded to use this parameter to replace those two
+                        parameters.
+                        It contains following configurations:
+
+                        - load_file (str): The path to load parallel strategy checkpoint. If the file name extension is
+                          `.json`, the file is loaded in JSON format. Otherwise, the file is loaded in ProtoBuf
+                          format.
+                          Default: ''
+
+                        - save_file (str): The path to save parallel strategy checkpoint. If the file name extension is
+                          `.json`, the file is saved in JSON format. Otherwise, the file is saved in ProtoBuf format.
+                          Default: ''
+
+                        - only_trainable_params (bool): Only save/load the strategy information for trainable parameter.
+                          Default: True.
+
     Raises:
         ValueError: If input key is not attribute in auto parallel context.
 
@@ -589,6 +703,8 @@ def set_auto_parallel_context(**kwargs):
         >>> ms.set_auto_parallel_context(parallel_optimizer_config=parallel_config, enable_parallel_optimizer=True)
         >>> config = {"allreduce": {"mode": "size", "config": 32}, "allgather": {"mode": "size", "config": 32}}
         >>> ms.set_auto_parallel_context(comm_fusion=config)
+        >>> stra_ckpt_dict = {"load_file": "./stra0.ckpt", "save_file": "./stra1.ckpt", "only_trainable_params": False}
+        >>> ms.set_auto_parallel_context(strategy_ckpt_config=stra_ckpt_dict)
     """
     _set_auto_parallel_context(**kwargs)
 
@@ -616,7 +732,7 @@ def get_auto_parallel_context(attr_key):
 
 def reset_auto_parallel_context():
     """
-    Reset auto parallel context attributes to the default values:
+    Reset auto parallel context attributes to the default values.
 
     - device_num: 1.
     - global_rank: 0.
@@ -637,6 +753,46 @@ def reset_auto_parallel_context():
     _reset_auto_parallel_context()
 
 
+@args_type_check(offload_config=dict)
+def set_offload_context(offload_config):
+    r"""
+    Set offload context.
+    Some configurations are offload specific, see the below table for details:
+
+    Args:
+        offload_config (dict): A dict contains the keys and values for setting the offload context
+                        configure.It supports the following keys.
+            enable_offload (bool):  The flag of whether enabling offload. Default: False.
+            offload_param (str):  The param for offload destination, cpu or disk.
+            offload_path (str):  The path of offload.
+            offload_checkpoint (str):  The checkpoint for offload destination, cpu or disk.
+            offload_ddr_size (int):  The ddr size for offload.
+            offload_disk_size (int): The disk size for offload.
+            enable_aio (bool): The flag of whether enabling aio. Default: True.
+            aio_block_size (int): The size of aio block.
+            aio_queue_depth (int): The depth of aio queue.
+            enable_pinned_mem (bool): The flag of whether enabling pinned memory.
+
+    Raises:
+        ValueError: If input key is not attribute in auto parallel context.
+
+    Examples:
+        >>> from mindspore import context
+        >>> context.set_offload_context(offload_config={"offload_param"="cpu"})
+    """
+    _set_offload_context(offload_config)
+
+
+def get_offload_context():
+    """
+    Get offload context.
+    Examples:
+        >>> from mindspore import context
+        >>> offload_config = context.get_offload_context()
+    """
+    return _get_offload_context()
+
+
 def _check_target_specific_cfgs(device, arg_key):
     """Checking whether a config is suitable for a specified device"""
     device_cfgs = {
@@ -648,7 +804,8 @@ def _check_target_specific_cfgs(device, arg_key):
         'auto_tune_mode': ['Ascend'],
         'max_device_memory': ['Ascend', 'GPU'],
         'mempool_block_size': ['GPU', 'Ascend'],
-        'disable_format_transform': ['GPU']
+        'disable_format_transform': ['GPU'],
+        'ascend_config': ['Ascend']
     }
     # configs not in map device_cfgs are supposed to be suitable for all devices
     if arg_key not in device_cfgs:
@@ -662,16 +819,15 @@ def _check_target_specific_cfgs(device, arg_key):
     return False
 
 
-@args_unreset_check(device_id=int, variable_memory_max_size=str, max_device_memory=str, mempool_block_size=str)
-@args_type_check(mode=int, precompile_only=bool, device_target=str, device_id=int, save_graphs=bool,
+@args_type_check(mode=int, precompile_only=bool, device_target=str, device_id=int, save_graphs=(bool, int),
                  save_graphs_path=str, enable_dump=bool, auto_tune_mode=str,
                  save_dump_path=str, enable_reduce_precision=bool, variable_memory_max_size=str,
-                 enable_auto_mixed_precision=bool,
+                 enable_auto_mixed_precision=bool, inter_op_parallel_num=int,
                  enable_graph_kernel=bool, reserve_class_name_in_scope=bool, check_bprop=bool,
                  max_device_memory=str, print_file_path=str, max_call_depth=int, env_config_path=str,
                  graph_kernel_flags=str, save_compile_cache=bool, runtime_num_threads=int, load_compile_cache=bool,
                  grad_for_scalar=bool, pynative_synchronize=bool, mempool_block_size=str, disable_format_transform=bool,
-                 op_timeout=int)
+                 op_timeout=int, deterministic=str, ascend_config=dict)
 def set_context(**kwargs):
     """
     Set context for running environment.
@@ -682,7 +838,7 @@ def set_context(**kwargs):
     Note:
         Attribute name is required for setting attributes.
         The mode is not recommended to be changed after net was initialized because the implementations of some
-        operations are different in graph mode and pynative mode. Default: GRAPH_MODE.
+        operations are different in graph mode and pynative mode. Default: PYNATIVE_MODE.
 
     Some configurations are device specific, see the below table for details:
 
@@ -709,6 +865,8 @@ def set_context(**kwargs):
     |                         +------------------------------+----------------------------+
     |                         |  save_dump_path              |  Ascend                    |
     |                         +------------------------------+----------------------------+
+    |                         |  deterministic               |  Ascend                    |
+    |                         +------------------------------+----------------------------+
     |                         |  print_file_path             |  Ascend                    |
     |                         +------------------------------+----------------------------+
     |                         |  env_config_path             |  CPU/GPU/Ascend            |
@@ -737,6 +895,8 @@ def set_context(**kwargs):
     |                         +------------------------------+----------------------------+
     |                         |  enable_compile_cache        |  CPU/GPU/Ascend            |
     |                         +------------------------------+----------------------------+
+    |                         |  inter_op_parallel_num       |  CPU/GPU/Ascend            |
+    |                         +------------------------------+----------------------------+
     |                         |  runtime_num_threads         |  CPU/GPU/Ascend            |
     |                         +------------------------------+----------------------------+
     |                         |  compile_cache_path          |  CPU/GPU/Ascend            |
@@ -744,6 +904,12 @@ def set_context(**kwargs):
     |                         |  disable_format_transform    |  GPU                       |
     |                         +------------------------------+----------------------------+
     |                         |  support_binary              |  CPU/GPU/Ascend            |
+    |                         +------------------------------+----------------------------+
+    |                         |  memory_optimize_level       |  CPU/GPU/Ascend            |
+    |                         +------------------------------+----------------------------+
+    |                         |  memory_offload              |  GPU/Ascend                |
+    |                         +------------------------------+----------------------------+
+    |                         |  ascend_config               |  Ascend                    |
     +-------------------------+------------------------------+----------------------------+
 
     Args:
@@ -760,14 +926,31 @@ def set_context(**kwargs):
             of the available memory of the device and mempool_block_size.
         op_timeout (int): Set the maximum duration of executing an operator in seconds.
             If the execution time exceeds this value, system will terminate the task. 0 means endless wait.
-            Default: 600.
-        save_graphs (bool): Whether to save graphs. Default: False.
-            When the `save_graphs` attribute is set as True, attribute of `save_graphs_path` is used to set the
-            intermediate compilation graph storage path. By default, the graphs are saved in the current directory.
+            Default: 1900.
+        save_graphs (bool or int): Whether to save intermediate compilation graphs. Default: 0.
+            Available values are:
+
+            - False or 0: disable saving of intermediate compilation graphs.
+            - 1: some intermediate files will be generated during graph compliation.
+            - True or 2: Generate more ir files related to backend process.
+            - 3: Generate visualization computing graphs and detailed frontend ir graphs.
+
+            When the `save_graphs` attribute is set as True, 1, 2 or 3, attribute of `save_graphs_path` is used
+            to set the intermediate compilation graph storage path. By default, the graphs are saved in the current
+            directory.
         save_graphs_path (str): Path to save graphs. Default: ".".
             If the specified directory does not exist, the system will automatically create the directory.
             During distributed training, graphs will be saved to the directory of
             `save_graphs_path/rank_${rank_id}/`. `rank_id` is the ID of the current device in the cluster.
+        deterministic (str): Whether to enable op run in deterministic mode. The value must be in the
+            range of ['ON', 'OFF'], and the default value is 'OFF'.
+
+            - "ON": Enable operator deterministic running mode.
+            - "OFF": Disable operator deterministic running mode.
+
+            When deterministic mode is on, model ops will be deterministic in Ascend. This means that if op run multiple
+            times with the same inputs on the same hardware, it will have the exact same outputs each time. This is
+            useful for debugging models.
         enable_dump (bool): This parameters is deprecated, and will be deleted in the next version.
         save_dump_path (str): This parameters is deprecated, and will be deleted in the next version.
         print_file_path (str): The path of saving print data. If this parameter is set, print data is saved to
@@ -785,7 +968,7 @@ def set_context(**kwargs):
               When set to false, the RDR will be turned off.
             - mode: sets the mode of RDR on exporting data. When set to 1, the RDR only exports data
               in the fault scenario. When set to 2, the RDR exports data in the fault scenario and the
-              normal end scenario. Default is 1.
+              normal end scenario. Default: 1.
             - path: sets the path where RDR saves data. The current path must be absolute.
 
             Memory reuse:
@@ -810,16 +993,15 @@ def set_context(**kwargs):
             be located, when the value is set to True, the operator is executed synchronously on the device. It will
             reduce the execution performance of the program. At this time, when an error occurs in the execution of
             the operator, the location of the error script code can be located according to the call stack of the error.
-        mode (int): Running in GRAPH_MODE(0) or PYNATIVE_MODE(1). Default: GRAPH_MODE(0).
-            GRAPH_MODE or PYNATIVE_MODE can be set by `mode` attribute and both modes support all backends, default
-            mode is GRAPH_MODE.
+        mode (int): Running in GRAPH_MODE(0) or PYNATIVE_MODE(1).
+            Both modes support all backends. Default: PYNATIVE_MODE.
         enable_graph_kernel (bool): Whether to enable graph kernel fusion to optimize network execution performance.
             Default: False.
             Indicates whether to enable image-computing convergence to optimize network execution performance.
             If enable_graph_kernel is set to True, acceleration can be enabled.
             For details of graph kernel fusion, please check
             `Enabling Graph Kernel Fusion
-            <https://www.mindspore.cn/tutorials/experts/en/r1.10/debug/graph_fusion_engine.html>`_.
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/debug/graph_fusion_engine.html>`_.
         graph_kernel_flags (str):
             Optimization options of graph kernel fusion, and the priority is higher when it conflicts
             with enable_graph_kernel. Only for experienced users.
@@ -854,7 +1036,7 @@ def set_context(**kwargs):
 
             For more information about the enable operator tuning tool settings, please check
             `Enable the operator optimization tool
-            <https://www.mindspore.cn/tutorials/experts/en/r1.10/debug/auto_tune.html>`_.
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/debug/auto_tune.html>`_.
         check_bprop (bool): Whether to check back propagation nodes. The checking ensures that the shape and dtype
             of back propagation node outputs is the same as input parameters. Default: False.
         max_call_depth (int): Specify the maximum depth of function call. Must be positive integer. Default: 1000.
@@ -872,12 +1054,14 @@ def set_context(**kwargs):
             the compile cache is loaded. Note that only limited automatic detection for the changes of
             python scripts is supported by now, which means that there is a correctness risk. Default: False.
             This is an experimental prototype that is subject to change and/or deletion.
-        compile_cache_path (str): Path to save the cache of the graph compiled by front-end. Default: ".".
+        compile_cache_path (str): Path to save the compile cache. Default: ".".
             If the specified directory does not exist, the system will automatically create the directory.
             The cache will be saved to the directory of `compile_cache_path/rank_${rank_id}/`. The `rank_id` is
             the ID of the current device in the cluster.
-        runtime_num_threads(int): The thread pool number of cpu kernel and actor used in runtime,
-            which must bigger than 0. Default value is 30, if you run many processes at
+        inter_op_parallel_num(int): The thread number of op parallel at the same time. Default value is 0,
+            which means use the default num.
+        runtime_num_threads(int): The thread pool number of cpu kernel used in runtime,
+            which must bigger than or equal to 0. Default value is 30, if you run many processes at
             the same time, you should set the value smaller to avoid thread contention.
         disable_format_transform (bool): Whether to disable the automatic format transform function from NCHW to NHWC.
             When the network training performance of fp16 is worse than fp32,
@@ -886,6 +1070,49 @@ def set_context(**kwargs):
             in graph mode, coulde set 'support_binary' to be True, and run once .py file. It would save the source
             of the interfaces would be compiled by MindSpore to the interfaces definition .py file that should be
             guaranteed to be writable. Then compile the .py file to the .pyc or .so file, and could run in Graph mode.
+        memory_optimize_level (str): The memory optimize level.
+            Default: O0. The value must be in ['O0', 'O1'].
+
+            - O0: priority performance option, disable SOMAS (Safe Optimized Memory Allocation Solver).
+            - O1: priority memory option, enable SOMAS.
+        memory_offload (str): Whether to enable the memory offload function. When it is enabled, the idle data will be
+            temporarily copied to the host side in the case of insufficient device memory. The value must be in the
+            range of ['ON', 'OFF'], and the default value is 'OFF'.
+
+            - ON: Enable the memory Offload function. On Ascend hardware platform, this parameter does not take effect
+              when the environment variable "GRAPH_OP_RUN=1" is not set; This parameter does not take effect when
+              memory_optimize_level is set 'O1'.
+            - OFF: Turn off the memory Offload function.
+        ascend_config (dict): Set the parameters specific to Ascend hardware platform. It is not set by default.
+            Currently, only setting `precision_mode` and `jit_compile` are supported on Ascend910B hardware platform.
+            The default value of `precision_mode` and `jit_compile` are experimental parameters, may change
+            in the future.
+
+            - precision_mode (str): Mixed precision mode setting, on Ascend910B hardware platform, the default
+              value of training network is based on the value of CANN, and the default value of inference network
+              is force_fp16. The value range is as follows:
+
+              - force_fp16: When the operator supports both float16 and float32, select float16 directly.
+              - allow_fp32_to_fp16: When the operator does not support the float32 data type, directly reduce
+                the precision of float16.
+              - allow_mix_precision: Automatic mixing precision, facing the whole network operator, according
+                to the built-in optimization strategy, automatically reduces the precision of some operators
+                to float16 or bfloat16.
+              - must_keep_origin_dtype: Keep the accuracy of the original drawing.
+              - force_fp32: When the input of the matrix calculation operator is float16 and the output supports
+                float16 and float32, output is forced to float32.
+              - force_lowerprecision: When the operator supports both float16 or bfloat16 and float32, select
+                float16 or bfloat16 directly.
+              - allow_fp32_to_bf16: When the operator does not support the float32 data type, directly reduce
+                the precision of bfloat16.
+              - allow_fp32_to_lowprecision: When the operator does not support the float32 data type, directly
+                reduce the precision of float16 or bfloat16.
+              - allow_mix_precision_fp16: Automatic mixing precision, facing the whole network operator, automatically
+                reduces the precision of some operators to float16 according to the built-in optimization strategy.
+              - allow_mix_precision_bf16: Automatic mixing precision, facing the whole network operator, according to
+                the built-in optimization strategy, automatically reduces the precision of some operators to bfloat16.
+
+            - jit_compile (bool): Whether to select online compilation. the default value is based on CANN.
 
     Raises:
         ValueError: If input key is not an attribute in context.
@@ -913,16 +1140,17 @@ def set_context(**kwargs):
         >>> ms.set_context(enable_compile_cache=True, compile_cache_path="./cache.ms")
         >>> ms.set_context(pynative_synchronize=True)
         >>> ms.set_context(runtime_num_threads=10)
+        >>> ms.set_context(inter_op_parallel_num=4)
         >>> ms.set_context(disable_format_transform=True)
+        >>> ms.set_context(memory_optimize_level='O0')
+        >>> ms.set_context(memory_offload='ON')
+        >>> ms.set_context(deterministic='ON')
+        >>> ms.set_context(ascend_config={"precision_mode": "force_fp16", "jit_compile": True})
     """
     ctx = _context()
     # set device target first
     if 'device_target' in kwargs:
         ctx.set_device_target(kwargs['device_target'])
-        device = ctx.get_param(ms_ctx_param.device_target)
-        if not device.lower() in __device_target__:
-            raise ValueError(f"For 'context.set_context', package type {__package_name__} support 'device_target' "
-                             f"type {__device_target__}, but got {device}.")
     device = ctx.get_param(ms_ctx_param.device_target)
     for key, value in kwargs.items():
         if key == 'enable_sparse':
@@ -933,6 +1161,15 @@ def set_context(**kwargs):
             logger.warning(f"For 'context.set_context', '{key}' parameter is deprecated. "
                            "For details, please see the interface parameter API comments")
             continue
+        if key in ('precision_mode', 'jit_compile'):
+            raise ValueError(f"Please set '{key}' through parameter ascend_config")
+        if key == 'save_graphs':
+            if value is True:
+                value = 2
+            if value is False:
+                value = 0
+            if value > 3:
+                raise ValueError(f"value for save_graphs should be 0-3 but got '{value}'")
         if not _check_target_specific_cfgs(device, key):
             continue
         if hasattr(ctx, key):
@@ -1020,6 +1257,7 @@ def set_ps_context(**kwargs):
     Set parameter server training mode context.
 
     Note:
+        Parameter server mode is only supported in graph mode.
         Some other environment variables should also be set for parameter server training mode.
         These environment variables are listed below:
 
@@ -1104,9 +1342,6 @@ _hccl_connect_timeout = '600'
 
 def _init_parallel_env():
     """Set hccl connect timeout."""
-    if 'ascend' not in __device_target__:
-        return
-
     if 'HCCL_CONNECT_TIMEOUT' not in os.environ:
         os.environ['HCCL_CONNECT_TIMEOUT'] = _hccl_connect_timeout