mindspore 2.0.0a0__cp37-none-any.whl → 2.0.0rc1__cp37-none-any.whl

mindspore/ops/operations/spectral_ops.py

@@ -14,7 +14,7 @@
 # ============================================================================
 
 """Spectral operators."""
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore.common import dtype as mstype
 from mindspore.ops.primitive import Primitive, prim_attr_register
 
@@ -23,10 +23,26 @@ class BartlettWindow(Primitive):
     r"""
     Bartlett window function.
 
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
     Refer to :func:`mindspore.ops.bartlett_window` for more details.
 
+    Args:
+        periodic (bool, optional): If True, returns a window to be used as periodic function.
+          If False, return a symmetric window. Default: True.
+        dtype (mindspore.dtype, optional): The desired datatype of returned tensor.
+            Only float16, float32 and float64 are allowed. Default: mstype.float32.
+
+    Inputs:
+        - **window_length** (Tensor) - The size of returned window, with data type int32, int64.
+          The input data should be an integer with a value of [0, 1000000].
+
+    Outputs:
+        A 1-D tensor of size `window_length` containing the window. Its datatype is set by the attr `dtype`.
+
     Supported Platforms:
-        ``GPU`` ``CPU``
+        ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
         >>> window_length = Tensor(5, mstype.int32)
@@ -50,10 +66,26 @@ class BlackmanWindow(Primitive):
     r"""
     Blackman window function.
 
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
     Refer to :func:`mindspore.ops.blackman_window` for more details.
 
+    Args:
+        periodic (bool, optional): If True, returns a window to be used as periodic function.
+            If False, return a symmetric window. Default: True.
+        dtype (mindspore.dtype, optional): the desired data type of returned tensor.
+            Only float16, float32 and float64 is allowed. Default: mstype.float32.
+
+    Inputs:
+        - **window_length** (Tensor) - the size of returned window, with data type int32, int64.
+          The input data should be an integer with a value of [0, 1000000].
+
+    Outputs:
+        A 1-D tensor of size `window_length` containing the window. Its datatype is set by the attr `dtype`.
+
     Supported Platforms:
-        ``GPU`` ``CPU``
+        ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
         >>> window_length = Tensor(10, mindspore.int32)

mindspore/ops/primitive.py

@@ -24,8 +24,9 @@ from mindspore.parallel._utils import _is_in_auto_parallel_mode, _is_in_data_par
 from mindspore.parallel._ps_context import _is_ps_mode, _is_role_sched
 from mindspore.common.parameter import Parameter
 from mindspore.common.api import _pynative_executor
+from mindspore.common._stub_tensor import _convert_stub
 from mindspore._c_expression import Primitive_, prim_type
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.ops import signature as sig
 
 
@@ -486,10 +487,10 @@ class PrimitiveWithCheck(Primitive):
         ...     def __init__(self):
         ...         pass
         ...     def check_shape(self, input_x):
-        ...         validator.check_int(len(input_x), 1, Rel.GE, 'input_x rank', self.name)
+        ...         Validator.check_int(len(input_x), 1, validator.GE, 'input_x rank', self.name)
         ...
         ...     def check_dtype(self, input_x):
-        ...         validator.check_subclass("input_x", input_x, mstype.tensor, self.name)
+        ...         Validator.check_subclass("input_x", input_x, mstype.tensor, self.name)
         ...
         >>> # init a Primitive obj
         >>> add = Flatten()
@@ -501,10 +502,18 @@ class PrimitiveWithCheck(Primitive):
 
     def __check__(self, *args):
         """Checking the input shape and the input type of ops is valid """
-        tracks = ['dtype', 'shape']
-        for track in tracks:
-            fn = getattr(self, 'check_' + track)
-            fn(*(x[track] for x in args))
+        check_dtype_fn = getattr(self, 'check_dtype')
+        check_dtype_fn(*(x['dtype'] for x in args))
+
+        is_shape_known = True
+        for x in args:
+            shape = x['shape']
+            if shape is None or -1 in shape or -2 in shape:
+                is_shape_known = False
+                break
+        if is_shape_known:
+            check_shape_fn = getattr(self, 'check_shape')
+            check_shape_fn(*(x['shape'] for x in args))
 
     def _clone(self):
         """
@@ -731,6 +740,24 @@ def prim_attr_register(fn):
     return deco
 
 
+def _check_contains_variable(item_dtype, item_value):
+    """
+    Check whether the item is or contains variable.
+    """
+    if isinstance(item_value, (list, tuple)):
+        for i, element in enumerate(item_value):
+            if _check_contains_variable(item_dtype[i], element):
+                return True
+    elif isinstance(item_value, dict):
+        for i in range(len(item_value)):
+            if _check_contains_variable(item_dtype[i], list(item_value.keys())[i]):
+                return True
+        for i in range(len(item_value)):
+            if _check_contains_variable(item_dtype[i], list(item_value.values())[i]):
+                return True
+    return item_dtype is not None and item_value is None
+
+
 def constexpr(fn=None, get_instance=True, name=None, reuse_result=True, check=True):
     """
     Creates a PrimitiveWithInfer operator that can infer the value at compile time. We can use it to define a function
@@ -778,13 +805,14 @@ def constexpr(fn=None, get_instance=True, name=None, reuse_result=True, check=Tr
                 PrimitiveWithInfer.__init__(self, op_name)
                 self.set_const_prim(True)
                 self.fn = fn
+                self.add_prim_attr('constexpr_prim', True)
                 if not reuse_result:
                     self.add_prim_attr('forbid_reuse_result', True)
 
             def __infer__(self, *args):
                 value_args = []
                 for item in args:
-                    if (item["dtype"] is not None and item["value"] is None and check):
+                    if _check_contains_variable(item["dtype"], item["value"]) and check:
                         logger.warning("The \"" + self.name + "\" is a constexpr function." \
                                                               " The input arguments must be all constant value.")
                     value_args.append(item["value"])
@@ -802,8 +830,82 @@ def constexpr(fn=None, get_instance=True, name=None, reuse_result=True, check=Tr
     return deco
 
 
-@_wrap_func
+def _primexpr(fn=None, get_instance=True, name=None, reuse_result=True):
+    """
+    _primexpr is similar as constexpr except that when the input to the function decorated by _primexpr contains
+    variable, the function will be compiled as graph.
+
+    _primexpr is only for internal use.
+
+    Args:
+        fn (function): A `fn` use as the infer_value of the output operator. Default: None.
+        get_instance (bool): If true, return the instance of operator,
+                             otherwise return the operator class. Default: True.
+        name (str): Defines the operator name. If `name` is None, use the function name as op name. Default: None.
+        reuse_result (bool): If true, the operator will be executed once and reuse the result next time,
+                             otherwise the operator will always be executed. Default: True.
+    """
+    def deco(fn):
+        """Decorator for CompileOp."""
+
+        class CompileOp(PrimitiveWithInfer):
+            """
+            CompileOp is a temporary operator used to execute the constexpr function.
+            """
+
+            def __init__(self):
+                op_name = name if name else fn.__name__
+                PrimitiveWithInfer.__init__(self, op_name)
+                self.set_const_prim(True)
+                self.fn = fn
+                self.add_prim_attr('constexpr_prim', True)
+                if not reuse_result:
+                    self.add_prim_attr('forbid_reuse_result', True)
+
+            def __infer__(self, *args):
+                value_args = []
+                for item in args:
+                    if _check_contains_variable(item["dtype"], item["value"]):
+                        return {'dtype': None, 'shape': None, 'value': None, 'fn': (fn,)}
+                    value_args.append(item["value"])
+                return {'dtype': None, 'shape': None, 'value': fn(*value_args)}
+
+            def __call__(self, *args, **kwargs):
+                return fn(*args, **kwargs)
+
+        if get_instance:
+            return CompileOp()
+        return CompileOp
+
+    if fn is not None:
+        return deco(fn)
+    return deco
+
+
+_RUN_OP_ASYNC = True
+
+
 def _run_op(obj, op_name, args):
     """Single op execution function supported by ge in PyNative mode."""
+    if _RUN_OP_ASYNC:
+        stub = _pynative_executor.run_op_async(obj, args)
+        return _convert_stub(stub)
+    return _run_op_sync(obj, op_name, args)
+
+
+@_wrap_func
+def _run_op_sync(obj, op_name, args):
+    """Single op execution function in synchronous mode."""
     output = _pynative_executor.real_run_op(obj, op_name, args)
     return output
+
+
+class _PrimitiveC(Primitive):
+    def __init__(self, name, attrs):
+        super().__init__(name)
+        for key, value in attrs.items():
+            super().add_prim_attr(key, value)
+
+
+def _get_primitivec(name, attrs):
+    return _PrimitiveC(name, attrs)

mindspore/parallel/_auto_parallel_context.py

@@ -15,13 +15,15 @@
 """Context of auto parallel"""
 from __future__ import absolute_import
 import os
+import copy
 import threading
 from mindspore import context
 import mindspore.log as logger
 from mindspore.parallel._dp_allreduce_fusion import _set_fusion_strategy_by_idx, _set_fusion_strategy_by_size
 from mindspore.parallel._ps_context import _is_role_pserver
 from mindspore._c_expression import AutoParallelContext
-from mindspore._checkparam import args_type_check, Validator
+from mindspore._checkparam import args_type_check
+from mindspore import _checkparam as Validator
 
 _MAX_GROUP_NAME_LEN = 127
 _DEFAULT_HCCL_FUSION_GROUP_NAME = "hccl_world_groupsum1"
@@ -40,6 +42,18 @@ class _ParallelFusionConfig:
     AUTO = "auto"
     INDEX = "index"
     SIZE = "size"
+    OPENSTATE = "openstate"
+    CONFIG = {"openstate": True,
+              "allreduce": {"mode": "auto", "config": None},
+              "allgather": {"mode": "auto", "config": None},
+              "reducescatter": {"mode": "auto", "config": None}}
+
+    @classmethod
+    def reset(cls):
+        cls.CONFIG = {"openstate": True,
+                      "allreduce": {"mode": "auto", "config": None},
+                      "allgather": {"mode": "auto", "config": None},
+                      "reducescatter": {"mode": "auto", "config": None}}
 
 
 class _ParallelOptimizerConfig:
@@ -117,6 +131,9 @@ class _AutoParallelContext:
             KeyError: When key of comm_fusion is not 'allreduce'.
         """
         self.check_context_handle()
+        config = copy.deepcopy(config)
+        if _ParallelFusionConfig.OPENSTATE not in config.keys():
+            config[_ParallelFusionConfig.OPENSTATE] = True
         for key in list(config.keys()):
             if key == _ParallelFusionConfig.ALLREDUCE:
                 self._set_allreduce_comm_fusion(config[key])
@@ -124,91 +141,18 @@ class _AutoParallelContext:
                 self._set_allgather_comm_fusion(config[key], key)
             elif key == _ParallelFusionConfig.REDUCESCATTER:
                 self._set_allgather_comm_fusion(config[key], key)
+            elif key == _ParallelFusionConfig.OPENSTATE:
+                self._set_openstate_comm_fusion(config[key])
             else:
-                raise KeyError("comm fusion type must be allreduce, allgather or reducescatter, but got {}".format(key))
+                raise KeyError("comm fusion type must be openstate,"
+                               "allreduce, allgather or reducescatter, but got {}".format(key))
+            if key in _ParallelFusionConfig.CONFIG:
+                _ParallelFusionConfig.CONFIG[key] = config[key]
 
     def get_comm_fusion(self):
         """Get comm fusion config."""
         self.check_context_handle()
-        mode = self._context_handle.get_fusion_mode()
-        if mode in (_ParallelFusionConfig.AUTO, _ParallelFusionConfig.SIZE):
-            config = self.fusion_threshold_mb()
-        if mode == _ParallelFusionConfig.INDEX:
-            config = self.get_all_reduce_fusion_split_indices()
-        return {_ParallelFusionConfig.ALLREDUCE: {_ParallelFusionConfig.MODE: mode,
-                                                  _ParallelFusionConfig.FUSION_CONFIG: config}}
-
-    def _set_allgather_comm_fusion(self, comm_fusion, comm_type="allgather"):
-        """
-        Set allgather and reducescatter fusion method for auto parallel.
-
-        Args:
-            comm_fusion (dict): A dict contains the methods and values for setting the fusion method. Currently it
-                                  supports four fusion methods: `auto` and `size`.
-            comm_type (str): The name of the communication operator, `allgather` or `reducescatter`.
-
-        Raises:
-            KeyError: When key of comm_fusion is not 'mode' or 'config'.
-            KeyError: When `mode` is not 'auto', 'size'.
-        """
-        self.check_context_handle()
-        if comm_type == "allgather" and not self.get_enable_all_gather_fusion():
-            return
-        if comm_type == "reducescatter" and not self.get_enable_reduce_scatter_fusion():
-            return
-        if not isinstance(comm_fusion, dict):
-            raise TypeError("For 'comm_fusion', {} config must be dict, but got the type : {}.".format(
-                comm_type, type(comm_fusion)))
-        if _ParallelFusionConfig.MODE not in comm_fusion:
-            raise KeyError("For 'comm_fusion', the key 'mode' should be contained.")
-        if _ParallelFusionConfig.FUSION_CONFIG not in comm_fusion:
-            raise KeyError("For 'comm_fusion', the key 'config' should be contained.")
-        check_mode = [_ParallelFusionConfig.AUTO, _ParallelFusionConfig.SIZE]
-        if comm_fusion[_ParallelFusionConfig.MODE] in check_mode:
-            self._context_handle.set_fusion_mode(comm_fusion[_ParallelFusionConfig.MODE])
-        else:
-            raise KeyError("fusion method mode must be auto or size, but got {}".format(
-                comm_fusion[_ParallelFusionConfig.MODE]))
-
-        fusion_threshold = 64
-        if comm_fusion[_ParallelFusionConfig.MODE] != _ParallelFusionConfig.AUTO:
-            fusion_threshold = comm_fusion[_ParallelFusionConfig.FUSION_CONFIG]
-        self.set_fusion_threshold_mb(fusion_threshold, comm_type)
-
-    def _set_allreduce_comm_fusion(self, comm_fusion):
-        """
-        Set fusion method for auto parallel.
-
-        Args:
-            comm_fusion (dict): A dict contains the methods and values for setting the fusion method. Currently it
-                                  supports four fusion methods: `auto`, `size` and `index`.
-
-        Raises:
-            KeyError: When key of comm_fusion is not 'mode' or 'config'.
-            KeyError: When `mode` is not 'auto', 'size' or 'index'.
-        """
-        self.check_context_handle()
-        if not self.get_enable_all_reduce_fusion():
-            return
-        if not isinstance(comm_fusion, dict):
-            raise TypeError("For 'comm_fusion', the 'allreduce' config must be dict, but got the type : {}.".format(
-                type(comm_fusion)))
-        if _ParallelFusionConfig.MODE not in comm_fusion:
-            raise KeyError("For 'comm_fusion', the key 'mode' should be contained.")
-        if _ParallelFusionConfig.FUSION_CONFIG not in comm_fusion:
-            raise KeyError("For 'comm_fusion', the key 'config' should be contained.")
-        check_mode = [_ParallelFusionConfig.AUTO, _ParallelFusionConfig.INDEX, _ParallelFusionConfig.SIZE]
-        if comm_fusion[_ParallelFusionConfig.MODE] in check_mode:
-            self._context_handle.set_fusion_mode(comm_fusion[_ParallelFusionConfig.MODE])
-        else:
-            raise KeyError("fusion method mode must be auto, index or size, but got {}".format(
-                comm_fusion[_ParallelFusionConfig.MODE]))
-        if comm_fusion[_ParallelFusionConfig.MODE] == _ParallelFusionConfig.AUTO:
-            self.set_fusion_threshold_mb(fusion_threshold=64)
-        if comm_fusion[_ParallelFusionConfig.MODE] == _ParallelFusionConfig.SIZE:
-            self.set_fusion_threshold_mb(comm_fusion[_ParallelFusionConfig.FUSION_CONFIG])
-        if comm_fusion[_ParallelFusionConfig.MODE] == _ParallelFusionConfig.INDEX:
-            self.set_all_reduce_fusion_split_indices(comm_fusion[_ParallelFusionConfig.FUSION_CONFIG])
+        return _ParallelFusionConfig.CONFIG
 
     def set_fusion_threshold_mb(self, fusion_threshold=64, comm_type="allreduce"):
         """
@@ -521,6 +465,9 @@ class _AutoParallelContext:
                 if not isinstance(dim, int):
                     raise TypeError("For 'set_auto_parallel_context', the element of argument "
                                     "'dataset_strategy' must be int type, but got the type : {} .".format(type(dim)))
+        if context.get_context('mode') == context.PYNATIVE_MODE:
+            raise ValueError("In PyNative mode, the setting value of 'dataset_strategy' must be either 'full_batch' "
+                             f"or 'data_parallel', but got {dataset_strategy}.")
         self._dataset_strategy_using_str = False
         self._context_handle.set_dataset_strategy(dataset_strategy)
 
@@ -531,7 +478,11 @@ class _AutoParallelContext:
             if self._context_handle.get_full_batch():
                 return "full_batch"
             return "data_parallel"
-        return self._context_handle.get_dataset_strategy()
+        dataset_strategy = self._context_handle.get_dataset_strategy()
+        if context.get_context('mode') == context.PYNATIVE_MODE:
+            raise ValueError("In PyNative mode, the value of 'dataset_strategy' must be either 'full_batch' "
+                             f"or 'data_parallel', but got the setting value is {dataset_strategy}.")
+        return dataset_strategy
 
     def set_grad_accumulation_step(self, grad_accumulation_step):
         """
@@ -567,6 +518,52 @@ class _AutoParallelContext:
         self.check_context_handle()
         return self._context_handle.get_strategy_ckpt_save_file()
 
+    def set_strategy_ckpt_config(self, strategy_ckpt_config):
+        """
+        Set strategy checkpoint config.
+
+        Args:
+            strategy_ckpt_config (dict): The strategy checkpoint config.
+        """
+        self.check_context_handle()
+        if not isinstance(strategy_ckpt_config, dict):
+            raise TypeError("For 'set_auto_parallel_context', the argument 'strategy_ckpt_config' "
+                            "must be dict, but got the type : {}.".format(type(strategy_ckpt_config)))
+        for config_name in strategy_ckpt_config:
+            unknown_config = []
+            if config_name not in ["load_file", "save_file", "only_trainable_params"]:
+                unknown_config.append(config_name)
+
+            if unknown_config:
+                raise ValueError("Unknown config: {}".format(unknown_config))
+        if "load_file" in strategy_ckpt_config:
+            load_file = strategy_ckpt_config.get("load_file")
+            if not isinstance(load_file, str):
+                raise TypeError("For 'set_auto_parallel_context().set_strategy_ckpt_config', "
+                                "the argument 'load_file' must be str, but got the type : {} .".format(type(load_file)))
+            self._context_handle.set_strategy_ckpt_load_file(load_file)
+        if "save_file" in strategy_ckpt_config:
+            save_file = strategy_ckpt_config.get("save_file")
+            if not isinstance(save_file, str):
+                raise TypeError("For 'set_auto_parallel_context().set_strategy_ckpt_config', "
+                                "the argument 'save_file' must be str, but got the type : {} .".format(type(save_file)))
+            self._context_handle.set_strategy_ckpt_save_file(save_file)
+        if "only_trainable_params" in strategy_ckpt_config:
+            only_trainable_params = strategy_ckpt_config.get("only_trainable_params")
+            if not isinstance(only_trainable_params, bool):
+                raise TypeError("For 'set_auto_parallel_context().set_strategy_ckpt_config', "
+                                "the argument 'only_trainable_params' must be bool,"
+                                " but got the type : {} .".format(type(only_trainable_params)))
+            self._context_handle.set_stra_file_only_trainable_params(only_trainable_params)
+
+    def get_strategy_ckpt_config(self):
+        """Get strategy checkpoint config."""
+        self.check_context_handle()
+        load_file = self._context_handle.get_strategy_ckpt_load_file()
+        save_file = self._context_handle.get_strategy_ckpt_save_file()
+        only_trainable_param = self._context_handle.get_stra_file_only_trainable_params()
+        return {"load_file": load_file, "save_file": save_file, "only_trainable_params": only_trainable_param}
+
     def set_group_ckpt_save_file(self, group_ckpt_save_file):
         """Set group checkpoint save path."""
         self.check_context_handle()
@@ -912,6 +909,7 @@ class _AutoParallelContext:
         return self._context_handle.get_optimizer_weight_shard_aggregated_save()
 
     def get_full_batch_is_set(self):
+        """Get full batch attr"""
         self.check_context_handle()
         return self._context_handle.get_full_batch_is_set()
 
@@ -919,6 +917,7 @@ class _AutoParallelContext:
         """Reset all settings."""
         self.check_context_handle()
         self._context_handle.reset()
+        _ParallelFusionConfig.reset()
 
     def _check_and_default_group(self, group):
         """Validate the given group, if group is empty, returns a default fusion group"""
@@ -936,6 +935,99 @@ class _AutoParallelContext:
                 group = _DEFAULT_NCCL_FUSION_GROUP_NAME
         return group
 
+    def _set_allgather_comm_fusion(self, comm_fusion, comm_type="allgather"):
+        """
+        Set allgather and reducescatter fusion method for auto parallel.
+
+        Args:
+            comm_fusion (dict): A dict contains the methods and values for setting the fusion method. Currently it
+                                  supports four fusion methods: `auto` and `size`.
+            comm_type (str): The name of the communication operator, `allgather` or `reducescatter`.
+
+        Raises:
+            KeyError: When key of comm_fusion is not 'mode' or 'config'.
+            KeyError: When `mode` is not 'auto', 'size'.
+        """
+        self.check_context_handle()
+        if comm_type == "allgather" and not self.get_enable_all_gather_fusion():
+            return
+        if comm_type == "reducescatter" and not self.get_enable_reduce_scatter_fusion():
+            return
+        if not isinstance(comm_fusion, dict):
+            raise TypeError("For 'comm_fusion', {} config must be dict, but got the type : {}.".format(
+                comm_type, type(comm_fusion)))
+        if _ParallelFusionConfig.MODE not in comm_fusion:
+            raise KeyError("For 'comm_fusion', the key 'mode' should be contained.")
+        if _ParallelFusionConfig.FUSION_CONFIG not in comm_fusion:
+            raise KeyError("For 'comm_fusion', the key 'config' should be contained.")
+        check_mode = [_ParallelFusionConfig.AUTO, _ParallelFusionConfig.SIZE]
+        if comm_fusion[_ParallelFusionConfig.MODE] in check_mode:
+            self._context_handle.set_fusion_mode(comm_fusion[_ParallelFusionConfig.MODE])
+        else:
+            raise KeyError("fusion method mode must be auto or size, but got {}".format(
+                comm_fusion[_ParallelFusionConfig.MODE]))
+
+        fusion_threshold = 64
+        if comm_fusion[_ParallelFusionConfig.MODE] != _ParallelFusionConfig.AUTO:
+            fusion_threshold = comm_fusion[_ParallelFusionConfig.FUSION_CONFIG]
+        self.set_fusion_threshold_mb(fusion_threshold, comm_type)
+
+    def _set_allreduce_comm_fusion(self, comm_fusion):
+        """
+        Set fusion method for auto parallel.
+
+        Args:
+            comm_fusion (dict): A dict contains the methods and values for setting the fusion method. Currently it
+                                  supports four fusion methods: `auto`, `size` and `index`.
+
+        Raises:
+            KeyError: When key of comm_fusion is not 'mode' or 'config'.
+            KeyError: When `mode` is not 'auto', 'size' or 'index'.
+        """
+        self.check_context_handle()
+        if not self.get_enable_all_reduce_fusion():
+            return
+        if not isinstance(comm_fusion, dict):
+            raise TypeError("For 'comm_fusion', the 'allreduce' config must be dict, but got the type : {}.".format(
+                type(comm_fusion)))
+        if _ParallelFusionConfig.MODE not in comm_fusion:
+            raise KeyError("For 'comm_fusion', the key 'mode' should be contained.")
+        if _ParallelFusionConfig.FUSION_CONFIG not in comm_fusion:
+            raise KeyError("For 'comm_fusion', the key 'config' should be contained.")
+        check_mode = [_ParallelFusionConfig.AUTO, _ParallelFusionConfig.INDEX, _ParallelFusionConfig.SIZE]
+        if comm_fusion[_ParallelFusionConfig.MODE] in check_mode:
+            self._context_handle.set_fusion_mode(comm_fusion[_ParallelFusionConfig.MODE])
+        else:
+            raise KeyError("fusion method mode must be auto, index or size, but got {}".format(
+                comm_fusion[_ParallelFusionConfig.MODE]))
+        if comm_fusion[_ParallelFusionConfig.MODE] == _ParallelFusionConfig.AUTO:
+            self.set_fusion_threshold_mb(fusion_threshold=64)
+        if comm_fusion[_ParallelFusionConfig.MODE] == _ParallelFusionConfig.SIZE:
+            self.set_fusion_threshold_mb(comm_fusion[_ParallelFusionConfig.FUSION_CONFIG])
+        if comm_fusion[_ParallelFusionConfig.MODE] == _ParallelFusionConfig.INDEX:
+            self.set_all_reduce_fusion_split_indices(comm_fusion[_ParallelFusionConfig.FUSION_CONFIG])
+
+    def _set_openstate_comm_fusion(self, openstate):
+        """
+        Set open state for comm fusion.
+
+        Args:
+            openstate (bool): The open state value to set the fusion method whether or not. Currently it
+                                  supports two states: `True`, or `Flase`.
+
+        Raises:
+            TypeError: When the value is not bool.
+        """
+        self.check_context_handle()
+        if not isinstance(openstate, bool):
+            raise TypeError("For 'comm_fusion', the 'openstate' must be bool, but got the type : {}.".format(
+                type(openstate)))
+        if not openstate:
+            self.set_enable_all_reduce_fusion(openstate)
+            self.set_enable_all_gather_fusion(openstate)
+            self.set_enable_reduce_scatter_fusion(openstate)
+
+
 
 _AUTO_PARALLEL_CONTEXT = None
 
@@ -978,6 +1070,7 @@ _set_auto_parallel_context_func_map = {
     "optimizer_weight_shard_aggregated_save": auto_parallel_context().set_optimizer_weight_shard_aggregated_save,
     "sharding_propagation": auto_parallel_context().set_sharding_propagation,
     "enable_alltoall": auto_parallel_context().set_enable_alltoall,
+    "strategy_ckpt_config": auto_parallel_context().set_strategy_ckpt_config,
     "comm_fusion": auto_parallel_context().set_comm_fusion}
 
 
@@ -1005,6 +1098,7 @@ _get_auto_parallel_context_func_map = {
     "sharding_propagation": auto_parallel_context().get_sharding_propagation,
     "enable_alltoall": auto_parallel_context().get_enable_alltoall,
     "comm_fusion": auto_parallel_context().get_comm_fusion,
+    "strategy_ckpt_config": auto_parallel_context().get_strategy_ckpt_config,
     "full_batch_is_set": auto_parallel_context().get_full_batch_is_set}
 
 
@@ -1014,7 +1108,8 @@ _get_auto_parallel_context_func_map = {
                  strategy_ckpt_save_file=str, full_batch=bool, enable_parallel_optimizer=bool,
                  grad_accumulation_step=int, all_reduce_fusion_config=list, group_ckpt_save_file=str,
                  communi_parallel_mode=str, optimizer_weight_shard_size=int, sharding_propagation=bool,
-                 optimizer_weight_shard_aggregated_save=bool, enable_alltoall=bool, comm_fusion=dict)
+                 optimizer_weight_shard_aggregated_save=bool, enable_alltoall=bool, comm_fusion=dict,
+                 strategy_ckpt_config=dict)
 
 def _set_auto_parallel_context(**kwargs):
     """
@@ -1091,12 +1186,23 @@ def _set_auto_parallel_context(**kwargs):
                     communication fusion config has two keys: "mode" and "config".
                     It supports following communication fusion types and configurations:
 
+                    - openstate: Whether turn on the communication fusion or not. If `openstate` is `True`, turn on
+                        the communication fusion, otherwise, turn off the communication fusion. Default: `True`.
+
                     - allreduce: if communication fusion type is `allreduce`. The `mode` contains: `auto`, `size`
                         and `index`. In `auto` mode, allreduce fusion is configured by gradients size, and the default
                         fusion threshold is `64` MB. In 'size' mode, allreduce fusion is configured by gradients size
                         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`.
+
 
     Raises:
         ValueError: If input key is not attribute in auto parallel context.