mindspore 2.2.14__cp39-cp39-win_amd64.whl → 2.3.0__cp39-cp39-win_amd64.whl

mindspore/context.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2023 Huawei Technologies Co., Ltd
+# 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.
@@ -34,6 +34,7 @@ from mindspore.parallel._auto_parallel_context import _set_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',
@@ -50,6 +51,10 @@ STRICT = 0
 COMPATIBLE = 1
 LAX = 2
 
+# Enumerate for the property 'debug_level'.
+RELEASE = 0
+DEBUG = 1
+
 
 def _make_directory(path):
     """Make directory."""
@@ -161,6 +166,9 @@ class _Context:
         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)
@@ -176,7 +184,11 @@ class _Context:
 
     def get_mode(self):
         """Get current mode."""
-        return self.get_param(ms_ctx_param.mode)
+        return self._mode
+
+    def get_jit_config(self):
+        """Get current jit_config."""
+        return self._jit_config
 
     def set_mode(self, mode):
         """
@@ -204,6 +216,7 @@ class _Context:
             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"""
@@ -212,6 +225,13 @@ class _Context:
                              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".
@@ -268,11 +288,16 @@ class _Context:
                             "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"`` .
-                - op_precision_mode (str): config file path.
                 - parallel_speed_up_json_path(Union[str, None]): The path to the parallel speed up json file.
                   If its value is None or '', it does not take effect. Default None.
+                - 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",
@@ -284,8 +309,15 @@ class _Context:
             'conv_allow_hf32': [True, False],
             'exception_dump': ["0", "1", "2"],
             'op_precision_mode': (str,),
+            'ge_options': (dict,),
             'parallel_speed_up_json_path': (str, None),
-            'topo_order': (dict,)
+            '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'),
@@ -294,8 +326,15 @@ class _Context:
             '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())
@@ -354,6 +393,31 @@ class _Context:
             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:
@@ -475,9 +539,13 @@ class _Context:
 
     def set_mempool_block_size(self, mempool_block_size):
         """Set the block size of memory pool."""
-        if _get_mode() == GRAPH_MODE:
+        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.")
+                           "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 "
@@ -563,8 +631,10 @@ class _Context:
         '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
@@ -612,6 +682,16 @@ class _Context:
             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)
@@ -620,6 +700,28 @@ class _Context:
                              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.
@@ -642,6 +744,30 @@ class _Context:
         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 == "":
@@ -652,23 +778,34 @@ class _Context:
                              f"{speedup_config_real_path} does not exist, please check whether the "
                              f"'parallel_speed_up_json_path' is correct.")
         try:
-            valid_option = {"recompute_comm_overlap": ms_ctx_param.recompute_comm_overlap,
-                            "matmul_grad_comm_overlap": ms_ctx_param.matmul_grad_comm_overlap,
-                            "enable_task_opt": ms_ctx_param.enable_task_opt,
-                            "enable_grad_comm_opt": ms_ctx_param.enable_grad_comm_opt,
-                            "interleaved_matmul_comm": ms_ctx_param.interleaved_matmul_comm,
-                            "enable_opt_shard_comm_opt": ms_ctx_param.enable_opt_shard_comm_opt,
-                            "interleaved_layernorm_comm": ms_ctx_param.interleaved_layernorm_comm}
+            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 k, v in speedup_config.items():
-                    if not isinstance(k, str):
-                        raise TypeError("key {} is not a str".format(k))
-                    if k not in valid_option:
-                        raise ValueError("key {} should be one of {}.".format(k, valid_option.keys()))
-                    if not isinstance(v, bool):
-                        raise TypeError("value {} is not a bool".format(v))
-                    self.set_param(valid_option.get(k), v)
+                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 {} "
@@ -705,8 +842,9 @@ def _context():
                  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_config=dict, parallel_optimizer_config=dict,
-                 comm_fusion=dict, strategy_ckpt_config=dict)
+                 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.
@@ -733,8 +871,11 @@ def set_auto_parallel_context(**kwargs):
     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:
@@ -744,6 +885,8 @@ def set_auto_parallel_context(**kwargs):
                      "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"`` .
@@ -758,15 +901,16 @@ def set_auto_parallel_context(**kwargs):
 
                      - auto_parallel: Achieving parallelism automatically.
         search_mode (str): There are three kinds of shard strategy search modes: ``"recursive_programming"`` ,
-                     ``"dynamic_programming"`` and ``"sharding_propagation"`` . Default: ``"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.
 
-                     - dynamic_programming: Dynamic programming search mode.
-
                      - 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
@@ -792,6 +936,9 @@ def set_auto_parallel_context(**kwargs):
                        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`` .
@@ -801,6 +948,8 @@ def set_auto_parallel_context(**kwargs):
                         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:
 
@@ -866,14 +1015,18 @@ def set_auto_parallel_context(**kwargs):
                         - 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: ''
+                          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: ''
+                          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.
@@ -885,8 +1038,8 @@ def set_auto_parallel_context(**kwargs):
         >>> 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="dynamic_programming")
-        >>> ms.set_auto_parallel_context(auto_parallel_search_mode="dynamic_programming")
+        >>> 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")
@@ -895,6 +1048,7 @@ def set_auto_parallel_context(**kwargs):
         >>> ms.set_auto_parallel_context(enable_alltoall=False)
         >>> ms.set_auto_parallel_context(all_reduce_fusion_config=[8, 160])
         >>> ms.set_auto_parallel_context(pipeline_stages=2)
+        >>> 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)
@@ -943,9 +1097,12 @@ def reset_auto_parallel_context():
     - 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
@@ -1035,6 +1192,23 @@ def _check_target_specific_cfgs(device, arg_key):
     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,
@@ -1043,8 +1217,8 @@ def _check_target_specific_cfgs(device, arg_key):
                  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,
-                 jit_enable_inplace_ops=bool, gpu_config=dict)
+                 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.
@@ -1093,6 +1267,8 @@ def set_context(**kwargs):
     |                         |  reserve_class_name_in_scope |  CPU/GPU/Ascend            |
     |                         +------------------------------+----------------------------+
     |                         |  pynative_synchronize        |  CPU/GPU/Ascend            |
+    |                         +------------------------------+----------------------------+
+    |                         |  debug_level                 |  CPU/GPU/Ascend            |
     +-------------------------+------------------------------+----------------------------+
     | Executive Control       |   mode                       |   CPU/GPU/Ascend           |
     |                         +------------------------------+----------------------------+
@@ -1133,6 +1309,8 @@ def set_context(**kwargs):
     |                         |  jit_syntax_level            |  CPU/GPU/Ascend            |
     |                         +------------------------------+----------------------------+
     |                         |  gpu_config                  |  GPU                       |
+    |                         +------------------------------+----------------------------+
+    |                         |  jit_config                  |  CPU/GPU/Ascend            |
     +-------------------------+------------------------------+----------------------------+
 
     Args:
@@ -1145,12 +1323,16 @@ def set_context(**kwargs):
             and max_device_memory. 'max_device_memory' should be set before the program runs.
         variable_memory_max_size (str): This parameter is deprecated, and will be removed in a future version.
             Please use parameter 'max_device_memory' instead.
-        mempool_block_size (str): Set the size of the memory pool block in PyNative mode for devices.
-            The format is "xxGB". Default: ``"1GB"`` . Minimum size is "1G". The actual used memory block size is the
-            minimum of the available memory of the device and mempool_block_size.
+        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.
-            Default: ``1900`` .
+            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:
 
@@ -1159,10 +1341,13 @@ def set_context(**kwargs):
             - 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: ".".
+        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.
@@ -1182,6 +1367,8 @@ def set_context(**kwargs):
             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")
 
@@ -1226,7 +1413,7 @@ def set_context(**kwargs):
             If enable_graph_kernel is set to ``True`` , acceleration can be enabled.
             For details of graph kernel fusion, please check
             `Enabling Graph Kernel Fusion
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/optimize/graph_fusion_engine.html>`_.
+            <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.
@@ -1282,7 +1469,7 @@ def set_context(**kwargs):
             the compile cache is loaded. Note that only limited automatic detection for the changes of
             python scripts is supported by now, which means that there is a correctness risk. Default: ``False`` .
             This is an experimental prototype that is subject to change and/or deletion.
-        compile_cache_path (str): Path to save the compile cache. Default: ".".
+        compile_cache_path (str): Path to save the compile cache. Default: ``"."``.
             If the specified directory does not exist, the system will automatically create the directory.
             The cache will be saved to the directory of `compile_cache_path/rank_${rank_id}/`. The `rank_id` is
             the ID of the current device in the cluster.
@@ -1299,16 +1486,18 @@ def set_context(**kwargs):
             of the interfaces would be compiled by MindSpore to the interfaces definition .py file that should be
             guaranteed to be writable. Then compile the .py file to the .pyc or .so file, and could run in Graph mode.
         memory_optimize_level (str): The memory optimize level.
-            Default: O0. The value must be in ['O0', 'O1'].
+            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).
-            - O1: priority memory option, enable SOMAS.
+            - 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 environment variable "GRAPH_OP_RUN=1" is not set; This parameter does not take effect when
+              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.
@@ -1319,22 +1508,27 @@ def set_context(**kwargs):
               is ``force_fp16`` . The value range is as follows:
 
               - force_fp16: When the operator supports both float16 and float32, select float16 directly.
-              - allow_fp32_to_fp16: When the operator does not support the float32 data type, directly reduce
-                the precision of float16.
+              - allow_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: When the operator does not support the float32 data type, directly reduce
-                the precision of bfloat16.
+              - 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. the default value is based on CANN.
+            - 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`` .
 
@@ -1350,24 +1544,66 @@ def set_context(**kwargs):
               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 and AICPU exception operators;
-              for ``"2"``, inputs will be dumped for AICore exception operators. Default: ``"2"`` .
+              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/r2.2/config/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 grad ops and communication ops if True.
+              - 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_task_opt (bool): Enable the optimization of the number of tasks for each communication if True.
+              - 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.
-              - interleaved_matmul_comm (bool): Enable interleaved optimization of Matmul-Comm if True. Default: False.
-              - interleaved_layernorm_comm (bool): Enable interleaved optimization of LayerNorm-Comm if True.
+              - 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.
@@ -1378,6 +1614,12 @@ def set_context(**kwargs):
               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.
@@ -1449,6 +1691,26 @@ def set_context(**kwargs):
             - 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.
 
@@ -1482,16 +1744,23 @@ def set_context(**kwargs):
         >>> 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"})
+        ...                "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, "
@@ -1501,9 +1770,7 @@ def set_context(**kwargs):
             logger.warning(f"For 'context.set_context', '{key}' parameter is deprecated. "
                            "For details, please see the interface parameter API comments")
             continue
-        if key in ('precision_mode', 'jit_compile', 'atomic_clean_policy', 'matmul_allow_hf32', 'conv_allow_hf32',
-                   'op_precision_mode'):
-            raise ValueError(f"Please set '{key}' through parameter ascend_config")
+        _check_key(key)
         if key == 'save_graphs':
             if value is True:
                 value = 2
@@ -1514,6 +1781,13 @@ def set_context(**kwargs):
         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):
@@ -1571,6 +1845,17 @@ def _get_mode():
     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.
@@ -1668,9 +1953,7 @@ def get_ps_context(attr_key):
 
 def reset_ps_context():
     """
-    Reset parameter server training mode context attributes to the default values:
-
-    - enable_ps: False.
+    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`.