mindspore 2.3.0__cp310-cp310-win_amd64.whl

mindspore/context.py

@@ -0,0 +1,1976 @@
+# Copyright 2020-2024 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.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""
+The context of mindspore, used to configure the current execution environment,
+includes the execution mode, execution backend and other feature switches.
+"""
+from __future__ import absolute_import
+
+import json
+import os
+import time
+import threading
+from collections import namedtuple
+from types import FunctionType
+
+from mindspore import log as logger
+from mindspore._c_expression import MSContext, ms_ctx_param
+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.parallel._offload_context import _set_offload_context, _get_offload_context
+from mindspore.hal.device import is_initialized
+
+__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'
+K_CONTEXT = None
+
+# Enumerate for the property 'jit_syntax_level'.
+STRICT = 0
+COMPATIBLE = 1
+LAX = 2
+
+# Enumerate for the property 'debug_level'.
+RELEASE = 0
+DEBUG = 1
+
+
+def _make_directory(path):
+    """Make directory."""
+    if path is None or not isinstance(path, str) or path.strip() == "":
+        raise ValueError(f"For 'context.set_context', the 'save_graphs_path' or the 'print_file_path' is invalid "
+                         f"type, it should be Non-empty string, but got '{path}'.")
+
+    path = os.path.realpath(path)
+    logger.debug("The absolute path is %r", path)
+
+    if not os.path.exists(path):
+        logger.debug("The directory(%s) doesn't exist, will create it", path)
+        try:
+            os.makedirs(path)
+        except FileExistsError:
+            logger.debug("The directory(%s) already exist.", path)
+        except PermissionError as e:
+            logger.critical(f"No write permission on the directory '{path}'', error = {e}")
+            raise ValueError(e.__str__() + f"\nNo write permission on the directory '{path}'.")
+    return path
+
+
+def _get_print_file_name(file_name):
+    """Add timestamp suffix to file name. Rename the file name:  file_name + "." + time(seconds)."""
+    time_second = str(int(time.time()))
+    file_name = file_name + "." + time_second
+    if os.path.exists(file_name):
+        raise ValueError("For 'context.set_context', the argument 'print_file_path' {} already exists, "
+                         "please check it".format(file_name))
+    return file_name
+
+
+class _ThreadLocalInfo(threading.local):
+    """
+    Thread local Info used for store thread local attributes.
+    """
+
+    def __init__(self):
+        super(_ThreadLocalInfo, self).__init__()
+        self._reserve_class_name_in_scope = True
+        self.debug_runtime = False
+
+    @property
+    def reserve_class_name_in_scope(self):
+        """Get whether to save the network class name in the scope."""
+        return self._reserve_class_name_in_scope
+
+    @reserve_class_name_in_scope.setter
+    def reserve_class_name_in_scope(self, reserve_class_name_in_scope):
+        """Set whether to save the network class name in the scope."""
+        self._reserve_class_name_in_scope = reserve_class_name_in_scope
+
+
+_ContextRecord = namedtuple(
+    "_ContextRecord", ["is_pynative_mode", "switch_context_fn"])
+
+
+class _ContextSwitchInfo(threading.local):
+    """
+    Record of context switch information.
+
+    Args:
+        is_pynative (bool): Whether to adopt the PyNative mode.
+    """
+
+    def __init__(self, is_pynative):
+        super(_ContextSwitchInfo, self).__init__()
+        self.context_stack = []
+        if is_pynative:
+            self.push(True, None)
+
+    def push(self, is_pynative, switch_context_fn):
+        """
+        Push a context switch record onto the stack.
+
+        Args:
+            is_pynative (bool): Whether context switch to PyNative mode.
+            switch_context_fn (Function): A callable that executes the context switch.
+        """
+        if isinstance(switch_context_fn, FunctionType):
+            switch_context_fn()
+        self.context_stack.append(
+            _ContextRecord(is_pynative, switch_context_fn))
+
+    def pop(self):
+        self.context_stack.pop()
+
+
+class _Context:
+    """
+    _Context is the environment in which operations are executed
+
+    Note:
+        Create a context through instantiating Context object is not recommended.
+        should use context() to get the context since Context is a singleton.
+    """
+    _instance = None
+    _instance_lock = threading.Lock()
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance_lock.acquire()
+            cls._instance = object.__new__(cls)
+            cls._instance_lock.release()
+        return cls._instance
+
+    def __init__(self):
+        self._thread_local_info = _ThreadLocalInfo()
+        self._context_switches = _ContextSwitchInfo(False)
+        self._context_handle = MSContext.get_instance()
+        self._support_binary = False
+        self.enable_compile_cache = None
+        self._mode = PYNATIVE_MODE
+        self._jit_config = {}
+
+    def __getattribute__(self, attr):
+        value = object.__getattribute__(self, attr)
+        if attr == "_context_handle" and value is None:
+            raise ValueError("Get {} failed, please check whether 'env_config_path' is correct.".format(attr))
+        return value
+
+    def get_param(self, param):
+        return self._context_handle.get_param(param)
+
+    def set_param(self, param, value):
+        self._context_handle.set_param(param, value)
+
+    def get_mode(self):
+        """Get current mode."""
+        return self._mode
+
+    def get_jit_config(self):
+        """Get current jit_config."""
+        return self._jit_config
+
+    def set_mode(self, mode):
+        """
+        Switch between Graph mode and PyNative mode.
+
+        Args:
+            mode (int): GRAPH_MODE or PYNATIVE_MODE.
+        """
+        if mode == PYNATIVE_MODE:
+            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, 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:
+                self.set_backend_policy("ge")
+            self._context_switches.push(False, None)
+        else:
+            raise ValueError(f"For 'context.set_context', the argument 'mode' should be context.GRAPH_MODE (0) "
+                             f"or context.PYNATIVE_MODE (1), but got {mode}.")
+        self.set_param(ms_ctx_param.mode, mode)
+        self._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_debug_level(self, level):
+        """"Set the debug level for graph compiling"""
+        if level != RELEASE and level != DEBUG:
+            raise ValueError(f"For 'context.set_debug_level', the argument 'level' should be context.RELEASE "
+                             f"or context.DEBUG, but got {level}.")
+        self.set_param(ms_ctx_param.debug_level, level)
+
+    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 (str): "force_fp16", "allow_fp32_to_fp16", "allow_mix_precision",
+                            "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): precision mode config file path.
+                - op_debug_option (str): Enable debugging options for Ascend operators,
+                  default not enabled, only supports ``"oom"`` currently.
+                  ``"oom"``: Detect memory out of bounds.
+                - ge_options (dict): Global or session CANN options.
+                - exception_dump (str): Enable exception dump for Ascend operators. ``"0"`` , ``"1"`` and ``"2"``.
+                  Default: ``"2"`` .
+                - 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.
+                - host_scheduling_max_threshold(int): The host scheduling max threshold.
+        """
+        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],
+            'exception_dump': ["0", "1", "2"],
+            'op_precision_mode': (str,),
+            'ge_options': (dict,),
+            'parallel_speed_up_json_path': (str, None),
+            'host_scheduling_max_threshold': (int,),
+            'cur_step_num': (int,),
+            'save_checkpoint_steps': (int,),
+            'need_ckpt': (bool,),
+            'last_triggered_step': (int,),
+            'topo_order': (dict,),
+            'op_debug_option': (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"),
+            'exception_dump': self._get_ascend_config_setter('exception_dump'),
+            'op_debug_option': self._set_op_debug_option,
+            'op_precision_mode': self._set_op_precision_mode,
+            'ge_options': self._set_ge_options,
+            'parallel_speed_up_json_path': self._set_speedup_config_path,
+            'host_scheduling_max_threshold': self._get_ascend_config_setter('host_scheduling_max_threshold', str),
+            'cur_step_num': self._set_cur_step_num,
+            'save_checkpoint_steps': self._set_save_checkpoint_steps,
+            'need_ckpt': self._set_need_ckpt,
+            'last_triggered_step': self._set_last_triggered_step,
+            'topo_order': self._set_topo_order
+        }
+        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_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_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_jit_config(self, jit_config):
+        """
+        Enable jit config.
+
+        Args:
+            jit_config (dict):
+
+                - jit_level (str): "O0", "O1" or "O2" to control the compilation optimization level.
+        """
+        jit_cfgs = {'jit_level': ["O0", "O1", "O2"], 'infer_boost': ["on", "off"]}
+        key_args_map = {'jit_level': ms_ctx_param.jit_level, 'infer_boost': ms_ctx_param.infer_boost}
+        for jit_key in jit_config:
+            if jit_key not in jit_cfgs:
+                raise ValueError(f"For 'context.set_context', the key of argument 'jit_config' must be one of "
+                                 f"{jit_cfgs}, but got {jit_key}.")
+            supported_value = jit_cfgs.get(jit_key)
+            if jit_config[jit_key] not in supported_value:
+                raise ValueError(f"For 'jit_cfgs', the value of argument {jit_key} must be one of "
+                                 f"{supported_value}, but got {jit_config[jit_key]}.")
+            self._jit_config = jit_config
+            self.set_param(key_args_map[jit_key], jit_config[jit_key])
+
+        if 'infer_boost' in jit_config and jit_config['infer_boost'] == "on" and jit_config['jit_level'] != "O0":
+            raise ValueError(f"Only jit_level set O0 can set infer_boost to on.")
+
+    def set_backend_policy(self, policy):
+        success = self._context_handle.set_backend_policy(policy)
+        if not success:
+            raise RuntimeError("Backend policy must be one of values in ['ge', 'vm', 'ms']. "
+                               "But got {}.".format(policy))
+
+    def set_save_graphs_path(self, save_graphs_path):
+        self.set_param(ms_ctx_param.save_graphs_path, _make_directory(save_graphs_path))
+
+    def set_device_target(self, target):
+        """
+        The target device to run, support "Ascend", "GPU", and "CPU".
+
+        Args:
+            target (str): "Ascend", "GPU", and "CPU".
+        """
+        valid_targets = ["CPU", "GPU", "Ascend", "Davinci"]
+        if target not in valid_targets:
+            raise ValueError(f"For 'context.set_context', the argument 'device_target' must be one of "
+                             f"{valid_targets}, but got {target}.")
+        if target == "Davinci":
+            target = "Ascend"
+            logger.warning("The device 'Davinci' is deprecated and will be removed in the next version. "
+                           "For 'context.set_context', please set the argument 'device_target' "
+                           "to 'CPU', 'GPU' or 'Ascend',if you set it to 'Davinci', it will be automatically "
+                           "changed to 'Ascend'.")
+        # If in Parameter Server mode, Ascend card should not be used by server and scheduler.
+        if _need_reset_device_target_for_ps(target):
+            logger.info("Reset device target to CPU when set_device_target.")
+            target = "CPU"
+        self.set_param(ms_ctx_param.device_target, target)
+        if self.enable_debug_runtime and target == "CPU":
+            self.set_backend_policy("vm")
+
+    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.aoe_tune_mode, tune_mode)
+        else:
+            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:
+            raise ValueError(f"For 'context.set_context', the argument 'device_id' must be in range [0, 4095], "
+                             f"but got {device_id}.")
+        self.set_param(ms_ctx_param.device_id, device_id)
+
+    def set_max_call_depth(self, max_call_depth):
+        if max_call_depth <= 0:
+            raise ValueError(f"For 'context.set_context', the argument 'max_call_depth' must be greater than 0, "
+                             f"but got {max_call_depth}.")
+        self.set_param(ms_ctx_param.max_call_depth, max_call_depth)
+
+    def set_profiling_options(self, option):
+        if not isinstance(option, str):
+            raise TypeError("For 'context.set_context', the argument 'profiling_option' must be string, "
+                            "but got {}.".format(type(option)))
+        self.set_param(ms_ctx_param.profiling_options, option)
+
+    def set_variable_memory_max_size(self, variable_memory_max_size):
+        """set values of variable_memory_max_size and graph_memory_max_size"""
+        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):
+            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 {}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 {}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"
+        self.set_param(ms_ctx_param.variable_memory_max_size, variable_memory_max_size_)
+        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):
+            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 {}."
+                             .format(max_device_memory))
+        max_device_memory_value = float(max_device_memory[:-2])
+        if max_device_memory_value == 0:
+            raise ValueError("For 'context.set_context', the argument 'max_device_memory' should not be \"0GB\".")
+        self.set_param(ms_ctx_param.max_device_memory, max_device_memory_value)
+
+    def set_mempool_block_size(self, mempool_block_size):
+        """Set the block size of memory pool."""
+        global_jit_config = get_jit_config()
+        is_force_kbk = False
+        if global_jit_config:
+            is_force_kbk = global_jit_config.get('jit_level') == "O0" or global_jit_config.get('jit_level') == "O1"
+        if _get_mode() == GRAPH_MODE and not is_force_kbk:
+            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 or set jit_level=O0/O1.")
+            return
+        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))
+        mempool_block_size_value = float(mempool_block_size[:-2])
+        if mempool_block_size_value < 1.0:
+            raise ValueError("For 'context.set_context',  the argument 'mempool_block_size' should be "
+                             "greater or equal to \"1GB\", "
+                             "but got {}GB".format(float(mempool_block_size[:-2])))
+        self.set_param(ms_ctx_param.mempool_block_size, mempool_block_size_value)
+
+    def set_print_file_path(self, file_path):
+        """Add timestamp suffix to file name. Sets print file path."""
+        print_file_path = os.path.realpath(file_path)
+        if os.path.isdir(print_file_path):
+            raise IOError("For 'context.set_context', the argument 'print_file_path' should be file path, "
+                          "but got directory {}.".format(file_path))
+
+        if os.path.exists(print_file_path):
+            _path, _file_name = os.path.split(print_file_path)
+            path = _make_directory(_path)
+            file_name = _get_print_file_name(_file_name)
+            full_file_name = os.path.join(path, file_name)
+        else:
+            full_file_name = print_file_path
+        self.set_param(ms_ctx_param.print_file_path, full_file_name)
+
+    def set_env_config_path(self, env_config_path):
+        """Check and set env_config_path."""
+        if not self._context_handle.enable_dump_ir():
+            raise ValueError("For 'context.set_context', the argument 'env_config_path' is not supported, please "
+                             "enable ENABLE_DUMP_IR with '-D on' and recompile source firstly.")
+        env_config_path = os.path.realpath(env_config_path)
+        if not os.path.isfile(env_config_path):
+            raise ValueError("For 'context.set_context', the 'env_config_path' file %r is not exists, "
+                             "please check whether 'env_config_path' is correct." % env_config_path)
+        try:
+            with open(env_config_path, 'r') as f:
+                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)) from exo
+        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 or equal to 0.")
+        self.set_param(ms_ctx_param.runtime_num_threads, runtime_num_threads)
+
+    def set_op_timeout(self, op_timeout):
+        """Set the maximum duration of executing an operator in seconds."""
+        if op_timeout < 0:
+            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,
+        'device_target': set_device_target,
+        'aoe_tune_mode': set_aoe_tune_mode,
+        'device_id': set_device_id,
+        'max_call_depth': set_max_call_depth,
+        'profiling_options': set_profiling_options,
+        'variable_memory_max_size': set_variable_memory_max_size,
+        'max_device_memory': set_max_device_memory,
+        '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,
+        '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,
+        'jit_syntax_level': set_jit_syntax_level,
+        'debug_level': set_debug_level,
+        'gpu_config': set_gpu_config,
+        'aoe_config': set_aoe_config,
+        'jit_config': set_jit_config,
+    }
+
+    @property
+    def reserve_class_name_in_scope(self):
+        """Get whether to save the network class name in the scope."""
+        return self._thread_local_info.reserve_class_name_in_scope
+
+    @reserve_class_name_in_scope.setter
+    def reserve_class_name_in_scope(self, reserve_class_name_in_scope):
+        """Set whether to save the network class name in the scope."""
+        if not isinstance(reserve_class_name_in_scope, bool):
+            raise ValueError("For 'context.set_context', the type of the property 'reserve_class_name_in_scope' must "
+                             "be bool, but got {}.".format(type(reserve_class_name_in_scope)))
+        self._thread_local_info.reserve_class_name_in_scope = reserve_class_name_in_scope
+
+    @property
+    def enable_ge(self):
+        return self._context_handle.get_backend_policy() == 'ge'
+
+    @property
+    def enable_debug_runtime(self):
+        return self._thread_local_info.debug_runtime
+
+    @enable_debug_runtime.setter
+    def enable_debug_runtime(self, enable):
+        thread_info = self._thread_local_info
+        thread_info.debug_runtime = enable
+
+    @property
+    def support_binary(self):
+        """Whether support run .pyc or .so in graph mode."""
+        return self._support_binary
+
+    @support_binary.setter
+    def support_binary(self, support: bool):
+        if not isinstance(support, bool):
+            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_debug_option(self, option_value):
+        valid_order = {'oom'}
+        if not isinstance(option_value, str):
+            raise TypeError(f"For 'ascend_config', the type of 'op_debug_option' must be str, "
+                            f"but got {type(option_value)}.")
+        if option_value not in valid_order:
+            raise ValueError(f"For 'ascend_config', the 'op_debug_option' supports being set to 'oom' currently, "
+                             f"but got {option_value}.")
+        self.set_param(ms_ctx_param.op_debug_option, option_value)
+
+    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_ge_options(self, ge_options):
+        """Set ge options."""
+        for level, options in ge_options.items():
+            if level not in ['global', 'session']:
+                raise ValueError(f"For 'ascend_config', the key of ge_options must be one of "
+                                 f"('global', 'session'), but got {level}.")
+
+            if not isinstance(options, dict):
+                raise TypeError(f"For 'ge_options', the type of {level} options must be dict, "
+                                f"but got {type(options)}. The error options: {options}.")
+
+            for key, value in options.items():
+                if not isinstance(key, str):
+                    raise TypeError(f"For 'ge_options', the type of key and value must be str, "
+                                    f"but got {type(key)}. The error key is {key}.")
+                if not isinstance(value, str):
+                    raise TypeError(f"For 'ge_options', the type of key and value must be str, "
+                                    f"but got {type(value)}. The error value is {value}")
+
+        options_str = json.dumps(ge_options)
+        self.set_param(ms_ctx_param.ge_options, options_str)
+
+    def _set_topo_order(self, topo_order):
+        """
+        Set topo order.
+
+        Args:
+            topo_order (dict):
+                key: str, the name of the graph.
+                value: str, the topo order of the graph, should be one of 'dfs', 'bfs', 'rdfs'.
+        """
+        valid_order = {'dfs', 'bfs', 'rdfs'}
+        if not isinstance(topo_order, dict):
+            raise TypeError(f"For 'ascend_config', the 'topo_order' should be a dict, "
+                            f"got '{type(topo_order)}'.")
+        for k, v in topo_order.items():
+            if not isinstance(k, str):
+                raise TypeError("key {} is not a str".format(k))
+            if v not in valid_order:
+                raise ValueError("value {} should be one of {}.".format(v, valid_order))
+
+        options_str = json.dumps(topo_order)
+        self.set_param(ms_ctx_param.topo_order, options_str)
+
+    def _set_need_ckpt(self, need_ckpt):
+        """Set need ckpt flag"""
+        if not isinstance(need_ckpt, bool):
+            raise TypeError(f"For step num, the value type should be int, but got {type(need_ckpt)}, {need_ckpt}")
+        self.set_param(ms_ctx_param.need_ckpt, need_ckpt)
+
+    def _set_cur_step_num(self, step_num):
+        """set current step num at every step begin"""
+        if not isinstance(step_num, int):
+            raise TypeError(f"For step num, the value type should be int, but got {type(step_num)}, {step_num}")
+        self.set_param(ms_ctx_param.cur_step_num, step_num)
+
+    def _set_save_checkpoint_steps(self, steps):
+        """set save checkpoint steps before run"""
+        if not isinstance(steps, int):
+            raise TypeError(f"For step num, the value type should be int, but got {type(steps)}, {steps}")
+        self.set_param(ms_ctx_param.save_checkpoint_steps, steps)
+
+    def _set_last_triggered_step(self, step):
+        """set last triggered save ckpt steps before run"""
+        if not isinstance(step, int):
+            raise TypeError(f"For step num, the value type should be int, but got {type(step)}, {step}")
+        self.set_param(ms_ctx_param.last_triggered_step, step)
+
+    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, bool),
+                            "matmul_grad_comm_overlap": (ms_ctx_param.matmul_grad_comm_overlap, bool),
+                            "enable_task_opt": (ms_ctx_param.enable_task_opt, bool),
+                            "enable_grad_comm_opt": (ms_ctx_param.enable_grad_comm_opt, bool),
+                            "recompute_allgather_overlap_fagrad":
+                                (ms_ctx_param.recompute_allgather_overlap_fagrad, bool),
+                            "interleaved_matmul_comm": (ms_ctx_param.interleaved_matmul_comm, bool),
+                            "bias_add_comm_swap": (ms_ctx_param.bias_add_comm_swap, bool),
+                            "enable_opt_shard_comm_opt": (ms_ctx_param.enable_opt_shard_comm_opt, bool),
+                            "enable_begin_end_inline_opt": (ms_ctx_param.enable_begin_end_inline_opt, bool),
+                            "enable_concat_eliminate_opt": (ms_ctx_param.enable_concat_eliminate_opt, bool),
+                            "interleaved_layernorm_comm": (ms_ctx_param.interleaved_layernorm_comm, bool),
+                            "compute_communicate_fusion_level":
+                                (ms_ctx_param.compute_communicate_fusion_level, int),
+                            "enable_flash_attention_load_balance":
+                                (ms_ctx_param.enable_flash_attention_load_balance, bool)}
+            with open(speedup_config_real_path, 'r') as f:
+                speedup_config = json.load(f)
+                for key, value in speedup_config.items():
+                    if not isinstance(key, str):
+                        raise TypeError("key {} is not a str".format(key))
+                    if key not in valid_option:
+                        raise ValueError("key {} should be one of {}.".format(key, valid_option.keys()))
+                    set_func, valid_type = valid_option.get(key)
+                    if not isinstance(value, valid_type):
+                        raise TypeError(f"The value type of {key} must be {valid_type}, "
+                                        f"but got value is {value} and type is {type(value)}.")
+                    self.set_param(set_func, value)
+        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():
+    """
+    Get the global _context, if context is not created, create a new one.
+
+    Returns:
+        _Context, the global context in PyNative mode.
+    """
+    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
+        if default_backend == 'debug':
+            K_CONTEXT.enable_debug_runtime = True
+            default_backend = 'vm'
+            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, pipeline_segments=int,
+                 pipeline_result_broadcast=bool, parallel_optimizer_config=dict,
+                 pipeline_config=dict,
+                 comm_fusion=dict, strategy_ckpt_config=dict, force_fp32_communication=bool)
+def set_auto_parallel_context(**kwargs):
+    r"""
+    Set auto parallel context, only data parallel supported on CPU.
+
+    Note:
+        Attribute name is required for setting attributes.
+        If a program has tasks on different parallel modes, before setting a new parallel mode for the
+        next task, interface :func:`mindspore.reset_auto_parallel_context` should be called to reset
+        the configuration.
+        Setting or changing parallel modes must be called before creating any Initializer, otherwise,
+        it may have RuntimeError when compiling the network.
+
+    Some configurations are parallel mode specific, see the below table for details:
+
+    ===========================  ===========================
+    Common                       AUTO_PARALLEL
+    ===========================  ===========================
+    device_num                   gradient_fp32_sync
+    global_rank                  loss_repeated_mean
+    gradients_mean               search_mode
+    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
+    pipeline_config              auto_parallel_search_mode
+    force_fp32_communication     pipeline_result_broadcast
+               \                 comm_fusion
+               \                 strategy_ckpt_config
+               \                 group_ckpt_save_file
+               \                 auto_pipeline
+    ===========================  ===========================
+
+    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`` .
+        gradients_mean (bool): Whether to perform mean operator after allreduce of gradients.
+                     "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`` .
+        loss_repeated_mean (bool) - Indicates whether the mean operator is executed backwards when the
+                     calculation is repeated. 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.
+
+                     - data_parallel: Distributes the data across different processors.
+
+                     - hybrid_parallel: Achieves data parallelism and model parallelism manually.
+
+                     - 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"`` ,
+                     ``"sharding_propagation"`` and ``"dynamic_programming"`` (Not recommended).
+                     Default: ``"recursive_programming"`` .
+
+                     - 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.
+
+                     - sharding_propagation: Propagate shardings from configured ops to non-configured ops.
+
+                     - dynamic_programming: Dynamic programming search mode.
+        auto_parallel_search_mode (str): This is the old version of 'search_mode'. Here, remaining this attribute is
+                     for forward compatibility, and this attribute will be deleted in a future MindSpore version.
+        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`` .
+        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 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
+                       `Lamb` and `AdamWeightDecay` in Ascend . Default: ``False`` .
+        force_fp32_communication (bool): A switch that determines whether reduce operators (AllReduce, ReduceScatter)
+                        are forced to use the fp32 data type for communication during communication. True is the enable
+                        switch. 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`` .
+        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.
+                        Default: ``1`` .
+        pipeline_result_broadcast (bool): A switch that broadcast the last stage result to all other stage in pipeline
+                        parallel inference. Default: ``False`` .
+        pipeline_config (dict): A dict contains the keys and values for setting the pipeline parallelism configuration.
+                        It supports the following keys:
+
+                        - pipeline_interleave(bool): Indicates whether to enable the interleaved execution mode.
+                        - pipeline_scheduler(str): Indicates the scheduling mode for pipeline parallelism. Only support
+                          ``gpipe/1f1b``.
+        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. 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
+                          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 ``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`` .
+
+                        - 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`` .
+
+                        - 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
+                          manually, and the fusion threshold must be larger than `0` MB. In `index` mode, it is same as
+                          `all_reduce_fusion_config`.
+
+                        - allgather: If communication fusion type is `allgather`. The `mode` contains: `auto`, `size`.
+                          In `auto` mode, AllGather fusion is configured by gradients size, and the default fusion
+                          threshold is `64` MB. In 'size' mode, AllGather fusion is configured by gradients size
+                          manually, and the fusion threshold must be larger than `0` MB.
+
+                        - 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`` .
+        group_ckpt_save_file (str): The path to save parallel group checkpoint.
+        auto_pipeline (bool): Set the pipeline stage number to automatic. Its value will be selected between 1 and the
+                        parameter `pipeline_stages`. This option requires the `parallel_mode` to be ``auto_parallel``
+                        and the `search_mode` to be ``recursive_programming``. Default: ``False`` .
+
+    Raises:
+        ValueError: If input key is not attribute in auto parallel context.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> ms.set_auto_parallel_context(device_num=8)
+        >>> ms.set_auto_parallel_context(global_rank=0)
+        >>> ms.set_auto_parallel_context(gradients_mean=True)
+        >>> ms.set_auto_parallel_context(gradient_fp32_sync=False)
+        >>> ms.set_auto_parallel_context(parallel_mode="auto_parallel")
+        >>> ms.set_auto_parallel_context(search_mode="recursive_programming")
+        >>> ms.set_auto_parallel_context(auto_parallel_search_mode="recursive_programming")
+        >>> ms.set_auto_parallel_context(parameter_broadcast=False)
+        >>> ms.set_auto_parallel_context(strategy_ckpt_load_file="./strategy_stage1.ckpt")
+        >>> ms.set_auto_parallel_context(strategy_ckpt_save_file="./strategy_stage1.ckpt")
+        >>> ms.set_auto_parallel_context(dataset_strategy=((1, 8), (1, 8)))
+        >>> ms.set_auto_parallel_context(enable_parallel_optimizer=False)
+        >>> 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)
+        >>> ms.set_auto_parallel_context(pipeline_stages=2, pipeline_result_broadcast=True)
+        >>> 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)
+        >>> 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)
+
+
+def get_auto_parallel_context(attr_key):
+    """
+    Get auto parallel context attribute value according to the key.
+
+    Args:
+        attr_key (str): The key of the attribute.
+
+    Returns:
+        Returns attribute value according to the key.
+
+    Raises:
+        ValueError: If input key is not attribute in auto parallel context.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> parallel_mode = ms.get_auto_parallel_context("parallel_mode")
+        >>> dataset_strategy = ms.get_auto_parallel_context("dataset_strategy")
+    """
+    return _get_auto_parallel_context(attr_key)
+
+
+def reset_auto_parallel_context():
+    """
+    Reset auto parallel context attributes to the default values.
+
+    - device_num: 1.
+    - global_rank: 0.
+    - gradients_mean: False.
+    - gradient_fp32_sync: True.
+    - parallel_mode: 'stand_alone'.
+    - search_mode: 'recursive_programming'.
+    - auto_parallel_search_mode: 'recursive_programming'.
+    - parameter_broadcast: False.
+    - strategy_ckpt_load_file: ''.
+    - strategy_ckpt_save_file: ''.
+    - full_batch: False.
+    - enable_parallel_optimizer: False.
+    - force_fp32_communication: False
+    - enable_alltoall: False.
+    - pipeline_stages: 1.
+    - pipeline_result_broadcast: False.
+    - fusion_threshold: 64.
+    - auto_pipeline: False.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> ms.reset_auto_parallel_context()
+    """
+    _reset_auto_parallel_context()
+
+
+@args_type_check(offload_config=dict)
+def set_offload_context(offload_config):
+    r"""
+    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.
+
+            - 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"})
+    """
+    _set_offload_context(offload_config)
+
+
+def 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()
+    """
+    return _get_offload_context()
+
+
+def _check_target_specific_cfgs(device, arg_key):
+    """Checking whether a config is suitable for a specified device"""
+    device_cfgs = {
+        'enable_graph_kernel': ['Ascend', 'GPU', 'CPU'],
+        'graph_kernel_flags': ['Ascend', 'GPU', 'CPU'],
+        'enable_reduce_precision': ['Ascend'],
+        'print_file_path': ['Ascend'],
+        'variable_memory_max_size': ['Ascend'],
+        'max_device_memory': ['Ascend', 'GPU'],
+        'mempool_block_size': ['GPU', 'Ascend'],
+        'disable_format_transform': ['GPU'],
+        '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:
+        return True
+    supported_devices = device_cfgs[arg_key]
+    if device in supported_devices:
+        return True
+    logger.warning(f"For 'context.set_context', when set the argument '{arg_key}', "
+                   f"the argument 'device_target' only supports devices in '{supported_devices}', "
+                   f"but got '{device}', ignore it.")
+    return False
+
+
+def _check_ascend_device_context_initialized(device_target, settings):
+    if device_target == 'Ascend' and is_initialized(device_target):
+        for key, _ in settings.items():
+            if key in ('ascend_config', 'deterministic', 'jit_compile', 'exception_dump', 'device_id'):
+                logger.warning(f"For 'context.set_context' in Ascend backend, the backend is already initialized, "
+                               "please set it before the definition of any Tensor and Parameter, and the "
+                               "instantiation and execution of any operation and net, otherwise the settings may not "
+                               "take effect. ")
+                break
+
+
+def _check_key(key):
+    if key in ('precision_mode', 'jit_compile', 'atomic_clean_policy', 'matmul_allow_hf32', 'conv_allow_hf32',
+               'op_precision_mode', 'host_scheduling_max_threshold', 'ge_options', 'op_debug_option'):
+        raise ValueError(f"Please set '{key}' through parameter ascend_config")
+
+
+@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, 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, jit_syntax_level=int, debug_level=int,
+                 jit_enable_inplace_ops=bool, gpu_config=dict, jit_config=dict, enable_compile_cache=bool)
+def set_context(**kwargs):
+    """
+    Set context for running environment.
+
+    Context should be configured before running your program. If there is no configuration,
+    it will be automatically set according to the device target by default.
+
+    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`` .
+
+    Some configurations are device specific, see the below table for details:
+
+    +-------------------------+------------------------------+----------------------------+
+    | Function Classification |   Configuration Parameters   |   Hardware Platform Support|
+    +=========================+==============================+============================+
+    | System Configuration    |   device_id                  |   CPU/GPU/Ascend           |
+    |                         +------------------------------+----------------------------+
+    |                         |   device_target              |   CPU/GPU/Ascend           |
+    |                         +------------------------------+----------------------------+
+    |                         |  max_device_memory           |  GPU/Ascend                |
+    |                         +------------------------------+----------------------------+
+    |                         |  variable_memory_max_size    |  Ascend                    |
+    |                         +------------------------------+----------------------------+
+    |                         |  mempool_block_size          |  GPU/Ascend                |
+    |                         +------------------------------+----------------------------+
+    |                         |  op_timeout                  |  Ascend                    |
+    +-------------------------+------------------------------+----------------------------+
+    | Debug Configuration     |  save_graphs                 |  CPU/GPU/Ascend            |
+    |                         +------------------------------+----------------------------+
+    |                         |  save_graphs_path            |  CPU/GPU/Ascend            |
+    |                         +------------------------------+----------------------------+
+    |                         |  enable_dump                 |  Ascend                    |
+    |                         +------------------------------+----------------------------+
+    |                         |  save_dump_path              |  Ascend                    |
+    |                         +------------------------------+----------------------------+
+    |                         |  deterministic               |  Ascend                    |
+    |                         +------------------------------+----------------------------+
+    |                         |  print_file_path             |  Ascend                    |
+    |                         +------------------------------+----------------------------+
+    |                         |  env_config_path             |  CPU/GPU/Ascend            |
+    |                         +------------------------------+----------------------------+
+    |                         |  precompile_only             |  CPU/GPU/Ascend            |
+    |                         +------------------------------+----------------------------+
+    |                         |  reserve_class_name_in_scope |  CPU/GPU/Ascend            |
+    |                         +------------------------------+----------------------------+
+    |                         |  pynative_synchronize        |  CPU/GPU/Ascend            |
+    |                         +------------------------------+----------------------------+
+    |                         |  debug_level                 |  CPU/GPU/Ascend            |
+    +-------------------------+------------------------------+----------------------------+
+    | Executive Control       |   mode                       |   CPU/GPU/Ascend           |
+    |                         +------------------------------+----------------------------+
+    |                         |  enable_graph_kernel         |  Ascend/GPU                |
+    |                         +------------------------------+----------------------------+
+    |                         |  graph_kernel_flags          |  Ascend/GPU                |
+    |                         +------------------------------+----------------------------+
+    |                         |  enable_reduce_precision     |  Ascend                    |
+    |                         +------------------------------+----------------------------+
+    |                         |  aoe_tune_mode               |  Ascend                    |
+    |                         +------------------------------+----------------------------+
+    |                         |  aoe_config                  |  Ascend                    |
+    |                         +------------------------------+----------------------------+
+    |                         |  check_bprop                 |  CPU/GPU/Ascend            |
+    |                         +------------------------------+----------------------------+
+    |                         |  max_call_depth              |  CPU/GPU/Ascend            |
+    |                         +------------------------------+----------------------------+
+    |                         |  grad_for_scalar             |  CPU/GPU/Ascend            |
+    |                         +------------------------------+----------------------------+
+    |                         |  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            |
+    |                         +------------------------------+----------------------------+
+    |                         |  disable_format_transform    |  GPU                       |
+    |                         +------------------------------+----------------------------+
+    |                         |  support_binary              |  CPU/GPU/Ascend            |
+    |                         +------------------------------+----------------------------+
+    |                         |  memory_optimize_level       |  CPU/GPU/Ascend            |
+    |                         +------------------------------+----------------------------+
+    |                         |  memory_offload              |  GPU/Ascend                |
+    |                         +------------------------------+----------------------------+
+    |                         |  ascend_config               |  Ascend                    |
+    |                         +------------------------------+----------------------------+
+    |                         |  jit_syntax_level            |  CPU/GPU/Ascend            |
+    |                         +------------------------------+----------------------------+
+    |                         |  gpu_config                  |  GPU                       |
+    |                         +------------------------------+----------------------------+
+    |                         |  jit_config                  |  CPU/GPU/Ascend            |
+    +-------------------------+------------------------------+----------------------------+
+
+    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`` .
+        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' 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 or jit level is 'O0'/'O1'
+            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.
+        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. The defaults for AI Core and AICPU operators vary on different hardware.
+            For more information,
+            please refer to `Ascend Community document about aclrtSetOpExecuteTimeOut
+            <https://www.hiascend.com/document/detail/en/CANNCommunityEdition/600alphaX/infacldevg/aclcppdevg/aclcppdevg_03_0069.html>`_.
+            Default: ``900`` .
+        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 compilation.
+            - True or 2: Generate more ir files related to backend process.
+            - 3: Generate visualization computing graphs and detailed frontend ir graphs.
+
+            When the network structure is complex, setting `save_graphs` attribute to ``2`` or ``3`` may take too long.
+            If you need quick problem locating, you can switch to ``1`` first.
+
+            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
+            a file by default, and print_file_path is not set, the screen will be displayed.
+            If the saved file already exists, the timestamp suffix will be added to the file. Saving data to a file
+            solves the problem of data loss in screen printing when a large amount of data is generated.
+            If it is not set, an error will be reported: prompt to set the upper absolute path.
+            When print data to file, the total output bytes of single print must be less then 2GB(limited by
+            protobuf).
+        env_config_path (str): Config path for DFX.
+            Through mindspore.set_context(env_config_path="./mindspore_config.json")
+
+            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`` .
+            - 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.
+
+        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.
+            For example:
+
+            Default/net-Net1/net-Net2 (reserve_class_name_in_scope=True)
+
+            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.
+        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/master/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,
+
+            .. 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.
+              Available values are:
+
+              - 0: disables graph kernel fusion;
+              - 1: enables the basic fusion of operators;
+              - 2: includes all optimizations of level 1,
+                and turns on more optimizations such as CSE, arithmetic simplification and so on;
+              - 3: includes all optimizations of level 2, and turns on more optimizations such as SitchingFusion,
+                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`` .
+
+        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`` .
+        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`` .
+            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,
+            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
+            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,
+            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 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`` ,
+            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, `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
+            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.
+            On Ascend hardware platform, default: ``O1``, on other hardware platforms, default: ``O0``.
+            The value must be in ['O0', 'O1'].
+
+            - O0: priority performance option, disable SOMAS (Safe Optimized Memory Allocation Solver)
+              and some other memory optimizations.
+            - O1: priority memory option, enable SOMAS and some other memory optimizations.
+        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 graph compilation level is not 'O0'; 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.
+            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, 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: For cube operators, use the float16. For vector operators,
+                prefer to keep the origin dtype, if the operator in model can support float32,
+                it will keep original dtype, otherwise it will reduce to 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.
+              - allow_fp32_to_bf16: For cube operators, use the bfloat16. For vector operators,
+                prefer to keep the origin dtype, if the operator in model can support float32,
+                it will keep original dtype, otherwise it will reduce to 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. When set to 'True', online compilation is
+              prioritized. When set to 'False', compiled operator binary files are prioritized to improve compilation
+              performance. The default settings are online compilation for static shape, and compiled operator binary
+              files for dynamic shape.
+            - 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/>`_ .
+            - exception_dump (str): Enable exception dump for Ascend operators, providing the input and output data for
+              failing Ascend operators. The value can be ``"0"`` , ``"1"`` and ``"2"``. For ``"0"`` , exception dump is
+              turned off; for ``"1"``, all inputs and outputs will be dumped for AICore exception operators;
+              for ``"2"``, inputs will be dumped for AICore exception operators, reducing the saved information
+              but improving performance. Default: ``"2"`` .
+            - op_precision_mode (str): Path to config file of op precision mode. For detailed information, please refer
+              to `Ascend community <https://www.hiascend.com/>`_ .
+            - op_debug_option (str): Enable debugging options for Ascend operators, default not enabled.
+              The value currently only supports being set to ``"oom"``.
+
+              - ``"oom"``: When there is a memory out of bounds during the execution of an operator,
+                AscendCL will return an error code of ``EZ9999``.
+
+            - ge_options (dict): Set options for CANN. The options are divided into two categories: global and session.
+              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/document/detail/zh/canncommercial/70RC1/inferapplicationdev/graphdevg/atlasgeapi_07_0119.html>`_ .
+              The configuration options in `ge_options` may be duplicated with the options in `ascend_config`. If the
+              same configuration options are set in both `ascend_config` and `ge_options`, the one set in `ge_options`
+              shall prevail.
+
+              - global (dict): Set global options.
+              - session (dict): Set session options.
+
+            - 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/master/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 dw matmul and
+                tensor parallel communication ops if True. Default: False.
+              - recompute_allgather_overlap_fagrad (bool): Enable overlap between duplicated allgather by recomputing
+                in sequence parallel and flashattentionscoregrad ops if True. Default: False.
+              - enable_task_opt (bool): Enable communication fusion to optimize the number of communication operator
+                tasks if True.
+                Default: False.
+              - enable_grad_comm_opt (bool): Enable overlap between dx ops and data parallel communication ops if True.
+                Currently, do not support
+                `LazyInline <https://www.mindspore.cn/docs/en/master/api_python/mindspore/mindspore.lazy_inline.html>`
+                Default: False.
+              - enable_opt_shard_comm_opt (bool): Enable overlap between forward ops
+                and optimizer parallel allgather communication if True. Currently, do not support
+                `LazyInline <https://www.mindspore.cn/docs/en/master/api_python/mindspore/mindspore.lazy_inline.html>`
+                Default: False.
+              - compute_communicate_fusion_level (int): Enable the fusion between compute and communicate.
+                Default: ``0``.
+
+                - 0: Disable fusion.
+
+                - 1: Apply fusion to forward nodes.
+
+                - 2: Apply fusion to backward nodes.
+
+                - 3: Apply fusion to all nodes.
+              - bias_add_comm_swap (bool): Enable node execution order swap communication operators and add operators
+                if ``True``. Only 1-dimension bias node is supported. Default: ``False``.
+            - host_scheduling_max_threshold(int): The max threshold to control whether the dynamic shape process is
+              used when run the static graph, the default value is 0. When the number of operations in the static graph
+              is less than the max threshold, this graph will be executed in dynamic shape process. In large model
+              scenarios, this approach can save stream resources. If the number of operations in the static graph is
+              greater than the maximum threshold, this graph will be executed in original static process.
+
+        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.
+
+        debug_level (int): Set config for debugging. Default value: ``RELEASE``.
+
+            - ``RELEASE``: Used for normally running, and some debug information will be discard to get a better
+              compiling performance.
+            - ``DEBUG``: Used for debugging when errors occur, more information will be record in compiling process.
+
+        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``.
+
+        jit_config (dict): Set the global jit config for compile, take effect in network defined in Cell or jit
+            decorators. It is not set by default.
+            The setting in context is the global jit config, while JitConfig is the local network's jit config.
+            When both exist simultaneously, the global jit config will not overwrite the local network's jit config.
+
+            - jit_level (str): Used to control the compilation optimization level. Default: ``""`` , The framework
+              automatically selects the execution method based on product, Altas training product is O2, and all other
+              products are O0. The value range is as follows:
+
+              - ``"O0"``: Except for optimizations that may affect functionality, all other optimizations are turned
+                off, adopt KernelByKernel execution mode.
+              - ``"O1"``: Using commonly used optimizations and automatic operator fusion optimizations,
+                adopt KernelByKernel execution mode.
+              - ``"O2"``: Ultimate performance optimization, adopt Sink execution mode.
+
+            - infer_boost (str): Used to control the infer mode. Default: ``"off"`` . The value range is as follows:
+
+              - ``"on"``: Enable infer mode, get better infer performance.
+              - ``"off"``: Disable infer mode, use forward to infer, performance is not good.
+
+    Raises:
+        ValueError: If input key is not an attribute in context.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> ms.set_context(mode=ms.PYNATIVE_MODE)
+        >>> ms.set_context(precompile_only=True)
+        >>> ms.set_context(device_target="Ascend")
+        >>> ms.set_context(device_id=0)
+        >>> ms.set_context(save_graphs=True, save_graphs_path="./model.ms")
+        >>> ms.set_context(enable_reduce_precision=True)
+        >>> ms.set_context(enable_graph_kernel=True)
+        >>> 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(grad_for_scalar=True)
+        >>> 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,
+        ...                "atomic_clean_policy": 1, "op_precision_mode": "./op_precision_config_file",
+        ...                "op_debug_option": "oom",
+        ...                "ge_options": {"global": {"ge.opSelectImplmode": "high_precision"},
+        ...                               "session": {"ge.exec.atomicCleanPolicy": "0"}}})
+        >>> ms.set_context(jit_syntax_level=ms.STRICT)
+        >>> ms.set_context(debug_level=ms.context.DEBUG)
+        >>> ms.set_context(gpu_config={"conv_fprop_algo": "performance", "conv_allow_tf32": True,
+        ...                "matmul_allow_tf32": True})
+        >>> ms.set_context(jit_config={"jit_level": "O0"})
+    """
+    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)
+    _check_ascend_device_context_initialized(device, kwargs)
+
+    for key, value in kwargs.items():
+        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
+        if key in ('enable_auto_mixed_precision', 'enable_dump', 'save_dump_path'):
+            logger.warning(f"For 'context.set_context', '{key}' parameter is deprecated. "
+                           "For details, please see the interface parameter API comments")
+            continue
+        _check_key(key)
+        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 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 key == 'debug_level' and value not in (RELEASE, DEBUG):
+            raise ValueError(f"For 'debug_level', the value should be context.DEBUG"
+                             f" or context.RELEASE, but got {value}.")
+        if key == 'enable_compile_cache':
+            setattr(ctx, key, value)
+            ctx.set_param(ms_ctx_param.__members__[key], int(value))
+            continue
+        if not _check_target_specific_cfgs(device, key):
+            continue
+        if hasattr(ctx, key):
+            setattr(ctx, key, value)
+            continue
+        if key in ctx.setters:
+            ctx.setters[key](ctx, value)
+            continue
+        # enum variables beginning with '_' are for internal use
+        if key in ms_ctx_param.__members__ and key[0] != '_':
+            ctx.set_param(ms_ctx_param.__members__[key], value)
+            continue
+        raise ValueError(f"For 'context.set_context', the keyword argument {key} is not recognized! For detailed "
+                         f"usage of 'set_context', please refer to the Mindspore official website.")
+
+
+def get_context(attr_key):
+    """
+    Get context attribute value according to the input key.
+    If some attributes are not set, they will be automatically obtained.
+
+    Args:
+        attr_key (str): The key of the attribute.
+
+    Returns:
+        Object, The value of given attribute key.
+
+    Raises:
+        ValueError: If input key is not an attribute in context.
+    Examples:
+        >>> import mindspore as ms
+        >>> ms.get_context("device_target")
+        >>> ms.get_context("device_id")
+    """
+    ctx = _context()
+    device = ctx.get_param(ms_ctx_param.device_target)
+    _ = _check_target_specific_cfgs(device, attr_key)
+    if hasattr(ctx, attr_key):
+        return getattr(ctx, attr_key)
+    # enum variables beginning with '_' are for internal use
+    if attr_key in ms_ctx_param.__members__ and attr_key[0] != '_':
+        return ctx.get_param(ms_ctx_param.__members__[attr_key])
+    raise ValueError(f"For 'context.get_context', the argument {attr_key} is not recognized! For detailed "
+                     f"usage of 'get_context', please refer to the Mindspore official website.")
+
+
+def _get_mode():
+    """
+    Get execution mode. Only for internal using.
+
+    Returns:
+        Object: The Value of execution mode.
+    """
+    ctx = _context()
+    return ctx.get_mode()
+
+
+def get_jit_config():
+    """
+    Get global jit config.
+
+    Returns:
+        Object: The Value of jit config.
+    """
+    ctx = _context()
+    return ctx.get_jit_config()
+
+
+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``.
+
+    - ``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.
+    """
+
+    STAND_ALONE = "stand_alone"
+    DATA_PARALLEL = "data_parallel"
+    HYBRID_PARALLEL = "hybrid_parallel"
+    SEMI_AUTO_PARALLEL = "semi_auto_parallel"
+    AUTO_PARALLEL = "auto_parallel"
+    MODE_LIST = [STAND_ALONE, DATA_PARALLEL, HYBRID_PARALLEL, SEMI_AUTO_PARALLEL, AUTO_PARALLEL]
+
+
+@args_type_check(enable_ps=bool)
+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:
+
+        - 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_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`` .
+        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: ``''`` .
+
+    Raises:
+        ValueError: If input key is not the attribute in parameter server training mode context.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> ms.set_ps_context(enable_ps=True, enable_ssl=True, client_password='123456', server_password='123456')
+    """
+    _set_ps_context(**kwargs)
+
+
+def get_ps_context(attr_key):
+    """
+    Get parameter server training mode context attribute value according to the key.
+
+    Args:
+        attr_key (str): The key of the attribute:
+
+            - 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: ``''`` .
+
+    Returns:
+        Returns attribute value according to the key.
+
+    Raises:
+        ValueError: If input key is not attribute in auto parallel context.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> ms.get_ps_context("enable_ps")
+    """
+    return _get_ps_context(attr_key)
+
+
+def reset_ps_context():
+    """
+    Reset parameter server training mode context attributes to the default values.
+
+    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()
+
+
+_hccl_connect_timeout = '600'
+
+
+def _init_parallel_env():
+    """Set hccl connect timeout."""
+    if 'HCCL_CONNECT_TIMEOUT' not in os.environ:
+        os.environ['HCCL_CONNECT_TIMEOUT'] = _hccl_connect_timeout
+
+
+_init_parallel_env()