mindspore 2.0.0rc1__cp38-none-any.whl → 2.2.0__cp38-none-any.whl

mindspore/context.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2022 Huawei Technologies Co., Ltd
+# Copyright 2020-2023 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.
@@ -35,16 +35,21 @@ from mindspore.parallel._ps_context import _set_ps_context, _get_ps_context, _re
     _need_reset_device_target_for_ps
 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', 'set_offload_context', 'get_offload_context']
+__all__ = ['GRAPH_MODE', 'PYNATIVE_MODE', 'STRICT', 'COMPATIBLE', 'LAX', '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', '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'
+_RE_PATTERN = r'[1-9][0-9]*(\.)?[0-9]*GB|0\.[0-9]*GB'
 K_CONTEXT = None
 
+# Enumerate for the property 'jit_syntax_level'.
+STRICT = 0
+COMPATIBLE = 1
+LAX = 2
+
 
 def _make_directory(path):
     """Make directory."""
@@ -156,7 +161,6 @@ class _Context:
         self._context_switches = _ContextSwitchInfo(False)
         self._context_handle = MSContext.get_instance()
         self._support_binary = False
-        self.enable_compile_cache = None
 
     def __getattribute__(self, attr):
         value = object.__getattribute__(self, attr)
@@ -201,6 +205,13 @@ class _Context:
                              f"or context.PYNATIVE_MODE (1), but got {mode}.")
         self.set_param(ms_ctx_param.mode, mode)
 
+    def set_jit_syntax_level(self, level):
+        """"Set the JIT syntax level for graph compiling"""
+        if level != STRICT and level != COMPATIBLE and level != LAX:
+            raise ValueError(f"For 'context.set_jit_syntax_level', the argument 'level' should be context.STRICT "
+                             f"or context.LAX, but got {level}.")
+        self.set_param(ms_ctx_param.jit_syntax_level, level)
+
     def set_memory_optimize_level(self, memory_optimize_level):
         """
         The memory optimize level, support "O0", "O1".
@@ -251,32 +262,91 @@ class _Context:
         Enable ascend config.
 
         Args:
-            ascend_config (dict): 'precision_mode'
+            ascend_config (dict):
                 - 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".
+                            "must_keep_origin_dtype", "force_fp32", "allow_fp32_to_bf16",
+                            "allow_mix_precision_fp16" and "allow_mix_precision_bf16".
+                - jit_compile (bool): ``False`` and ``True``.
+                - atomic_clean_policy (int): ``0`` and ``1``. Default: ``1`` .
+                - op_precision_mode (str): config file path.
+                - parallel_speed_up_json_path(Union[str, None]): The path to the parallel speed up json file.
+                  If its value is None or '', it does not take effect. Default None.
         """
-
-        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:
+        ascend_cfg_modes = {
+            'precision_mode': ["force_fp16", "allow_fp32_to_fp16", "allow_mix_precision", "must_keep_origin_dtype",
+                               "force_fp32", "allow_fp32_to_bf16", "allow_mix_precision_fp16",
+                               "allow_mix_precision_bf16"],
+            'jit_compile': [True, False],
+            'atomic_clean_policy': [0, 1],
+            'matmul_allow_hf32': [True, False],
+            'conv_allow_hf32': [True, False],
+            'op_precision_mode': (str,),
+            'parallel_speed_up_json_path': (str, None)
+        }
+        ascend_cfg_setters = {
+            'precision_mode': self._get_ascend_config_setter('precision_mode'),
+            'jit_compile': self._get_ascend_config_setter('jit_compile', lambda v: "1" if v else "0"),
+            'atomic_clean_policy': self._get_ascend_config_setter('atomic_clean_policy', str),
+            'matmul_allow_hf32': self._get_ascend_config_setter('matmul_allow_hf32', lambda v: "1" if v else "0"),
+            'conv_allow_hf32': self._get_ascend_config_setter('conv_allow_hf32', lambda v: "1" if v else "0"),
+            'op_precision_mode': self._set_op_precision_mode,
+            'parallel_speed_up_json_path': self._set_speedup_config_path
+        }
+        ascend_cfg_set = tuple(ascend_cfg_modes.keys())
+        for ascend_key, ascend_value in ascend_config.items():
+            if ascend_key not in ascend_cfg_set:
                 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:
+                                 f"{ascend_cfg_set}, but got {ascend_key}.")
+            supported_modes = ascend_cfg_modes.get(ascend_key)
+            if isinstance(supported_modes, list) and ascend_value 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")
+                                 f"{supported_modes}, but got {ascend_value}.")
+            if isinstance(supported_modes, tuple) and not isinstance(ascend_value, supported_modes):
+                raise TypeError(f"For 'ascend_config', the type of argument {ascend_key} must be one of "
+                                f"{supported_modes}, but got {type(ascend_value)}.")
+            cfg_setter = ascend_cfg_setters.get(ascend_key)
+            cfg_setter(ascend_value)
+
+    def set_gpu_config(self, gpu_config):
+        """
+        Enable gpu config.
+
+        Args:
+            gpu_config (dict):
+
+                - conv_fprop_algo (str): "normal", "performance" or user specifies conv forward algorithm directly.
+                - conv_dgrad_algo (str): "normal", "performance" or user specifies conv data grad algorithm directly.
+                - conv_wgrad_algo (str): "normal", "performance" or user specifies conv weight grad algorithm directly.
+                - conv_allow_tf32 (bool): ``False`` and ``True``.
+                - matmul_allow_tf32 (bool): ``False`` and ``True``.
+        """
+
+        gpu_cfgs = {'conv_fprop_algo': ["normal", "performance", "implicit_gemm", "precomp_gemm", "gemm", "direct",
+                                        "fft", "fft_tiling", "winograd", "winograd_nonfused"],
+                    'conv_dgrad_algo': ["normal", "performance", "algo_0", "algo_1", "fft", "fft_tiling", "winograd",
+                                        "winograd_nonfused"],
+                    'conv_wgrad_algo': ["normal", "performance", "algo_0", "algo_1", "fft", "algo_3", "fft_tiling",
+                                        "winograd_nonfused"],
+                    'conv_allow_tf32': [True, False],
+                    'matmul_allow_tf32': [True, False]}
+        for gpu_key in gpu_config:
+            if gpu_key not in gpu_cfgs:
+                raise ValueError(f"For 'context.set_context', the key of argument 'gpu_config' must be one of "
+                                 f"{gpu_cfgs}, but got {gpu_key}.")
+            supported_value = gpu_cfgs.get(gpu_key)
+            if gpu_config[gpu_key] not in supported_value:
+                raise ValueError(f"For 'gpu_config', the value of argument {gpu_key} must be one of "
+                                 f"{supported_value}, but got {gpu_config[gpu_key]}.")
+            if gpu_key == 'conv_fprop_algo':
+                self.set_param(ms_ctx_param.conv_fprop_algo, gpu_config[gpu_key])
+            if gpu_key == 'conv_dgrad_algo':
+                self.set_param(ms_ctx_param.conv_dgrad_algo, gpu_config[gpu_key])
+            if gpu_key == 'conv_wgrad_algo':
+                self.set_param(ms_ctx_param.conv_wgrad_algo, gpu_config[gpu_key])
+            if gpu_key == 'conv_allow_tf32':
+                self.set_param(ms_ctx_param.conv_allow_tf32, gpu_config[gpu_key])
+            if gpu_key == 'matmul_allow_tf32':
+                self.set_param(ms_ctx_param.matmul_allow_tf32, gpu_config[gpu_key])
 
     def set_backend_policy(self, policy):
         success = self._context_handle.set_backend_policy(policy)
@@ -312,13 +382,42 @@ class _Context:
         if self.enable_debug_runtime and target == "CPU":
             self.set_backend_policy("vm")
 
-    def set_auto_tune_mode(self, tune_mode):
-        candidate = ["NO_TUNE", "RL", "GA", "RL,GA", "GA,RL"]
+    def set_aoe_tune_mode(self, tune_mode):
+        """
+        Set aoe tune mode, support "online" and "offline".
+
+        Args:
+            tune_mode (str): "online" and "offline".
+        """
+        candidate = ["online", "offline"]
         if tune_mode in candidate:
-            self.set_param(ms_ctx_param.auto_tune_mode, tune_mode)
+            self.set_param(ms_ctx_param.aoe_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}.")
+            raise ValueError(f"For 'context.set_context', the argument 'aoe_tune_mode' must be in "
+                             f"['online', 'offline'], but got {tune_mode}.")
+
+    def set_aoe_config(self, aoe_config):
+        """
+        Enable aoe config.
+
+        Args:
+            aoe_config (dict):
+                - job_type (str): ``"1"``, ``"2"``. Default: ``"2"`` .
+                  - ``"1"``: subgraph tuning.
+                  - ``"2"``: operator tuning.
+        """
+
+        aoe_cfgs = {'job_type': ["1", "2"]}
+        for aoe_config_key in aoe_config:
+            if aoe_config_key not in aoe_cfgs:
+                raise ValueError(f"For 'context.set_context', the key of argument 'aoe_config' must be one of "
+                                 f"{aoe_cfgs}, but got {aoe_config_key}.")
+            supported_value = aoe_cfgs.get(aoe_config_key)
+            if aoe_config[aoe_config_key] not in supported_value:
+                raise ValueError(f"For 'aoe_config', the value of argument {aoe_config_key} must be one of "
+                                 f"{supported_value}, but got {aoe_config[aoe_config_key]}.")
+            if aoe_config_key == 'job_type':
+                self.set_param(ms_ctx_param.aoe_job_type, aoe_config[aoe_config_key])
 
     def set_device_id(self, device_id):
         if device_id < 0 or device_id > 4095:
@@ -343,14 +442,14 @@ class _Context:
         logger.warning("For 'context.set_context', the parameter 'variable_memory_max_size' is deprecated, "
                        "and will be removed in a future "
                        "version. Please use parameter 'max_device_memory' instead.")
-        if not Validator.check_str_by_regular(variable_memory_max_size, _re_pattern):
+        if not Validator.check_str_by_regular(variable_memory_max_size, _RE_PATTERN):
             raise ValueError("For 'context.set_context', the argument 'variable_memory_max_size' should be in correct"
                              " format! It must be a string ending with 'GB', in addition to that, it must contain "
-                             "only numbers or decimal points, such as \"5GB\" or \"3.5GB\", but got {}."
+                             "only numbers or decimal points, such as \"5GB\" or \"3.5GB\", but got {}GB."
                              .format(variable_memory_max_size))
         if float(variable_memory_max_size[:-2]) > _DEVICE_APP_MEMORY_SIZE:
             raise ValueError("For 'context.set_context', the argument 'variable_memory_max_size' should not be "
-                             "greater than 31GB, but got {}.".format(variable_memory_max_size))
+                             "greater than 31GB, but got {}GB.".format(variable_memory_max_size))
         variable_memory_max_size_ = variable_memory_max_size[:-2] + " * 1024 * 1024 * 1024"
         graph_memory_max_size = _DEVICE_APP_MEMORY_SIZE - int(variable_memory_max_size[:-2])
         graph_memory_max_size_ = str(graph_memory_max_size) + " * 1024 * 1024 * 1024"
@@ -358,7 +457,7 @@ class _Context:
         self.set_param(ms_ctx_param._graph_memory_max_size, graph_memory_max_size_)
 
     def set_max_device_memory(self, max_device_memory):
-        if not Validator.check_str_by_regular(max_device_memory, _re_pattern):
+        if not Validator.check_str_by_regular(max_device_memory, _RE_PATTERN):
             raise ValueError("For 'context.set_context', the argument 'max_device_memory' should be in correct "
                              " format! It must be a string ending with 'GB', in addition to that, it must contain "
                              "only numbers or decimal points, such as \"5GB\" or \"3.5GB\", but got {}."
@@ -374,7 +473,7 @@ class _Context:
             logger.warning("Graph mode doesn't support to set parameter 'mempool_block_size' of context currently, "
                            "you can use context.set_context to set pynative mode.")
             return
-        if not Validator.check_str_by_regular(mempool_block_size, _re_pattern):
+        if not Validator.check_str_by_regular(mempool_block_size, _RE_PATTERN):
             raise ValueError("For 'context.set_context', the argument 'mempool_block_size' should be in "
                              "correct format! Such as \"10GB\", "
                              "but got {}".format(mempool_block_size))
@@ -416,7 +515,7 @@ class _Context:
         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))
+                                        "or may not have permission to read it.".format(env_config_path)) from exo
         self.set_param(ms_ctx_param.env_config_path, env_config_path)
 
     def set_runtime_num_threads(self, runtime_num_threads):
@@ -441,8 +540,8 @@ class _Context:
         'mode': set_mode,
         'save_graphs_path': set_save_graphs_path,
         'device_target': set_device_target,
+        'aoe_tune_mode': set_aoe_tune_mode,
         'device_id': set_device_id,
-        'auto_tune_mode': set_auto_tune_mode,
         'max_call_depth': set_max_call_depth,
         'profiling_options': set_profiling_options,
         'variable_memory_max_size': set_variable_memory_max_size,
@@ -456,7 +555,10 @@ class _Context:
         'op_timeout': set_op_timeout,
         'memory_offload': set_memory_offload,
         'deterministic': set_deterministic,
-        'ascend_config': set_ascend_config
+        'ascend_config': set_ascend_config,
+        'jit_syntax_level': set_jit_syntax_level,
+        'gpu_config': set_gpu_config,
+        'aoe_config': set_aoe_config,
     }
 
     @property
@@ -496,6 +598,55 @@ class _Context:
             raise TypeError(f"The attribute 'support_binary' should be a bool, but got {type(support)}.")
         self._support_binary = support
 
+    def _get_ascend_config_setter(self, ascend_key, trans_fn=None):
+        def _config_setter(ascend_value):
+            self.set_param(ms_ctx_param.__members__[ascend_key], trans_fn(ascend_value))
+
+        if trans_fn is None:
+            trans_fn = lambda x: x
+        return _config_setter
+
+    def _set_op_precision_mode(self, ascend_value):
+        op_precision_path = ascend_value
+        real_path = os.path.realpath(op_precision_path)
+        if not os.path.exists(real_path):
+            raise ValueError(f"For 'ascend_config', the 'op_precision_mode' is invalid path, "
+                             f"got '{op_precision_path}'.")
+        self.set_param(ms_ctx_param.op_precision_mode, ascend_value)
+
+    def _set_speedup_config_path(self, speedup_config_path):
+        """"Check and set speedup config for auto parallel."""
+        if speedup_config_path is None or speedup_config_path == "":
+            return
+        speedup_config_real_path = os.path.abspath(speedup_config_path)
+        if not os.path.exists(speedup_config_real_path):
+            raise ValueError(f"For 'ascend_config', the path to parallel_speed_up_json: "
+                             f"{speedup_config_real_path} does not exist, please check whether the "
+                             f"'parallel_speed_up_json_path' is correct.")
+        try:
+            valid_option = {"recompute_comm_overlap": ms_ctx_param.recompute_comm_overlap,
+                            "matmul_grad_comm_overlap": ms_ctx_param.matmul_grad_comm_overlap,
+                            "enable_task_opt": ms_ctx_param.enable_task_opt,
+                            "enable_grad_comm_opt": ms_ctx_param.enable_grad_comm_opt,
+                            "interleaved_matmul_comm": ms_ctx_param.interleaved_matmul_comm,
+                            "interleaved_layernorm_comm": ms_ctx_param.interleaved_layernorm_comm}
+            with open(speedup_config_real_path, 'r') as f:
+                speedup_config = json.load(f)
+                for k, v in speedup_config.items():
+                    if not isinstance(k, str):
+                        raise TypeError("key {} is not a str".format(k))
+                    if k not in valid_option:
+                        raise ValueError("key {} should be one of {}.".format(k, valid_option.keys()))
+                    if not isinstance(v, bool):
+                        raise TypeError("value {} is not a bool".format(v))
+                    self.set_param(valid_option.get(k), v)
+        except (TypeError, ValueError) as exo:
+            raise ValueError(str(exo) + "\nFor 'context.set_context', "
+                                        "open or load the 'speedup_config_path' file {} "
+                                        "failed, please check whether 'speedup_config_path' is json file and correct, "
+                                        "or may not have permission to read it.".format(speedup_config_real_path)) \
+                                        from exo
+
 
 def _context():
     """
@@ -524,8 +675,9 @@ def _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, strategy_ckpt_config=dict)
+                 all_reduce_fusion_config=list, pipeline_stages=int, pipeline_segments=int,
+                 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.
@@ -546,26 +698,26 @@ def set_auto_parallel_context(**kwargs):
     device_num                   gradient_fp32_sync
     global_rank                  loss_repeated_mean
     gradients_mean               search_mode
-    parallel_mode                strategy_ckpt_load_file
-    all_reduce_fusion_config     strategy_ckpt_save_file
-    enable_parallel_optimizer    dataset_strategy
-    parallel_optimizer_config    pipeline_stages
-    enable_alltoall              grad_accumulation_step
+    parallel_mode                parameter_broadcast
+    all_reduce_fusion_config     strategy_ckpt_load_file
+    enable_parallel_optimizer    strategy_ckpt_save_file
+    parallel_optimizer_config    dataset_strategy
+    enable_alltoall              pipeline_stages
                \                 auto_parallel_search_mode
                \                 comm_fusion
                \                 strategy_ckpt_config
     ===========================  ===========================
 
     Args:
-        device_num (int): Available device number, the value must be in [1, 4096]. Default: 1.
-        global_rank (int): Global rank id, the value must be in [0, 4095]. Default: 0.
+        device_num (int): Available device number, the value must be in [1, 4096]. Default: ``1`` .
+        global_rank (int): Global rank id, the value must be in [0, 4095]. Default: ``0`` .
         gradients_mean (bool): Whether to perform mean operator after allreduce of gradients.
-                     "stand_alone" do not support gradients_mean. Default: False.
+                     "stand_alone" do not support gradients_mean. Default: ``False`` .
         gradient_fp32_sync (bool): Run allreduce of gradients in fp32. "stand_alone", "data_parallel"
-                     and "hybrid_parallel" do not support gradient_fp32_sync. Default: True.
-        parallel_mode (str): There are five kinds of parallel modes, "stand_alone", "data_parallel",
-                     "hybrid_parallel", "semi_auto_parallel" and "auto_parallel". Note the pynative mode only supports
-                     the "stand_alone" and "data_parallel" mode. Default: "stand_alone".
+                     and "hybrid_parallel" do not support gradient_fp32_sync. Default: ``True`` .
+        parallel_mode (str): There are five kinds of parallel modes, ``"stand_alone"`` , ``"data_parallel"`` ,
+                     ``"hybrid_parallel"`` , ``"semi_auto_parallel"`` and ``"auto_parallel"`` . Note the pynative mode
+                     only supports the ``"stand_alone"`` and ``"data_parallel"`` mode. Default: ``"stand_alone"`` .
 
                      - stand_alone: Only one processor is working.
 
@@ -576,10 +728,12 @@ def set_auto_parallel_context(**kwargs):
                      - semi_auto_parallel: Achieves data and model parallelism by setting parallel strategies.
 
                      - auto_parallel: Achieving parallelism automatically.
-        search_mode (str): There are three kinds of shard strategy search modes: "recursive_programming",
-                     "dynamic_programming" and "sharding_propagation". Default: "dynamic_programming".
+        search_mode (str): There are three kinds of shard strategy search modes: ``"recursive_programming"`` ,
+                     ``"dynamic_programming"`` and ``"sharding_propagation"`` . Default: ``"recursive_programming"`` .
 
-                     - recursive_programming: Recursive programming search mode.
+                     - recursive_programming: Recursive programming search mode. In order to obtain optimal performance,
+                       it is recommended that users set the batch size to be greater than or equal to the product of
+                       the number of devices and the number of multi-copy parallelism.
 
                      - dynamic_programming: Dynamic programming search mode.
 
@@ -589,17 +743,18 @@ def set_auto_parallel_context(**kwargs):
         parameter_broadcast (bool): Whether to broadcast parameters before training. Before training, in order to have
                      the same network initialization parameter values for all devices, broadcast the parameters
                      on device 0 to other devices. Parameter broadcasting in different parallel modes is different,
-                     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.
+                     ``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. The parameter is not to be
-                       recommended currently, it is better using 'strategy_ckpt_config' to replace it. Default: ''
+                       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".
+                       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 execution mode is 'GRAPH_MODE' and dataset load into net by model
                        parallel strategy likes ds_stra ((1, 8), (1, 8)), it requires using
@@ -607,45 +762,51 @@ def set_auto_parallel_context(**kwargs):
         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
-                       `Lamb` and `AdamWeightDecay` in Ascend . Default: False.
+                       `Lamb` and `AdamWeightDecay` in Ascend . Default: ``False`` .
         enable_alltoall (bool): A switch that allows AllToAll operators to be generated during communication. If its
-                        value is False, there will be a combination of operators such as AllGather, Split and Concat
-                        instead of AllToAll. Default: False.
+                        value is ``False`` , there will be a combination of operators such as AllGather, Split and
+                        Concat instead of AllToAll. Default: ``False`` .
         all_reduce_fusion_config (list): Set allreduce fusion strategy by parameters indices. Only support ReduceOp.SUM
                        and HCCL_WORLD_GROUP/NCCL_WORLD_GROUP. No Default, if it is not set, the fusion is closed.
         pipeline_stages (int): Set the stage information for pipeline parallel. This indicates how the devices are
                         distributed alone in the pipeline. The total devices will be divided into 'pipeline_stags'
-                        stages. Currently, this could only be used when parallel mode semi_auto_parallel is enabled.
-                        Default: 1.
-        grad_accumulation_step (int): Set the accumulation steps of gradients in auto and semi auto parallel mode.
-                        This should be a positive int. Default: 1.
+                        stages.
+                        Default: ``1`` .
         parallel_optimizer_config (dict): A dict contains the keys and values for setting the parallel optimizer
                         configure. The configure provides more detailed behavior control about parallel training
-                        when parallel optimizer is enabled. Currently it supports the key `gradient_accumulation_shard`.
-                        The configure will be effective when we use
+                        when parallel optimizer is enabled. The configure will be effective when we use
                         mindspore.set_auto_parallel_context(enable_parallel_optimizer=True).
                         It supports the following keys.
 
-                        - gradient_accumulation_shard(bool): If true, the accumulation gradient parameters will be
+                        - gradient_accumulation_shard(bool): If ``true`` , the accumulation gradient parameters will be
                           sharded across the data parallel devices. This will
                           introduce additional communication(ReduceScatter) at
                           each step when accumulate the gradients, but saves a
                           lot of device memories, thus can make model be trained
                           with larger batch size. This configure is effective only
                           when the model runs on pipeline training or gradient
-                          accumulation with data parallel. Default True.
+                          accumulation with data parallel. Default ``False`` .
 
                         - parallel_optimizer_threshold(int): Set the threshold of parallel optimizer. When parallel
                           optimizer is enabled, parameters with size smaller than this threshold will not be sharded
                           across the devices. Parameter size = shape[0] \* ... \* shape[n] \* size(dtype). Non-negative.
-                          Unit: KB. Default: 64.
+                          Unit: KB. Default: ``64`` .
+
+                        - optimizer_weight_shard_size(int): Set the optimizer weight shard group size, if you want to
+                          specific the maximum group size across devices when the parallel optimizer is enabled.
+                          The numerical range can be (0, device_num]. If pipeline parallel is enabled, the numerical
+                          range is (0, device_num/stage]. If the size of data parallel communication domain
+                          of the parameter cannot be divided by `optimizer_weight_shard_size`, then the specified
+                          communication group size will not take effect. Default value is ``-1`` , which means the
+                          optimizer weight shard group size will be the size of data parallel group of each parameter.
 
         comm_fusion (dict): A dict contains the types and configurations for setting the communication fusion. each
                         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`.
+                        - 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
@@ -677,7 +838,7 @@ def set_auto_parallel_context(**kwargs):
                           Default: ''
 
                         - only_trainable_params (bool): Only save/load the strategy information for trainable parameter.
-                          Default: True.
+                          Default: ``True`` .
 
     Raises:
         ValueError: If input key is not attribute in auto parallel context.
@@ -699,7 +860,8 @@ def set_auto_parallel_context(**kwargs):
         >>> ms.set_auto_parallel_context(enable_alltoall=False)
         >>> ms.set_auto_parallel_context(all_reduce_fusion_config=[8, 160])
         >>> ms.set_auto_parallel_context(pipeline_stages=2)
-        >>> parallel_config = {"gradient_accumulation_shard": True, "parallel_optimizer_threshold": 24}
+        >>> parallel_config = {"gradient_accumulation_shard": True, "parallel_optimizer_threshold": 24,
+        ...                    "optimizer_weight_shard_size": 2}
         >>> 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)
@@ -739,8 +901,8 @@ def reset_auto_parallel_context():
     - gradients_mean: False.
     - gradient_fp32_sync: True.
     - parallel_mode: 'stand_alone'.
-    - search_mode: 'dynamic_programming'.
-    - auto_parallel_search_mode: 'dynamic_programming'.
+    - search_mode: 'recursive_programming'.
+    - auto_parallel_search_mode: 'recursive_programming'.
     - parameter_broadcast: False.
     - strategy_ckpt_load_file: ''.
     - strategy_ckpt_save_file: ''.
@@ -749,6 +911,10 @@ def reset_auto_parallel_context():
     - enable_alltoall: False.
     - pipeline_stages: 1.
     - fusion_threshold: 64.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> ms.reset_auto_parallel_context()
     """
     _reset_auto_parallel_context()
 
@@ -756,36 +922,51 @@ def 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:
+    Configure heterogeneous training detailed parameters to adjust the offload strategy.
+
+    Note:
+        The offload configuration is only used if the memory offload feature is enabled
+        via mindspore.set_context(memory_offload="ON").
 
     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.
+            configure.It supports the following keys.
+
+            - offload_path (str):  The path of offload, relative paths are supported. Default: ``"./offload"``.
+            - offload_cpu_size (str):  The cpu memory size for offload. The format is "xxGB".
+            - offload_disk_size (str): The disk size for offload. The format is "xxGB"
+            - hbm_ratio (float): The ratio that can be used based on the maximum device memory.
+              The range is (0,1], Default: ``1.0``.
+            - cpu_ratio (float): The ratio that can be used based on the maximum host memory.
+              The range is (0,1], Default: ``1.0``.
+            - enable_pinned_mem (bool): The flag of whether enabling Pinned Memory. Default: ``True``.
+            - enable_aio (bool): The flag of whether enabling aio. Default: ``True``.
+            - aio_block_size (str): The size of aio block. The format is "xxGB".
+            - aio_queue_depth (int): The depth of aio queue.
+            - offload_param (str):  The param for offload destination, cpu or disk, Default: ``""``.
+            - offload_checkpoint (str):  The checkpoint for offload destination, only valid if recompute is turned on,
+              cpu or disk, Default: ``""``.
+            - auto_offload (bool): The flag of whether auto offload. Default: ``True``.
+            - host_mem_block_size (str): The memory block size of host memory pool. The format is "xxGB"
 
     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"})
+        >>> context.set_offload_context(offload_config={"offload_param":"cpu"})
     """
     _set_offload_context(offload_config)
 
 
 def get_offload_context():
     """
-    Get offload context.
+    Gets the offload configuration parameters. Configure through interface mindspore.set_offload_context().
+    If the user is not set, the default configuration is obtained.
+
+    Returns:
+        Dict, heterogeneous training offload detailed configuration parameters.
+
     Examples:
         >>> from mindspore import context
         >>> offload_config = context.get_offload_context()
@@ -801,11 +982,11 @@ def _check_target_specific_cfgs(device, arg_key):
         'enable_reduce_precision': ['Ascend'],
         'print_file_path': ['Ascend'],
         'variable_memory_max_size': ['Ascend'],
-        'auto_tune_mode': ['Ascend'],
         'max_device_memory': ['Ascend', 'GPU'],
         'mempool_block_size': ['GPU', 'Ascend'],
         'disable_format_transform': ['GPU'],
-        'ascend_config': ['Ascend']
+        'ascend_config': ['Ascend'],
+        'gpu_config': ['GPU'],
     }
     # configs not in map device_cfgs are supposed to be suitable for all devices
     if arg_key not in device_cfgs:
@@ -820,14 +1001,15 @@ def _check_target_specific_cfgs(device, arg_key):
 
 
 @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_graphs_path=str, enable_dump=bool, aoe_tune_mode=str, aoe_config=dict,
                  save_dump_path=str, enable_reduce_precision=bool, variable_memory_max_size=str,
                  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, deterministic=str, ascend_config=dict)
+                 op_timeout=int, deterministic=str, ascend_config=dict, jit_syntax_level=int,
+                 jit_enable_inplace_ops=bool, gpu_config=dict)
 def set_context(**kwargs):
     """
     Set context for running environment.
@@ -838,7 +1020,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: PYNATIVE_MODE.
+        operations are different in graph mode and pynative mode. Default: ``PYNATIVE_MODE`` .
 
     Some configurations are device specific, see the below table for details:
 
@@ -875,7 +1057,7 @@ def set_context(**kwargs):
     |                         +------------------------------+----------------------------+
     |                         |  reserve_class_name_in_scope |  CPU/GPU/Ascend            |
     |                         +------------------------------+----------------------------+
-    |                         |  pynative_synchronize        |  GPU/Ascend                |
+    |                         |  pynative_synchronize        |  CPU/GPU/Ascend            |
     +-------------------------+------------------------------+----------------------------+
     | Executive Control       |   mode                       |   CPU/GPU/Ascend           |
     |                         +------------------------------+----------------------------+
@@ -885,7 +1067,9 @@ def set_context(**kwargs):
     |                         +------------------------------+----------------------------+
     |                         |  enable_reduce_precision     |  Ascend                    |
     |                         +------------------------------+----------------------------+
-    |                         |  auto_tune_mode              |  Ascend                    |
+    |                         |  aoe_tune_mode               |  Ascend                    |
+    |                         +------------------------------+----------------------------+
+    |                         |  aoe_config                  |  Ascend                    |
     |                         +------------------------------+----------------------------+
     |                         |  check_bprop                 |  CPU/GPU/Ascend            |
     |                         +------------------------------+----------------------------+
@@ -910,47 +1094,52 @@ def set_context(**kwargs):
     |                         |  memory_offload              |  GPU/Ascend                |
     |                         +------------------------------+----------------------------+
     |                         |  ascend_config               |  Ascend                    |
+    |                         +------------------------------+----------------------------+
+    |                         |  jit_syntax_level            |  CPU/GPU/Ascend            |
+    |                         +------------------------------+----------------------------+
+    |                         |  gpu_config                  |  GPU                       |
     +-------------------------+------------------------------+----------------------------+
 
     Args:
         device_id (int): ID of the target device, the value must be in [0, device_num_per_host-1],
-            while device_num_per_host should be no more than 4096. Default: 0.
+            while device_num_per_host should be no more than 4096. Default: ``0`` .
         device_target (str): The target device to run, support "Ascend", "GPU", and "CPU".
             If device target is not set, the version of MindSpore package is used.
-        max_device_memory (str): Set the maximum memory available for devices. The format is "xxGB". Default: "1024GB".
-            The actual used memory size is the minimum of the available memory of the device and max_device_memory.
+        max_device_memory (str): Set the maximum memory available for devices. The format is "xxGB".
+            Default: ``" 1024GB"`` . The actual used memory size is the minimum of the available memory of the device
+            and max_device_memory. 'max_device_memory' should be set before the program runs.
         variable_memory_max_size (str): This parameter is deprecated, and will be removed in a future version.
             Please use parameter 'max_device_memory' instead.
         mempool_block_size (str): Set the size of the memory pool block in PyNative mode for devices.
-            The format is "xxGB". Default: "1GB". Minimum size is "1G". The actual used memory block size is the minimum
-            of the available memory of the device and mempool_block_size.
+            The format is "xxGB". Default: ``"1GB"`` . Minimum size is "1G". The actual used memory block size is the
+            minimum 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: 1900.
-        save_graphs (bool or int): Whether to save intermediate compilation graphs. Default: 0.
+            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.
+            - 1: some intermediate files will be generated during graph compilation.
             - 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.
+            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'.
+            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.
+            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
@@ -964,23 +1153,23 @@ def set_context(**kwargs):
             configure RDR:
 
             - enable: controls whether the RDR is enabled to collect the key data during training and
-              save key data in the fault scenario. When set to true, the RDR will be turned on.
-              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: 1.
+              save key data in the fault scenario. When set to ``true`` , the RDR will be turned on.
+              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: ``1`` .
             - path: sets the path where RDR saves data. The current path must be absolute.
 
             Memory reuse:
 
-            - mem_Reuse: controls whether the memory reuse function is turned on. When set to True,
-            - the memory reuse function is turned on. When set to False, the memory reuse function is turned off.
+            - mem_Reuse: controls whether the memory reuse function is turned on. When set to ``True`` ,
+              the memory reuse function is turned on. When set to ``False`` , the memory reuse function is turned off.
 
-        precompile_only (bool): Whether to only precompile the network. Default: False.
-            If set to True, the network will only be compiled, not executed.
-        reserve_class_name_in_scope (bool) : Whether to save the network class name in the scope. Default: True.
+        precompile_only (bool): Whether to only precompile the network. Default: ``False`` .
+            If set to ``True`` , the network will only be compiled, not executed.
+        reserve_class_name_in_scope (bool) : Whether to save the network class name in the scope. Default: ``True`` .
             Each node has a scope. A scope of a subnode is the name of its parent node. If reserve_class_name_in_scope
-            is set to True, the class name will be saved after keyword 'net-' in the scope.
+            is set to ``True`` , the class name will be saved after keyword 'net-' in the scope.
             For example:
 
             Default/net-Net1/net-Net2 (reserve_class_name_in_scope=True)
@@ -988,27 +1177,34 @@ def set_context(**kwargs):
             Default/net/net (reserve_class_name_in_scope=False)
 
         pynative_synchronize (bool): Whether to enable synchronous execution of the device in PyNative mode.
-            Default: False. When the value is set to False, the operator is executed asynchronously on the device.
-            When an error occurs in the execution of the operator, the specific error script code location cannot
-            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.
+            Default: ``False`` . When the value is set to ``False`` , the operator is executed asynchronously on the
+            device. When an error occurs in the execution of the operator, the specific error script code location
+            cannot 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).
-            Both modes support all backends. Default: PYNATIVE_MODE.
+            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.
+            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.
+            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/r2.0/debug/graph_fusion_engine.html>`_.
+            <https://www.mindspore.cn/tutorials/experts/en/r2.2/optimize/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.
-            For example, mindspore.set_context(graph_kernel_flags="--opt_level=2 --dump_as_text"). Some general options:
+            For example,
+
+            .. code-block::
+
+                mindspore.set_context(graph_kernel_flags="--opt_level=2 --dump_as_text")
+
+            Some general options:
 
             - opt_level: Set the optimization level.
-              Default: 2. Graph kernel fusion can be enabled equivalently by setting opt_level greater than 0.
+              Default: ``2`` . Graph kernel fusion can be enabled equivalently by setting opt_level greater than 0.
               Available values are:
 
               - 0: disables graph kernel fusion;
@@ -1019,55 +1215,52 @@ def set_context(**kwargs):
                 ParallelFusion and so on. Optimizations of this level are radical and unstable in some scenarios.
                 Be caution when using this level.
 
-            - dump_as_text: dumps detail info as text files. Default: false.
+            - dump_as_text: dumps detail info as text files. Default: ``False`` .
 
-            More options can refer to the implementation code.
         enable_reduce_precision (bool): Whether to enable precision reduction.
             If the operator does not support the user-specified precision, the precision will
-            be changed automatically. Default: True.
-        auto_tune_mode (str): The mode of auto tune when op building, get the best tiling performance.
-            Default: NO_TUNE. The value must be in ['RL', 'GA', 'RL,GA'].
-
-            - RL: Reinforcement Learning tune.
-            - GA: Genetic Algorithm tune.
-            - RL,GA: When both RL and GA optimization are enabled, the tool automatically selects RL or GA based on
-              different types of operators in the network model. The sequence of RL and GA is not differentiated.
-              (Automatic selection).
-
-            For more information about the enable operator tuning tool settings, please check
-            `Enable the operator optimization tool
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0/debug/auto_tune.html>`_.
+            be changed automatically. Default: ``True`` .
+        aoe_tune_mode (str): AOE tuning mode setting, which is not set by default.
+            When set to ``"online"`` , the tuning in online function is turned on.
+            When set to ``"offline"`` , ge graph will be save for offline tuning.
+        aoe_config (dict): Set the parameters specific to Ascend Optimization Engine. It is not set by default.
+
+            - job_type (str): Mode type setting, default value is ``"2"``.
+
+              - ``"1"``: subgraph tuning;
+              - ``"2"``: operator tuning.
+
         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.
+            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`` .
             The max_call_depth parameter needs to be set when the nested call is too deep or the number
             of subgraphs is too large. If max_call_depth is set larger than before, the system max stack depth should be
             set larger too, otherwise a `core dumped` exception may be raised because of system stack overflow.
-        grad_for_scalar (bool):  Whether to get gradient for scalar. Default: False.
-            When grad_for_scalar is set to True, the function's scalar input can be derived.
-            The default value is False. Because the back-end does not support scaling operations currently,
+        grad_for_scalar (bool):  Whether to get gradient for scalar. Default: ``False`` .
+            When grad_for_scalar is set to ``True`` , the function's scalar input can be derived.
+            The default value is ``False`` . Because the back-end does not support scaling operations currently,
             this interface only supports simple operations that can be deduced by the front-end.
         enable_compile_cache (bool): Whether to save or load the cache of the graph compiled by front-end.
-            After enable_compile_cache is set to True, during the first execution, a hardware-independent
+            After enable_compile_cache is set to ``True`` , during the first execution, a hardware-independent
             compilation cache is generated and exported to a MINDIR file. When the network is executed again,
-            if enable_compile_cache is still set to True and the network scripts are not changed,
+            if enable_compile_cache is still set to ``True`` and the network scripts are not changed,
             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.
+            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 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.
-        inter_op_parallel_num(int): The thread number of op parallel at the same time. Default value is 0,
+        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
+            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,
-            `disable_format_transform` can be set to True to try to improve training performance. Default: False.
+            When the network training performance of fp16 is worse than fp32, `disable_format_transform` can be set to
+            ``True`` to try to improve training performance. Default: ``False`` .
         support_binary (bool): Whether to support run .pyc or .so in graph mode. If want to support run .so or .pyc
-            in graph mode, coulde set 'support_binary' to be True, and run once .py file. It would save the source
+            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.
@@ -1077,20 +1270,20 @@ def set_context(**kwargs):
             - 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'.
+            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.
+            Currently, configurations except `parallel_speed_up_json_path` and `precision_mode.force_fp32` are only
+            supported on Ascend910B hardware platform. The default value of `precision_mode`, `jit_compile` and
+            `atomic_clean_policy` 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:
+              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
@@ -1101,18 +1294,123 @@ def set_context(**kwargs):
               - 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.
+            - atomic_clean_policy (int): The policy for cleaning memory occupied by atomic operators in the network.
+              Default: ``1`` .
+
+              - 0: The memory occupied by all atomic operators in the network is cleaned centrally.
+              - 1: Memory is not cleaned centrally and each atomic operator in the network is cleaned separately.
+                When the memory of the network exceeds the limit, you may try this cleaning policy, but it may cause
+                performance loss.
+            - matmul_allow_hf32 (bool): Whether to convert FP32 to HF32 for Matmul operators. Default value: ``False``.
+              This is an experimental prototype that is subject to change and/or deletion.
+              For detailed information, please refer to `Ascend community <https://www.hiascend.com/>`_ .
+            - conv_allow_hf32 (bool): Whether to convert FP32 to HF32 for Conv operators. Default value: ``True``.
+              This is an experimental prototype that is subject to change and/or deletion.
+              For detailed information, please refer to `Ascend community <https://www.hiascend.com/>`_ .
+            - op_precision_mode (str): Path to config file of op precision mode. For detailed information, please refer
+              to `Ascend community <https://www.hiascend.com/>`_ .
+            - parallel_speed_up_json_path(Union[str, None]): The path to the parallel speed up json file, configuration
+              can refer to `parallel_speed_up.json
+              <https://gitee.com/mindspore/mindspore/blob/r2.2/config/parallel_speed_up.json>`_ .
+              If its value is None or '', it does not take effect. Default None.
+
+              - recompute_comm_overlap (bool): Enable overlap between recompute ops and communication ops if True.
+                Default: False.
+              - matmul_grad_comm_overlap (bool): Enable overlap between grad ops and communication ops if True.
+                Default: False.
+              - enable_task_opt (bool): Enable the optimization of the number of tasks for each communication if True.
+                Default: False.
+              - interleaved_matmul_comm (bool): Enable interleaved optimization of Matmul-Comm if True. Default: False.
+              - interleaved_layernorm_comm (bool): Enable interleaved optimization of LayerNorm-Comm if True.
+                Default: False.
+
+        jit_syntax_level (int): Set JIT syntax level for graph compiling, triggered by GRAPH_MODE and @jit decorator.
+            The value must be ``STRICT`` or ``LAX`` . Default: ``LAX`` . All levels support all backends.
+
+            - ``STRICT`` : Only basic syntax is supported, and execution performance is optimal. Can be used for MindIR
+              load and export.
+            - ``LAX`` : Compatible with all Python syntax as much as possible. However, execution performance may be
+              affected and not optimal. Cannot be used for MindIR load and export due to some syntax that may not be
+              able to be exported.
+
+        gpu_config (dict): Set the parameters specific to gpu hardware platform. It is not set by default.
+            Currently, only setting `conv_fprop_algo` and `conv_dgrad_algo` and `conv_wgrad_algo` and `conv_allow_tf32`
+            and `matmul_allow_tf32` are supported on GPU hardware platform.
+
+            - conv_fprop_algo (str): Specifies convolution forward algorithm and the default value is 'normal',
+              The value range is as follows:
+
+              - normal: Use the heuristic search algorithm.
+              - performance: Use the trial search algorithm.
+              - implicit_gemm: This algorithm expresses the convolution as a matrix product without actually explicitly
+                forming the matrix that holds the input tensor data.
+              - implicit_precomp_gemm: This algorithm expresses convolution as a matrix product without actually
+                explicitly forming the matrix that holds the input tensor data, but still needs some memory workspace to
+                precompute some indices in order to facilitate the implicit construction of the matrix that holds the
+                input tensor data.
+              - gemm: This algorithm expresses the convolution as an explicit matrix product. A significant memory
+                workspace is needed to store the matrix that holds the input tensor data.
+              - direct: This algorithm expresses the convolution as a direct convolution (for example, without
+                implicitly or explicitly doing a matrix multiplication).
+              - fft: This algorithm uses the Fast-Fourier Transform approach to compute the convolution. A significant
+                memory workspace is needed to store intermediate results.
+              - fft_tiling: This algorithm uses the Fast-Fourier Transform approach but splits the inputs into tiles.
+                A significant memory workspace is needed to store intermediate results but less than fft algorithm for
+                large size images.
+              - winograd: This algorithm uses the Winograd Transform approach to compute the convolution. A reasonably
+                sized workspace is needed to store intermediate results.
+              - winograd_nonfused: This algorithm uses the Winograd Transform approach to compute the convolution. A
+                significant workspace may be needed to store intermediate results.
+            - conv_dgrad_algo (str): Specifies convolution data grad algorithm and the default value is 'normal',
+              The value range is as follows:
+
+              - normal: Use the heuristic search algorithm.
+              - performance: Use the trial search algorithm.
+              - algo_0: This algorithm expresses the convolution as a sum of matrix products without actually explicitly
+                forming the matrix that holds the input tensor data. The sum is done using the atomic add operation,
+                thus the results are non-deterministic.
+              - algo_1: This algorithm expresses the convolution as a matrix product without actually explicitly forming
+                the matrix that holds the input tensor data. The results are deterministic.
+              - fft: This algorithm uses a Fast-Fourier Transform approach to compute the convolution. A significant
+                memory workspace is needed to store intermediate results. The results are deterministic.
+              - fft_tiling: This algorithm uses the Fast-Fourier Transform approach but splits the inputs into tiles.
+                A significant memory workspace is needed to store intermediate results but less than fft for large size
+                images. The results are deterministic.
+              - winograd: This algorithm uses the Winograd Transform approach to compute the convolution. A reasonably
+                sized workspace is needed to store intermediate results. The results are deterministic.
+              - winograd_nonfused: This algorithm uses the Winograd Transform approach to compute the convolution.
+                A significant workspace may be needed to store intermediate results. The results are deterministic.
+            - conv_wgrad_algo (str): Specifies convolution filter grad algorithm and the default value is 'normal',
+              The value range is as follows:
+
+              - normal: Use the heuristic search algorithm.
+              - performance: Use the trial search algorithm.
+              - algo_0: This algorithm expresses the convolution as a sum of matrix products without actually explicitly
+                forming the matrix that holds the input tensor data. The sum is done using the atomic add operation,
+                thus the results are non-deterministic.
+              - algo_1: This algorithm expresses the convolution as a matrix product without actually explicitly forming
+                the matrix that holds the input tensor data. The results are deterministic.
+              - fft: This algorithm uses a Fast-Fourier Transform approach to compute the convolution. A significant
+                memory workspace is needed to store intermediate results. The results are deterministic.
+              - algo_3: This algorithm is similar to algo_0 but uses some small workspace to precompute some indices.
+                The results are also non-deterministic.
+              - winograd_nonfused: This algorithm uses the Winograd Transform approach to compute the convolution.
+                A significant workspace may be needed to store intermediate results. The results are deterministic.
+              - fft_tiling: This algorithm uses the Fast-Fourier Transform approach but splits the inputs into tiles.
+                A significant memory workspace is needed to store intermediate results but less than fft for large size
+                images. The results are deterministic.
+            - conv_allow_tf32 (bool): The flag below controls to allow Tensor core TF32 computation on CUDNN and the
+              default value is ``True``.
+            - matmul_allow_tf32 (bool): The flag below controls to allow Tensor core TF32 computation on CUBLAS and the
+              default value is ``False``.
 
     Raises:
         ValueError: If input key is not an attribute in context.
@@ -1129,13 +1427,14 @@ def set_context(**kwargs):
         >>> ms.set_context(graph_kernel_flags="--opt_level=2 --dump_as_text")
         >>> ms.set_context(reserve_class_name_in_scope=True)
         >>> ms.set_context(variable_memory_max_size="6GB")
+        >>> ms.set_context(aoe_tune_mode="online")
+        >>> ms.set_context(aoe_config={"job_type": "2"})
         >>> ms.set_context(check_bprop=True)
         >>> ms.set_context(max_device_memory="3.5GB")
         >>> ms.set_context(mempool_block_size="1GB")
         >>> ms.set_context(print_file_path="print.pb")
         >>> ms.set_context(max_call_depth=80)
         >>> ms.set_context(env_config_path="./env_config.json")
-        >>> ms.set_context(auto_tune_mode="GA,RL")
         >>> ms.set_context(grad_for_scalar=True)
         >>> ms.set_context(enable_compile_cache=True, compile_cache_path="./cache.ms")
         >>> ms.set_context(pynative_synchronize=True)
@@ -1145,7 +1444,11 @@ def set_context(**kwargs):
         >>> 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})
+        >>> ms.set_context(ascend_config={"precision_mode": "force_fp16", "jit_compile": True,
+        ...                "atomic_clean_policy": 1, "op_precision_mode": "./op_precision_config_file"})
+        >>> ms.set_context(jit_syntax_level=ms.STRICT)
+        >>> ms.set_context(gpu_config={"conv_fprop_algo": "performance", "conv_allow_tf32": True,
+        ...                "matmul_allow_tf32": True})
     """
     ctx = _context()
     # set device target first
@@ -1153,7 +1456,7 @@ def set_context(**kwargs):
         ctx.set_device_target(kwargs['device_target'])
     device = ctx.get_param(ms_ctx_param.device_target)
     for key, value in kwargs.items():
-        if key == 'enable_sparse':
+        if key in ('enable_sparse', 'auto_tune_mode'):
             logger.warning(f"For 'context.set_context', '{key}' parameter is deprecated, "
                            "and will be removed in the next version.")
             continue
@@ -1161,7 +1464,8 @@ 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'):
+        if key in ('precision_mode', 'jit_compile', 'atomic_clean_policy', 'matmul_allow_hf32', 'conv_allow_hf32',
+                   'op_precision_mode'):
             raise ValueError(f"Please set '{key}' through parameter ascend_config")
         if key == 'save_graphs':
             if value is True:
@@ -1170,6 +1474,9 @@ def set_context(**kwargs):
                 value = 0
             if value > 3:
                 raise ValueError(f"value for save_graphs should be 0-3 but got '{value}'")
+        if key == 'jit_syntax_level' and value not in (STRICT, COMPATIBLE, LAX):
+            raise ValueError(f"For 'jit_syntax_level', the value should be context.STRICT"
+                             f" or context.LAX, but got {value}.")
         if not _check_target_specific_cfgs(device, key):
             continue
         if hasattr(ctx, key):
@@ -1231,16 +1538,16 @@ class ParallelMode:
     """
     Parallel mode options.
 
-    There are five kinds of parallel modes, "STAND_ALONE", "DATA_PARALLEL",
-    "HYBRID_PARALLEL", "SEMI_AUTO_PARALLEL" and "AUTO_PARALLEL". Default: "STAND_ALONE".
+    There are five kinds of parallel modes, ``STAND_ALONE``, ``DATA_PARALLEL``,
+    ``HYBRID_PARALLEL``, ``SEMI_AUTO_PARALLEL`` and ``AUTO_PARALLEL``. Default: ``STAND_ALONE``.
 
-    - STAND_ALONE: Only one processor is working.
-    - DATA_PARALLEL: Distributes the data across different processors.
-    - HYBRID_PARALLEL: Achieves data parallelism and model parallelism manually.
-    - SEMI_AUTO_PARALLEL: Achieves data parallelism and model parallelism by setting parallel strategies.
-    - AUTO_PARALLEL: Achieves parallelism automatically.
+    - ``STAND_ALONE``: Only one processor is working.
+    - ``DATA_PARALLEL``: Distributes the data across different processors.
+    - ``HYBRID_PARALLEL``: Achieves data parallelism and model parallelism manually.
+    - ``SEMI_AUTO_PARALLEL``: Achieves data parallelism and model parallelism by setting parallel strategies.
+    - ``AUTO_PARALLEL``: Achieves parallelism automatically.
 
-    MODE_LIST: The list of all supported parallel modes.
+    ``MODE_LIST``: The list of all supported parallel modes.
     """
 
     STAND_ALONE = "stand_alone"
@@ -1261,32 +1568,26 @@ def set_ps_context(**kwargs):
         Some other environment variables should also be set for parameter server training mode.
         These environment variables are listed below:
 
-        MS_SERVER_NUM: Server number
-
-        MS_WORKER_NUM: Worker number
-
-        MS_SCHED_HOST: Scheduler IP address
-
-        MS_SCHED_PORT: Scheduler port
+        - MS_SERVER_NUM: Server number
+        - MS_WORKER_NUM: Worker number
+        - MS_SCHED_HOST: Scheduler IP address
+        - MS_SCHED_PORT: Scheduler port
+        - MS_ROLE: The role of this process:
 
-        MS_ROLE: The role of this process:
-
-        MS_SCHED: represents the scheduler,
-
-        MS_WORKER: represents the worker,
-
-        MS_PSERVER/MS_SERVER: represents the Server
+          - MS_SCHED: represents the scheduler,
+          - MS_WORKER: represents the worker,
+          - MS_PSERVER/MS_SERVER: represents the Server
 
     Args:
         enable_ps (bool): Whether to enable parameter server training mode.
                           Only after enable_ps is set True, the environment variables will be effective.
-                          Default: False.
+                          Default: ``False`` .
         config_file_path (string): Configuration file path used by recovery, parameter server training mode only
-                                   supports Server disaster recovery currently. Default: ''.
-        scheduler_manage_port (int): Scheduler manage port used to scale out/in. Default: 11202.
-        enable_ssl (bool): Set PS SSL mode enabled or disabled. Default: False.
-        client_password (str): Password to decrypt the secret key stored in the client certificate. Default: ''.
-        server_password (str): Password to decrypt the secret key stored in the server certificate. Default: ''.
+                                   supports Server disaster recovery currently. Default: ``''`` .
+        scheduler_manage_port (int): Scheduler manage port used to scale out/in. Default: ``11202`` .
+        enable_ssl (bool): Set PS SSL mode enabled or disabled. Default: ``False`` .
+        client_password (str): Password to decrypt the secret key stored in the client certificate. Default: ``''`` .
+        server_password (str): Password to decrypt the secret key stored in the server certificate. Default: ``''`` .
 
     Raises:
         ValueError: If input key is not the attribute in parameter server training mode context.
@@ -1305,13 +1606,15 @@ def get_ps_context(attr_key):
     Args:
         attr_key (str): The key of the attribute:
 
-            - enable_ps (bool): Whether to enable parameter server training mode.
+            - enable_ps (bool): Whether to enable parameter server training mode. Default: ``False`` .
             - config_file_path (string): Configuration file path used by recovery, parameter server training mode only
-              supports Server disaster recovery currently. Default: ''.
-            - scheduler_manage_port (int): Scheduler manage port used to scale out/in. Default: 11202.
-            - enable_ssl (bool): Set PS SSL mode enabled or disabled. Default: False.
-            - client_password (str): Password to decrypt the secret key stored in the client certificate. Default: ''.
-            - server_password (str): Password to decrypt the secret key stored in the server certificate. Default: ''.
+              supports Server disaster recovery currently. Default: ``''`` .
+            - scheduler_manage_port (int): Scheduler manage port used to scale out/in. Default: ``11202`` .
+            - enable_ssl (bool): Set PS SSL mode enabled or disabled. Default: ``False`` .
+            - client_password (str): Password to decrypt the secret key stored in the client certificate.
+              Default: ``''`` .
+            - server_password (str): Password to decrypt the secret key stored in the server certificate.
+              Default: ``''`` .
 
     Returns:
         Returns attribute value according to the key.
@@ -1333,6 +1636,10 @@ def reset_ps_context():
     - enable_ps: False.
 
     Meaning of each field and its default value refer to :func:`mindspore.set_ps_context`.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> ms.reset_ps_context()
     """
     _reset_ps_context()