mindspore 2.0.0rc1__cp38-cp38-manylinux1_x86_64.whl → 2.2.0__cp38-cp38-manylinux1_x86_64.whl

mindspore/parallel/_auto_parallel_context.py

@@ -62,6 +62,7 @@ class _ParallelOptimizerConfig:
     """
     GRADIENT_ACCUMULATION_SHARD = "gradient_accumulation_shard"
     PARALLEL_OPTIMIZER_THRESHOLD = "parallel_optimizer_threshold"
+    OPTIMIZER_WEIGHT_SHARD_SIZE = "optimizer_weight_shard_size"
 
 
 class _AutoParallelContext:
@@ -176,7 +177,6 @@ class _AutoParallelContext:
         if comm_type == _ParallelFusionConfig.REDUCESCATTER:
             self._context_handle.set_reducescatter_fusion_threshold_mb(fusion_threshold)
 
-
     def fusion_threshold_mb(self):
         """Get all reduce threshold."""
         self.check_context_handle()
@@ -229,6 +229,22 @@ class _AutoParallelContext:
         self.check_context_handle()
         return self._context_handle.get_pipeline_stage_split_num()
 
+    def set_pipeline_segments(self, segments):
+        """Set the segments of the pipeline"""
+        if isinstance(segments, bool) or not isinstance(segments, int):
+            raise TypeError("For 'set_auto_parallel_context', the argument 'pipeline_segments' "
+                            "must be int, but got the type : {}.".format(type(segments)))
+        if segments < 1:
+            raise ValueError("For 'set_auto_parallel_context', the argument 'pipeline_segments' "
+                             "should be greater or equal 1, but got the value of segments : {}.".format(segments))
+        self.check_context_handle()
+        self._context_handle.set_pipeline_segment_split_num(segments)
+
+    def get_pipeline_segments(self):
+        """Get the stages of the pipeline"""
+        self.check_context_handle()
+        return self._context_handle.get_pipeline_segment_split_num()
+
     def set_gradients_mean(self, gradients_mean):
         """
         Set gradients_mean flag.
@@ -370,7 +386,7 @@ 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.
+        will search the desired strategies. Default: ``False``.
         This attribute is replaced by context.set_auto_parallel_context(search_mode="sharding_propagation").
 
         Args:
@@ -491,6 +507,9 @@ class _AutoParallelContext:
         Args:
             grad_accumulation_step (int): The grad accumulation step.
         """
+        if grad_accumulation_step > 1:
+            raise ValueError("The interface is deprecated. To use gradient accumulation, "
+                             "please use GradAccumulationCell in mindspore.nn.wrap.cell_wrapper.")
         self.check_context_handle()
         Validator.check_positive_int(grad_accumulation_step)
         self._context_handle.set_grad_accumulation_step(grad_accumulation_step)
@@ -758,6 +777,11 @@ class _AutoParallelContext:
                             .format(type(enable_parallel_optimizer)))
         self._context_handle.set_enable_parallel_optimizer(enable_parallel_optimizer)
 
+    def get_enable_fold_pipeline(self):
+        """Get parallel optimizer flag."""
+        self.check_context_handle()
+        return self._context_handle.get_enable_fold_pipeline()
+
     def get_enable_parallel_optimizer(self):
         """Get parallel optimizer flag."""
         self.check_context_handle()
@@ -767,8 +791,6 @@ class _AutoParallelContext:
         r"""
         Set the configure for parallel optimizer. The configure provides more detailed behavior control about parallel
         training when parallel optimizer is enabled.
-        Currently it supports the key `gradient_accumulation_shard`. The configure will be effective
-        when we use context.set_auto_parallel_context(enable_parallel_optimizer=True).
 
         Args:
             parallel_optimizer_config(dict): A dict contains the keys and values for setting the parallel optimizer
@@ -786,14 +808,21 @@ class _AutoParallelContext:
                                                  enabled, parameters with size smaller than this threshold will not be
                                                  sharded across the devices. Parameter size = shape[0] \* ... \*
                                                  shape[n] \* size(dtype). Non-negative. Unit: KB. Default: 64.
+            - optimizer_weight_shard_size(int): Set the optimizer weight shard group size if you want to specific the
+                                                maximum group size across devices when the parallel optimizer is
+                                                enabled. The numerical range can be (0, device_num]. Default value
+                                                is -1, which means the optimizer weight shard group size will
+                                                the data parallel group of each parameter. Default -1.
+
         """
         self.check_context_handle()
         grad_shard_name = _ParallelOptimizerConfig.GRADIENT_ACCUMULATION_SHARD
         threshold_name = _ParallelOptimizerConfig.PARALLEL_OPTIMIZER_THRESHOLD
+        optimizer_weight_shard_size_name = _ParallelOptimizerConfig.OPTIMIZER_WEIGHT_SHARD_SIZE
 
         for config_name in parallel_optimizer_config:
             unknown_config = []
-            if config_name not in [grad_shard_name, threshold_name]:
+            if config_name not in [grad_shard_name, threshold_name, optimizer_weight_shard_size_name]:
                 unknown_config.append(config_name)
 
             if unknown_config:
@@ -811,6 +840,11 @@ class _AutoParallelContext:
             self._context_handle.set_parallel_optimizer_threshold(
                 parallel_optimizer_config[threshold_name])
 
+        if optimizer_weight_shard_size_name in parallel_optimizer_config:
+            value = parallel_optimizer_config[optimizer_weight_shard_size_name]
+            Validator.check_positive_int(value)
+            self.set_optimizer_weight_shard_size(value)
+
     def get_grad_accumulation_shard(self):
         """Get grad accumulation shard."""
         self.check_context_handle()
@@ -824,7 +858,7 @@ class _AutoParallelContext:
     def set_enable_alltoall(self, enable_a2a):
         """
         Set the value of enabling AllToAll. If False, AllGather and Split are used to circumvent AllToAll.
-        Default: False.
+        Default: ``False``.
 
         Args:
             enable_a2a (bool): Enable/disable AllToAll.
@@ -890,6 +924,13 @@ class _AutoParallelContext:
         self.check_context_handle()
         return self._context_handle.get_optimizer_weight_shard_size()
 
+    def set_ops_strategy_json_config(self, type, path, mode):
+        """
+         Set configuration of saving ops strategy in file .json.
+        """
+        self.check_context_handle()
+        self._context_handle.set_ops_strategy_json_config(type, path, mode)
+
     def set_optimizer_weight_shard_aggregated_save(self, optimizer_weight_shard_aggregated_save):
         """
         Set optimizer_weight_shard_aggregated_save.
@@ -1027,7 +1068,28 @@ class _AutoParallelContext:
             self.set_enable_all_gather_fusion(openstate)
             self.set_enable_reduce_scatter_fusion(openstate)
 
+def _set_ops_strategy_json_config(type="SAVE", path="", mode="all"):
+    """
+    Set strategy json configuration.
 
+    Args:
+        type (str): The parameter for choosing save or load .json file.
+        path (str): Path to save or load parallel strategy json.
+        mode (str): The parameter for choosing save all or important operators.
+
+    Raises:
+        KeyError: When type is not 'SAVE' or 'LOAD'.
+        KeyError: When mode is not 'all' or 'principal'.
+    """
+    dir_path = os.path.dirname(path)
+    if dir_path and not os.path.exists(dir_path):
+        os.makedirs(dir_path)
+    check_type = ["SAVE", "LOAD"]
+    check_mode = ["all", "principal"]
+    if type in check_type and mode in check_mode:
+        auto_parallel_context().set_ops_strategy_json_config(type, path, mode)
+    else:
+        raise KeyError("Type must be 'SAVE' or 'LOAD' and mode must be 'all' or 'principal'")
 
 _AUTO_PARALLEL_CONTEXT = None
 
@@ -1052,6 +1114,7 @@ _set_auto_parallel_context_func_map = {
     "gradient_fp32_sync": auto_parallel_context().set_gradient_fp32_sync,
     "loss_repeated_mean": auto_parallel_context().set_loss_repeated_mean,
     "pipeline_stages": auto_parallel_context().set_pipeline_stages,
+    "pipeline_segments": auto_parallel_context().set_pipeline_segments,
     "parallel_mode": auto_parallel_context().set_parallel_mode,
     "search_mode": auto_parallel_context().set_strategy_search_mode,
     "auto_parallel_search_mode": auto_parallel_context().set_auto_parallel_search_mode,
@@ -1073,7 +1136,6 @@ _set_auto_parallel_context_func_map = {
     "strategy_ckpt_config": auto_parallel_context().set_strategy_ckpt_config,
     "comm_fusion": auto_parallel_context().set_comm_fusion}
 
-
 _get_auto_parallel_context_func_map = {
     "device_num": auto_parallel_context().get_device_num,
     "global_rank": auto_parallel_context().get_global_rank,
@@ -1110,7 +1172,6 @@ _get_auto_parallel_context_func_map = {
                  communi_parallel_mode=str, optimizer_weight_shard_size=int, sharding_propagation=bool,
                  optimizer_weight_shard_aggregated_save=bool, enable_alltoall=bool, comm_fusion=dict,
                  strategy_ckpt_config=dict)
-
 def _set_auto_parallel_context(**kwargs):
     """
     Set auto parallel context.
@@ -1121,11 +1182,11 @@ def _set_auto_parallel_context(**kwargs):
     Args:
         device_num (int): Available device number, the value must be in [1, 4096]. Default: 1.
         global_rank (int): Global rank id, the value must be in [0, 4095]. Default: 0.
-        gradients_mean (bool): Whether to perform mean operator after all-reduce of mirror. Default: False.
+        gradients_mean (bool): Whether to perform mean operator after all-reduce of mirror. Default: ``False``.
         loss_repeated_mean (bool): Whether to perform mean operator in backward in the case of repeated
-                        calculations. Default: True.
+                        calculations. Default: ``True``.
         gradient_fp32_sync (bool): Gradients allreduce by fp32 even though gradients is fp16 if this flag is True.
-                        Default: True.
+                        Default: ``True``.
         parallel_mode (str): There are five kinds of parallel modes, "stand_alone", "data_parallel",
                      "hybrid_parallel", "semi_auto_parallel" and "auto_parallel". Default: "stand_alone".
 
@@ -1151,13 +1212,13 @@ def _set_auto_parallel_context(**kwargs):
                      for forward compatibility, and this attribute will be deleted in a future MindSpore version.
         parameter_broadcast (bool): Indicating whether to broadcast parameters before training.
                        "stand_alone", "semi_auto_parallel" and "auto_parallel" do not support parameter
-                       broadcast. Default: False.
+                       broadcast. Default: ``False``.
         strategy_ckpt_load_file (str): The path to load parallel strategy checkpoint. Default: ''
         strategy_ckpt_save_file (str): The path to save parallel strategy checkpoint. Default: ''
         group_ckpt_save_file (str): The path to save parallel group checkpoint. Default: ''
-        full_batch (bool): Whether to load the whole batch on each device. Default: False.
+        full_batch (bool): Whether to load the whole batch on each device. Default: ``False``.
         dataset_strategy Union[str, tuple]: Dataset sharding strategy. Default: "data_parallel".
-        enable_parallel_optimizer (bool): Enable using optimizer segmentation or not. Default: False.
+        enable_parallel_optimizer (bool): Enable using optimizer segmentation or not. Default: ``False``.
         all_reduce_fusion_config (list): Set allreduce fusion strategy by parameters indices.
         pipeline_stages (int): Set the stage information for pipeline parallel. This indicates how
                         the devices are distributed alone the pipeline. The total devices will be divided into
@@ -1175,13 +1236,13 @@ def _set_auto_parallel_context(**kwargs):
                                     It should be larger than one and less than or equal with the data parallel size.
                                     Default: -1, which means fully use parallel optimizer in data parallel dimension.
         optimizer_weight_shard_aggregated_save (bool): Whether to integrated save weight shard when enable parallel
-                                                       optimizer. Default: False.
+                                                       optimizer. Default: ``False``.
         sharding_propagation (bool): 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.
+                                    search the desired strategies. Default: ``False``.
         enable_alltoall (bool): Set the value of enabling AllToAll. If False, AllGather and Split are used to
-                                circumvent AllToAll. Default: False.
+                                circumvent AllToAll. Default: ``False``.
         comm_fusion (dict): A dict contains the types and configurations for setting the communication fusion. each
                     communication fusion config has two keys: "mode" and "config".
                     It supports following communication fusion types and configurations:
@@ -1246,8 +1307,8 @@ def _reset_auto_parallel_context():
     - strategy_ckpt_load_file: ""
     - strategy_ckpt_save_file: ""
     - enable_parallel_optimizer: False
-    - search_mode: dynamic_programming
-    - auto_parallel_search_mode: dynamic_programming
+    - search_mode: 'recursive_programming
+    - auto_parallel_search_mode: 'recursive_programming
     - sharding_propagation: False
     - pipeline_stages: 0
     - gradient_accumulation_shard: True

mindspore/parallel/_cost_model_context.py

@@ -452,6 +452,31 @@ class _CostModelContext:
             raise ValueError("Context handle is none in context!!!")
         return self._context_handle.get_costmodel_allreduce_fusion_allreduce_bandwidth()
 
+    def set_rp_matmul_mem_coef(self, coef):
+        """
+        Set the matmul memory coef which is used in the RP algorithm.
+
+        Args:
+            coef (int): The coefficient of memory cost in RP.
+
+        Raises:
+            ValueError: If context handle is none.
+        """
+        if self._context_handle is None:
+            raise ValueError("Context handle is none in context!!!")
+        self._context_handle.set_rp_matmul_mem_coef(coef)
+
+    def get_rp_matmul_mem_coef(self):
+        """
+        Get the matmul memory coef which is used in the RP algorithm.
+
+        Raises:
+            ValueError: If context handle is none.
+        """
+        if self._context_handle is None:
+            raise ValueError("Context handle is none in context!!!")
+        return self._context_handle.get_rp_matmul_mem_coef()
+
     def set_costmodel_allreduce_fusion_computation_time_parameter(self, computation_time_parameter):
         """
         Set costmodel allreduce fusion computation time parameter.
@@ -656,3 +681,20 @@ def _get_algo_single_loop():
     Get the flag of whether or not generating a single suite of OperatorInfos in for-loop.
     """
     return cost_model_context().get_dp_algo_single_loop()
+
+
+def _set_rp_matmul_mem_coef(coef):
+    """
+    Set the matmul memory coef which is used in the RP algorithm.
+
+    Args:
+        coef (int): The coefficient of memory cost in RP.
+    """
+    cost_model_context().set_rp_matmul_mem_coef(coef)
+
+
+def _get_rp_matmul_mem_coef():
+    """
+    Get the matmul memory coef which is used in the RP algorithm.
+    """
+    return cost_model_context().get_rp_matmul_mem_coef()

mindspore/parallel/_offload_context.py

@@ -22,21 +22,27 @@ from mindspore._c_expression import OffloadContext
 from mindspore._checkparam import args_type_check
 from mindspore import _checkparam as Validator
 
+K_RE_PATTERN = r'[1-9][0-9]*(\.)?[0-9]*GB|0\.[0-9]*GB'
+K_GBTOBYTE = 1 << 30
+
 
 class _OffloadConfig:
     """
     The key of the Offload Config.
     """
-    ENABLE_OFFLOAD = "enable_offload"
     OFFLOAD_PARAM = "offload_param"
     OFFLOAD_PATH = "offload_path"
+    OFFLOAD_CPU_SIZE = "offload_cpu_size"
     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"
+    AUTO_OFFLOAD = "auto_offload"
+    CPU_RATIO = "cpu_ratio"
+    HBM_RATIO = "hbm_ratio"
+    HOST_MEM_BLOCk_SIZE = "host_mem_block_size"
 
 
 class _OffloadContext:
@@ -70,97 +76,164 @@ class _OffloadContext:
         if self._context_handle is None:
             raise ValueError("Context handle is none in context!!!")
 
+    def set_offload_param(self, offload_param):
+        """Set offload_param"""
+        if not isinstance(offload_param, str):
+            raise TypeError("For 'set_offload_param', "
+                            "the argument 'offload_param' must be str, but got the type : {}."
+                            .format(type(offload_param)))
+        Validator.check_string(offload_param.lower(), ["cpu", "disk"], "offload_param", "set_offload_param")
+        self._context_handle.set_offload_param(offload_param.lower())
+
+    def set_offload_checkpoint(self, offload_checkpoint):
+        """Set offload_checkpoint"""
+        if not isinstance(offload_checkpoint, str):
+            raise TypeError("For 'set_offload_checkpoint', "
+                            "the argument 'offload_checkpoint' must be str, but got the type : {}."
+                            .format(type(offload_checkpoint)))
+        Validator.check_string(offload_checkpoint.lower(), ["cpu", "disk"], "offload_checkpoint",
+                               "set_offload_checkpoint")
+        self._context_handle.set_offload_checkpoint(offload_checkpoint.lower())
+
+    def set_offload_path(self, offload_path):
+        """Set offload_path"""
+        if not isinstance(offload_path, str):
+            raise TypeError("For 'set_offload_path', "
+                            "the argument 'offload_path' must be str, but got the type : {}."
+                            .format(type(offload_path)))
+        self._context_handle.set_offload_path(offload_path)
+
+    def set_offload_cpu_size(self, offload_cpu_size):
+        """Set offload_cpu_size"""
+        if not isinstance(offload_cpu_size, str):
+            raise TypeError("For 'set_offload_cpu_size', "
+                            "the argument 'offload_cpu_size' must be str, but got the type : {}."
+                            .format(type(offload_cpu_size)))
+        if not Validator.check_str_by_regular(offload_cpu_size, K_RE_PATTERN):
+            raise ValueError("The argument 'offload_cpu_size' should be in correct "
+                             " format! It must be a string ending with 'GB', in addition to that, it must contain "
+                             "only numbers or decimal points, such as \"5GB\" or \"3.5GB\", but got {}."
+                             .format(offload_cpu_size))
+        ddr_size = float(offload_cpu_size[:-2])
+        self._context_handle.set_offload_cpu_size(int(ddr_size * K_GBTOBYTE))
+
+    def set_offload_disk_size(self, offload_disk_size):
+        """Set offload_disk_size"""
+        if not isinstance(offload_disk_size, str):
+            raise TypeError("For 'set_offload_disk_size', "
+                            "the argument 'offload_disk_size' must be str, but got the type : {}."
+                            .format(type(offload_disk_size)))
+        if not Validator.check_str_by_regular(offload_disk_size, K_RE_PATTERN):
+            raise ValueError("The argument 'offload_disk_size' should be in correct "
+                             " format! It must be a string ending with 'GB', in addition to that, it must contain "
+                             "only numbers or decimal points, such as \"5GB\" or \"3.5GB\", but got {}."
+                             .format(offload_disk_size))
+        disk_size = float(offload_disk_size[:-2])
+        self._context_handle.set_offload_disk_size(int(disk_size * K_GBTOBYTE))
+
+    def set_enable_aio(self, enable_aio):
+        """Set enable_aio"""
+        Validator.check_bool(enable_aio, "enable_aio", "set_enable_aio")
+        self._context_handle.set_enable_aio(enable_aio)
+
+    def set_aio_block_size(self, aio_block_size):
+        """Set aio_block_size"""
+        if not isinstance(aio_block_size, str):
+            raise TypeError("For 'set_aio_block_size', "
+                            "the argument 'aio_block_size' must be str, but got the type : {}."
+                            .format(type(aio_block_size)))
+        if not Validator.check_str_by_regular(aio_block_size, K_RE_PATTERN):
+            raise ValueError("The argument 'aio_block_size' should be in correct "
+                             " format! It must be a string ending with 'GB', in addition to that, it must contain "
+                             "only numbers or decimal points, such as \"5GB\" or \"3.5GB\", but got {}."
+                             .format(aio_block_size))
+        aio_size = float(aio_block_size[:-2])
+        self._context_handle.set_aio_block_size(int(aio_size * K_GBTOBYTE))
+
+    def set_aio_queue_depth(self, aio_queue_depth):
+        """Set aio_queue_depth"""
+        Validator.check_positive_int(
+            aio_queue_depth, "aio_queue_depth", "set_aio_queue_depth")
+        self._context_handle.set_aio_queue_depth(aio_queue_depth)
+
+    def set_enable_pinned_mem(self, enable_pinned_mem):
+        """Set enable_pinned_mem"""
+        Validator.check_bool(
+            enable_pinned_mem, "enable_pinned_mem", "set_enable_pinned_mem")
+        self._context_handle.set_enable_pinned_mem(enable_pinned_mem)
+
+    def set_auto_offload(self, auto_offload):
+        """Set auto_offload"""
+        Validator.check_bool(auto_offload, "auto_offload", "set_auto_offload")
+        self._context_handle.set_auto_offload(auto_offload)
+
+    def set_host_mem_block_size(self, host_mem_block_size):
+        """Set host_mem_block_size"""
+        if not isinstance(host_mem_block_size, str):
+            raise TypeError("For 'set_host_mem_block_size', "
+                            "the argument 'host_mem_block_size' must be str, but got the type : {}."
+                            .format(type(host_mem_block_size)))
+        if not Validator.check_str_by_regular(host_mem_block_size, K_RE_PATTERN):
+            raise ValueError("The argument 'host_mem_block_size' should be in correct "
+                             " format! It must be a string ending with 'GB', in addition to that, it must contain "
+                             "only numbers or decimal points, such as \"5GB\" or \"3.5GB\", but got {}."
+                             .format(host_mem_block_size))
+        block_size = float(host_mem_block_size[:-2])
+        self._context_handle.set_host_mem_block_size(
+            int(block_size * K_GBTOBYTE))
+
+    def set_cpu_ratio(self, cpu_ratio):
+        """Set cpu_ratio"""
+        Validator.check_float_range(
+            cpu_ratio, 0, 1, Validator.INC_RIGHT, 'cpu_ratio')
+        self._context_handle.set_cpu_ratio(cpu_ratio)
+
+    def set_hbm_ratio(self, hbm_ratio):
+        """Set hbm_ratio"""
+        Validator.check_float_range(
+            hbm_ratio, 0, 1, Validator.INC_RIGHT, 'hbm_ratio')
+        self._context_handle.set_hbm_ratio(hbm_ratio)
+
     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]:
+            if config_name not in [_OffloadConfig.OFFLOAD_PARAM, _OffloadConfig.OFFLOAD_PATH,
+                                   _OffloadConfig.CPU_RATIO, _OffloadConfig.HOST_MEM_BLOCk_SIZE,
+                                   _OffloadConfig.HBM_RATIO, _OffloadConfig.OFFLOAD_CPU_SIZE,
+                                   _OffloadConfig.OFFLOAD_DISK_SIZE, _OffloadConfig.ENABLE_AIO,
+                                   _OffloadConfig.AIO_BLOCK_SIZE, _OffloadConfig.AIO_QUEUE_DEPTH,
+                                   _OffloadConfig.ENABLE_PINNED_MEM, _OffloadConfig.AUTO_OFFLOAD,
+                                   _OffloadConfig.OFFLOAD_CHECKPOINT]:
                 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])
+            func = _set_offload_context_func_map.get(config_name, None)
+            if not func:
+                raise ValueError(
+                    "Can not find set function: {}".format(config_name))
+            func(offload_config[config_name])
 
     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_CPU_SIZE: self._context_handle.offload_cpu_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()
+            _OffloadConfig.ENABLE_PINNED_MEM: self._context_handle.enable_pinned_mem(),
+            _OffloadConfig.AUTO_OFFLOAD: self._context_handle.auto_offload(),
+            _OffloadConfig.HOST_MEM_BLOCk_SIZE: self._context_handle.host_mem_block_size(),
+            _OffloadConfig.CPU_RATIO: self._context_handle.cpu_ratio(),
+            _OffloadConfig.HBM_RATIO: self._context_handle.hbm_ratio(),
+            _OffloadConfig.OFFLOAD_CHECKPOINT: self._context_handle.offload_checkpoint()
         }
         return offload_config
 
@@ -183,3 +256,20 @@ def _set_offload_context(offload_config):
 
 def _get_offload_context():
     return offload_context().offload_config()
+
+
+_set_offload_context_func_map = {
+    _OffloadConfig.OFFLOAD_PARAM: offload_context().set_offload_param,
+    _OffloadConfig.OFFLOAD_PATH: offload_context().set_offload_path,
+    _OffloadConfig.OFFLOAD_CPU_SIZE: offload_context().set_offload_cpu_size,
+    _OffloadConfig.OFFLOAD_DISK_SIZE: offload_context().set_offload_disk_size,
+    _OffloadConfig.ENABLE_AIO: offload_context().set_enable_aio,
+    _OffloadConfig.AIO_BLOCK_SIZE: offload_context().set_aio_block_size,
+    _OffloadConfig.AIO_QUEUE_DEPTH: offload_context().set_aio_queue_depth,
+    _OffloadConfig.ENABLE_PINNED_MEM: offload_context().set_enable_pinned_mem,
+    _OffloadConfig.AUTO_OFFLOAD: offload_context().set_auto_offload,
+    _OffloadConfig.HOST_MEM_BLOCk_SIZE: offload_context().set_host_mem_block_size,
+    _OffloadConfig.CPU_RATIO: offload_context().set_cpu_ratio,
+    _OffloadConfig.HBM_RATIO: offload_context().set_hbm_ratio,
+    _OffloadConfig.OFFLOAD_CHECKPOINT: offload_context().set_offload_checkpoint
+}

mindspore/parallel/_parallel_serialization.py

@@ -330,8 +330,8 @@ def _rank_list_for_transform_parallel_checkpoint(rank_id, src_strategy_list, dst
         device_list = list(range(0, np.prod(from_tensor_layout[0])))
         param_rank_list = _get_needed_rank_list_by_layouts(from_tensor_layout, to_tensor_layout, device_list, rank_id)
         param_rank_list_new = [rank % from_device_num for rank in param_rank_list]
-        param_rank_list_new = set(param_rank_list_new)
-        result_list.update(param_rank_list_new)
+        param_rank_set_new = set(param_rank_list_new)
+        result_list.update(param_rank_set_new)
     return list(result_list)
 
 

mindspore/parallel/_ps_context.py

@@ -114,10 +114,10 @@ def _set_ps_context(**kwargs):
     Args:
         enable_ps (bool): Whether to enable parameter server training mode.
                           Only after enable_ps is set True, the environment variables will be effective.
-                          Default: False.
+                          Default: ``False``.
         config_file_path (string): Configuration file path used by recovery. Default: ''.
         scheduler_manage_port (int): scheduler manage port used to scale out/in. Default: 11202.
-        enable_ssl (bool): Set PS SSL mode enabled or disabled. Default: False.
+        enable_ssl (bool): Set PS SSL mode enabled or disabled. Default: ``False``.
         client_password (str): Password to decrypt the secret key stored in the client certificate. Default: ''.
         server_password (str): Password to decrypt the secret key stored in the server certificate. Default: ''.
 
@@ -180,8 +180,8 @@ def _insert_hash_table_size(name, cache_vocab_size, embedding_size, vocab_size,
     ps_context().insert_hash_table_size(name, cache_vocab_size, embedding_size, vocab_size, param_key)
 
 
-def _reinsert_hash_table_size(new_name, cur_name, cache_vocab_size, embedding_size):
-    ps_context().reinsert_hash_table_size(new_name, cur_name, cache_vocab_size, embedding_size)
+def _reinsert_hash_table_size(new_name, cur_name):
+    ps_context().reinsert_hash_table_size(new_name, cur_name)
 
 
 def _insert_accumu_init_info(name, init_val):
@@ -228,3 +228,15 @@ def _enable_distributed_mindrt():
     This method is used to distinguish from old distributed training mode.
     '''
     return ps_context().enable_distributed_mindrt()
+
+
+def _set_checkpoint_load_status(status):
+    return ps_context().set_checkpoint_load_status(status)
+
+
+def _store_warm_up_ptr_by_tensor(param_key, tensor):
+    return ps_context().store_warm_up_ptr_by_tensor(param_key, tensor)
+
+
+def _store_warm_up_ptr_by_tensor_list(param_key, key_tensor, value_tensor, status_tensor):
+    return ps_context().store_warm_up_ptr_by_tensor_list(param_key, key_tensor, value_tensor, status_tensor)

mindspore/parallel/_recovery_context.py

@@ -65,7 +65,8 @@ def _set_recovery_context(**kwargs):
 
     Args:
         ckpt_path (string): Set the recovery path used to save checkpoint. Default: ''.
-        need_reset (bool): Set whether should call reset minddata and load ckpt for disaster recovery. Default: False.
+        need_reset (bool): Set whether should call reset minddata and load ckpt for disaster recovery.
+            Default: ``False``.
 
     Raises:
         ValueError: If input key is not the attribute in recovery context.