mindspore 1.10.0__cp37-none-any.whl → 2.0.0rc1__cp37-none-any.whl

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, _enable_distributed_mindrt
+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"):
         """
@@ -373,8 +317,6 @@ class _AutoParallelContext:
     def get_parallel_mode(self):
         """Get parallel mode."""
         self.check_context_handle()
-        if _is_role_pserver() and not _enable_distributed_mindrt():
-            return context.ParallelMode.STAND_ALONE
         return self._context_handle.get_parallel_mode()
 
     def set_strategy_search_mode(self, search_mode):
@@ -429,13 +371,14 @@ class _AutoParallelContext:
         Set the value of sharding strategy propagation in AUTO_PARALLEL mode. If True, the strategy-configured operators
         will propagate the strategies to other operators with minimum redistribution cost; otherwise, the algorithm
         will search the desired strategies. Default: False.
-        This attribute is replaced by context.set_auto_parallel(search_mode="sharding_propagation").
+        This attribute is replaced by context.set_auto_parallel_context(search_mode="sharding_propagation").
 
         Args:
             sharding_propagation (bool): Enable/disable strategy propagation.
         """
-        logger.warning("This attribute is replaced by context.set_auto_parallel(search_mode='sharding_propagation'), "
-                       "and this attribute will be deleted in a future MindSpore version.")
+        logger.warning("This attribute is replaced by "
+                       "context.set_auto_parallel_context(search_mode='sharding_propagation'), and this attribute will"
+                       " be deleted in a future MindSpore version.")
         self.check_context_handle()
         if not isinstance(sharding_propagation, bool):
             raise TypeError("For 'set_auto_parallel_context().set_sharding_propagation', "
@@ -522,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)
 
@@ -532,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):
         """
@@ -568,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()
@@ -907,18 +903,21 @@ class _AutoParallelContext:
             raise TypeError('optimizer_weight_shard_aggregated_save is invalid type')
         self._context_handle.set_optimizer_weight_shard_aggregated_save(optimizer_weight_shard_aggregated_save)
 
-
     def get_optimizer_weight_shard_aggregated_save(self):
         """Get optimizer_weight_shard_size."""
         self.check_context_handle()
         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()
 
     def reset(self):
         """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}
 
 
@@ -1004,7 +1097,9 @@ _get_auto_parallel_context_func_map = {
     "optimizer_weight_shard_aggregated_save": auto_parallel_context().get_optimizer_weight_shard_aggregated_save,
     "sharding_propagation": auto_parallel_context().get_sharding_propagation,
     "enable_alltoall": auto_parallel_context().get_enable_alltoall,
-    "comm_fusion": auto_parallel_context().get_comm_fusion}
+    "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}
 
 
 @args_type_check(device_num=int, global_rank=int, gradients_mean=bool, gradient_fp32_sync=bool,
@@ -1013,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):
     """
@@ -1090,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.

mindspore/parallel/_cell_wrapper.py

@@ -14,10 +14,12 @@
 # ============================================================================
 """Cell of auto parallel"""
 from __future__ import absolute_import
+from __future__ import division
+
 from mindspore.nn.cell import Cell
 from mindspore.ops.operations.comm_ops import AllGather
 from mindspore.communication import GlobalComm
-from mindspore.common import ms_function
+from mindspore.common import jit
 
 _ALLGATHER_CELL = None
 
@@ -33,7 +35,7 @@ class AllGatherCell(Cell):
         self.allgather = AllGather(group)
         self.add_flags(skip_auto_parallel_compile=True)
 
-    @ms_function()
+    @jit()
     def construct(self, x):
         x = self.allgather(x)
 

mindspore/parallel/_cost_model_context.py

@@ -14,7 +14,10 @@
 # ============================================================================
 """Context of cost_model in auto_parallel"""
 from __future__ import absolute_import
+from __future__ import division
+
 import threading
+
 from mindspore._c_expression import CostModelContext
 from mindspore._checkparam import args_type_check
 

mindspore/parallel/_dp_allreduce_fusion.py

@@ -13,8 +13,9 @@
 # limitations under the License.
 # ============================================================================
 """Data parallel allreduce fusion"""
-
 from __future__ import absolute_import
+from __future__ import division
+
 import ctypes
 
 _MAX_GROUP_NAME_LEN = 127

mindspore/parallel/_offload_context.py

@@ -0,0 +1,185 @@
+# Copyright 2023 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# 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.
+# ============================================================================
+"""Context of offload"""
+from __future__ import absolute_import
+from __future__ import division
+
+import threading
+
+from mindspore._c_expression import OffloadContext
+from mindspore._checkparam import args_type_check
+from mindspore import _checkparam as Validator
+
+
+class _OffloadConfig:
+    """
+    The key of the Offload Config.
+    """
+    ENABLE_OFFLOAD = "enable_offload"
+    OFFLOAD_PARAM = "offload_param"
+    OFFLOAD_PATH = "offload_path"
+    OFFLOAD_CHECKPOINT = "offload_checkpoint"
+    OFFLOAD_DDR_SIZE = "offload_ddr_size"
+    OFFLOAD_DISK_SIZE = "offload_disk_size"
+    ENABLE_AIO = "enable_aio"
+    AIO_BLOCK_SIZE = "aio_block_size"
+    AIO_QUEUE_DEPTH = "aio_queue_depth"
+    ENABLE_PINNED_MEM = "enable_pinned_mem"
+
+
+class _OffloadContext:
+    """
+    _OffloadContext is the configuration for offload.
+
+    Note:
+        Create a context through instantiating Context object is not recommended.
+        Should use offload_context() to get the context since Context is singleton.
+    """
+    _instance = None
+    _instance_lock = threading.Lock()
+
+    def __new__(cls):
+        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._context_handle = OffloadContext.get_instance()
+
+    def check_context_handle(self):
+        """
+        Check context handle.
+
+        Raises:
+            ValueError: If the context handle is none.
+        """
+        if self._context_handle is None:
+            raise ValueError("Context handle is none in context!!!")
+
+    def set_offload_config(self, offload_config):
+        """Set offfload context"""
+        self.check_context_handle()
+        enable_offload = _OffloadConfig.ENABLE_OFFLOAD
+        offload_param = _OffloadConfig.OFFLOAD_PARAM
+        offload_path = _OffloadConfig.OFFLOAD_PATH
+        offload_checkpoint = _OffloadConfig.OFFLOAD_CHECKPOINT
+        offload_ddr_size = _OffloadConfig.OFFLOAD_DDR_SIZE
+        offload_disk_size = _OffloadConfig.OFFLOAD_DISK_SIZE
+        enable_aio = _OffloadConfig.ENABLE_AIO
+        aio_block_size = _OffloadConfig.AIO_BLOCK_SIZE
+        aio_queue_depth = _OffloadConfig.AIO_QUEUE_DEPTH
+        enable_pinned_mem = _OffloadConfig.ENABLE_PINNED_MEM
+
+        for config_name in offload_config:
+            unknown_config = []
+            if config_name not in [enable_offload, offload_param, offload_path, offload_checkpoint,
+                                   offload_ddr_size, offload_disk_size, enable_aio, aio_block_size,
+                                   aio_queue_depth, enable_pinned_mem]:
+                unknown_config.append(config_name)
+
+            if unknown_config:
+                raise ValueError("Unknown config: {}".format(unknown_config))
+
+        if enable_offload in offload_config:
+            Validator.check_bool(
+                offload_config[enable_offload], enable_offload, enable_offload)
+            self._context_handle.set_enable_offload(
+                offload_config[enable_offload])
+
+        if offload_param in offload_config:
+            Validator.check_string(
+                offload_config[offload_param].lower(), ["cpu", "disk"])
+            self._context_handle.set_offload_param(
+                offload_config[offload_param].lower())
+
+        if offload_path in offload_config:
+            if not isinstance(offload_config[offload_path], str):
+                raise TypeError("For 'set_offload_path', "
+                                "the argument 'offload_path' must be str, but got the type : {}."
+                                .format(type(offload_config[offload_path])))
+            self._context_handle.set_offload_path(
+                offload_config[offload_path])
+        if offload_checkpoint in offload_config:
+            Validator.check_string(
+                offload_config[offload_checkpoint].lower(), ["cpu", "disk"])
+            self._context_handle.set_offload_checkpoint(
+                offload_config[offload_checkpoint].lower())
+
+        if offload_ddr_size in offload_config:
+            Validator.check_positive_int(offload_config[offload_ddr_size])
+            self._context_handle.set_offload_ddr_size(
+                offload_config[offload_ddr_size])
+
+        if offload_disk_size in offload_config:
+            Validator.check_positive_int(offload_config[offload_disk_size])
+            self._context_handle.set_offload_disk_size(
+                offload_config[offload_disk_size])
+        if enable_aio in offload_config:
+            Validator.check_bool(
+                offload_config[enable_aio], enable_aio, enable_aio)
+            self._context_handle.set_enable_aio(
+                offload_config[enable_aio])
+        if aio_block_size in offload_config:
+            Validator.check_positive_int(offload_config[aio_block_size])
+            self._context_handle.set_aio_block_size(
+                offload_config[aio_block_size])
+        if aio_queue_depth in offload_config:
+            Validator.check_positive_int(offload_config[aio_queue_depth])
+            self._context_handle.set_aio_queue_depth(
+                offload_config[aio_queue_depth])
+        if enable_pinned_mem in offload_config:
+            Validator.check_bool(
+                offload_config[enable_pinned_mem], enable_pinned_mem, enable_pinned_mem)
+            self._context_handle.set_enable_pinned_mem(
+                offload_config[enable_pinned_mem])
+
+    def offload_config(self):
+        """Get config of offload"""
+        self.check_context_handle()
+        offload_config = {
+            _OffloadConfig.ENABLE_OFFLOAD: self._context_handle.enable_offload(),
+            _OffloadConfig.OFFLOAD_PARAM: self._context_handle.offload_param(),
+            _OffloadConfig.OFFLOAD_PATH: self._context_handle.offload_path(),
+            _OffloadConfig.OFFLOAD_CHECKPOINT: self._context_handle.offload_checkpoint(),
+            _OffloadConfig.OFFLOAD_DDR_SIZE: self._context_handle.offload_ddr_size(),
+            _OffloadConfig.OFFLOAD_DISK_SIZE: self._context_handle.offload_disk_size(),
+            _OffloadConfig.ENABLE_AIO: self._context_handle.enable_aio(),
+            _OffloadConfig.AIO_BLOCK_SIZE: self._context_handle.aio_block_size(),
+            _OffloadConfig.AIO_QUEUE_DEPTH: self._context_handle.aio_queue_depth(),
+            _OffloadConfig.ENABLE_PINNED_MEM: self._context_handle.enable_pinned_mem()
+        }
+        return offload_config
+
+
+_OFFLOAD_CONTEXT = None
+
+
+def offload_context():
+    """Get offload_context. if it is not created, create a new one."""
+    global _OFFLOAD_CONTEXT
+    if _OFFLOAD_CONTEXT is None:
+        _OFFLOAD_CONTEXT = _OffloadContext()
+    return _OFFLOAD_CONTEXT
+
+
+@args_type_check(offload_config=dict)
+def _set_offload_context(offload_config):
+    offload_context().set_offload_config(offload_config)
+
+
+def _get_offload_context():
+    return offload_context().offload_config()