mindspore 2.4.10__cp39-cp39-win_amd64.whl → 2.6.0__cp39-cp39-win_amd64.whl

mindspore/train/_utils.py

@@ -16,22 +16,23 @@
 from __future__ import absolute_import
 
 import os
-import threading
-from datetime import datetime
+import sys
 import json
 from collections.abc import Iterable
 
+import time
 import numpy as np
 
 from mindspore.common.tensor import Tensor
-from mindspore._c_expression import Tensor as Tensor_
+from mindspore._c_expression import TensorPy as Tensor_
+from mindspore._c_expression import MSContext, ms_ctx_param
 from mindspore.common.dtype import dtype_to_nptype, pytype_to_dtype
 from mindspore.common import dtype as mstype
 from mindspore import context
 from mindspore import log as logger
 from mindspore import _checkparam as Validator
 from mindspore.common.api import _cell_graph_executor
-from mindspore.communication import get_group_size
+from mindspore.communication.management import get_rank, get_group_size
 from mindspore.train.mind_ir_pb2 import ModelProto as mindir_model
 from mindspore.train.checkpoint_pb2 import Checkpoint
 from mindspore.train.node_strategy_pb2 import ParallelStrategyMap as ckpt_strategy
@@ -65,6 +66,11 @@ def _get_types_and_shapes(dataset):
     return dataset_types, dataset_shapes
 
 
+def enable_data_broadcast():
+    """Get status to indicate if enable dataset broadcast."""
+    return MSContext.get_instance().get_param(ms_ctx_param.dataset_broadcast_opt_level) > 0
+
+
 def _exec_datagraph(exec_dataset, dataset_size, phase='dataset', create_data_info_queue=False):
     """Initialize and execute the dataset graph."""
     batch_size = exec_dataset.get_batch_size()
@@ -77,15 +83,12 @@ def _exec_datagraph(exec_dataset, dataset_size, phase='dataset', create_data_inf
     if queue_name is None:
         queue_name = str("")
 
+    # Don't enable dynamic shape(multi-subgraph) feature in pp/data_broadcast mode,
+    # otherwise get_data_info will stuck since some rank do not consume data.
     use_pipeline_parallel = (context.get_auto_parallel_context("pipeline_stages") > 1)
+    data_broadcast = enable_data_broadcast()
 
-    # temp env to disable dynamic feature of sink size 1
-    dynamic_sink1_env = os.getenv("MS_DEV_DYNAMIC_SINK1", None)
-    dynamic_sink1 = True
-    if dynamic_sink1_env and dynamic_sink1_env.strip() in ['False', 'false']:
-        dynamic_sink1 = False
-
-    if use_pipeline_parallel or not dynamic_sink1:
+    if use_pipeline_parallel or data_broadcast:
         create_data_info_queue = False
 
     exec_dataset = exec_dataset.device_que(send_epoch_end=send_epoch_end,
@@ -320,9 +323,15 @@ def parse_strategy_ckpt(file_name):
 def _get_strategy_opt_shard(param_redundancy_dict, parameter_layout_opt_shard):
     """Strategy ckpt append opt shard."""
     for key, value in parameter_layout_opt_shard.items():
-        if value[1] not in (-1, 0):
-            opt_para_num = value[1]
+        if value[1] != 0:
             param_redundancy_ranks = param_redundancy_dict.get(key)
+            if value[1] != -1:
+                opt_para_num = value[1]
+            elif param_redundancy_ranks:
+                opt_para_num = len(param_redundancy_ranks) * len(param_redundancy_ranks[0]) // value[0]
+            else:
+                raise ValueError(f"For get_parameter_redundancy, the format of the parallel communication domain for "
+                                 f"the optimizer is incorrect.")
             res = []
             for param_ranks in param_redundancy_ranks:
                 if len(param_ranks) % opt_para_num == 0:
@@ -374,20 +383,40 @@ def _get_parameter_redundancy_without_opt_shard(parameter_layout, param_redundan
         param_redundancy_dict[key] = tuple(redundancy_list)
 
 
-def get_parameter_redundancy(layout_obj, initial_rank=0):
+def _get_initial_rank(parameter_layout):
+    """Get the initial rank of pp."""
+    for k, _ in parameter_layout.items():
+        dev_matrix = parameter_layout[k][0]
+        break
+    dev_num = 1
+    if dev_matrix:
+        for i in dev_matrix:
+            dev_num *= i
+    rank_id = get_rank()
+    initial_rank = (rank_id // dev_num) * dev_num
+    return initial_rank
+
+
+def _get_pp_size_from_redundancy_map(param_redundancy):
+    """Get pp size from redundancy map."""
+    for _, v in param_redundancy.items():
+        return len(v) * len(v[0])
+
+
+def get_parameter_redundancy(layout_obj, initial_rank=None):
     """
     Get parameter redundancy map.
 
     Args:
         layout_obj (Union[str, layout): File name of `strategy.ckpt` or net.parameter_layout_dict.
-        initial_rank (int): Start rank id for each pipeline. Default: 0.
+        initial_rank (int): Start rank id for each pipeline. Default: ``None``.
 
     Returns:
         Dict, dict of parameter redundancy info.
 
     Examples:
         >>> from mindspore.train.utils import get_parameter_redundancy
-        >>> param_redundancy_dict = get_parameter_redundancy("/path/to/strategy.ckpt")
+        >>> param_redundancy_dict = get_parameter_redundancy("/path/to/strategy.ckpt", initial_rank=0)
         {'param1': ((0, 1, 2, 3, 4, 5, 6, 7),),
          'param2': ((0, 4, 8, 12), (1, 5, 9, 13), (2, 6, 10, 14), (3, 7, 11, 15)),
          'param3': ((0, 4, 8, 12), (1, 5, 9, 13), (2, 6, 10, 14), (3, 7, 11, 15)),
@@ -404,7 +433,8 @@ def get_parameter_redundancy(layout_obj, initial_rank=0):
         from mindspore.communication.management import get_process_group_ranks
         groups_ranks = (tuple(get_process_group_ranks()),)
         param_redundancy_dict = {param.name: groups_ranks for _, param in layout_obj.parameters_and_names()}
-        return param_redundancy_dict
+        sorted_param_redundancy_dict = {key: param_redundancy_dict[key] for key in sorted(param_redundancy_dict.keys())}
+        return sorted_param_redundancy_dict
     else:
         parameter_layout = {}
         for k, v in layout_obj.items():
@@ -412,6 +442,9 @@ def get_parameter_redundancy(layout_obj, initial_rank=0):
 
     param_redundancy_dict = {}
 
+    if initial_rank is None:
+        initial_rank = _get_initial_rank(parameter_layout)
+
     _get_parameter_redundancy_without_opt_shard(parameter_layout, param_redundancy_dict, initial_rank)
 
     if isinstance(layout_obj, str):
@@ -419,7 +452,8 @@ def get_parameter_redundancy(layout_obj, initial_rank=0):
     else:
         _get_layout_opt_shard(layout_obj, param_redundancy_dict)
 
-    return param_redundancy_dict
+    sorted_param_redundancy_dict = {key: param_redundancy_dict[key] for key in sorted(param_redundancy_dict.keys())}
+    return sorted_param_redundancy_dict
 
 
 def _collect_settings_by_rank(redundancy_map):
@@ -514,12 +548,47 @@ def parse_hccl_file(hccl_file_path):
     return rankid_dict
 
 
-def vlog_print(level, module, file, line, message):
-    '''Read environment variable VLOG_v and print to log'''
-    if os.environ.get("VLOG_v") == level:
-        now = datetime.now()
-        formatted_time = now.strftime("%Y-%m-%d-%H:%M:%S.%f")[:-3] + f".{now.microsecond // 1000}"
-        path = 'mindspore' + file.split("mindspore")[-1]
-        pid = os.getpid()
-        thread_id = threading.get_ident()
-        print(f"[V{level}] {module}({pid},{thread_id},python):{formatted_time} [{path}:{line}] {message}", flush=True)
+def _progress_bar(iterable, total=None):
+    """
+    Decorate an iterable object, returning an iterator which acts exactly
+    like the original iterable, but prints a dynamically updating
+    progressbar every time a value is requested.
+    """
+    if total is None:
+        total = len(iterable)
+
+    start_time = time.time()
+
+    def print_progress_bar(iteration):
+        percent = f"{100 * (iteration / float(total)):.1f}"
+        bar_length = 40
+        filled_length = int(bar_length * iteration // total)
+        bar = '█' * filled_length + '-' * (bar_length - filled_length)
+
+        elapsed_time = time.time() - start_time
+        estimated_total_time = elapsed_time / iteration * total
+        remaining_time = estimated_total_time - elapsed_time
+
+        elapsed_time_str = time.strftime("%H:%M:%S", time.gmtime(elapsed_time))
+        remaining_time_str = time.strftime("%H:%M:%S", time.gmtime(remaining_time))
+
+        sys.stdout.reconfigure(encoding="utf-8")
+        print(f'\r{percent}%|{bar}|[{elapsed_time_str}<{remaining_time_str}]', end='')
+        if iteration == total:
+            print()
+
+    for i, item in enumerate(iterable, start=1):
+        yield item
+        print_progress_bar(i)
+
+
+def _load_and_transform(path, name_map, load_func, transform_func):
+    if load_func is not None:
+        param_dict = load_func(path)
+    else:
+        param_dict = path
+    transform_dict = {}
+    for k, v in param_dict.items():
+        new_name = name_map.get(k, k) if name_map is not None else k
+        transform_dict[new_name] = transform_func(v, new_name)
+    return transform_dict

mindspore/train/amp.py

@@ -101,6 +101,7 @@ AMP_AUTO_BLACK_LIST = [
     P.LayerNorm,
     gen.LayerNormExt,
     P.BatchNorm,
+    gen.BatchNormExt,
     gen.GroupNorm,
     P.KLDivLoss,
     P.SmoothL1Loss,
@@ -112,6 +113,7 @@ AMP_AUTO_BLACK_LIST = [
     P.Pdist,
     P.Cdist,
     P.Renorm,
+    gen.MSELossExt,
 ]
 
 # Indicates which inputs of primitives need to be converted
@@ -428,15 +430,15 @@ def auto_mixed_precision(network, amp_level="O0", dtype=mstype.float16):
 
     ``Pow``, ``ACos``, ``Asin``, ``Cosh``, ``Erfinv``, ``Exp``, ``Expm1``, ``Log``, ``Log1p``, ``Reciprocal``,
     ``Rsqrt``, ``Sinh``, ``Tan``, ``Softplus``, ``SoftplusExt``, ``LayerNorm``, ``LayerNormExt``, ``BatchNorm``,
-    ``GroupNorm``, ``KLDivLoss``, ``SmoothL1Loss``, ``MultilabelMarginLoss``, ``SoftMarginLoss``,
+    ``BatchNormExt``, ``GroupNorm``, ``KLDivLoss``, ``SmoothL1Loss``, ``MultilabelMarginLoss``, ``SoftMarginLoss``,
     ``TripletMarginLoss``, ``MultiMarginLoss``, ``BCEWithLogitsLoss``, ``Pdist``, ``Cdist``, ``Renorm``,
     ``ReduceProd``, ``Softmax``, ``LogSoftmax``, ``CumProd``, ``CumSum``, ``CumsumExt``, ``ProdExt``, ``SumExt``,
-    ``Norm``
+    ``Norm``, ``MSELossExt``
 
     Operators in `promote_list` are:
 
     ``Addcdiv``, ``Addcmul``, ``Cross``, ``_PyboostCrossPrim``, ``Dot``, ``GridSampler2D``, ``GridSampler3D``,
-    ``BiasAdd``
+    ``BiasAdd``, ``AddN``, ``Concat``
 
     For details on automatic mixed precision, refer to
     `Automatic Mix Precision <https://www.mindspore.cn/tutorials/en/master/beginner/mixed_precision.html>`_ .
@@ -636,7 +638,7 @@ def _add_loss_network(network, loss_fn, cast_model_type):
 
 
 def _is_grad_accumulation(mcell):
-    if mcell.cls_name == "GradAccumulationCell":
+    if mcell.cls_name == "GradAccumulationCell" or mcell.cls_name == "GradAccumulation":
         return True
     for cell in mcell.cells():
         if _is_grad_accumulation(cell):
@@ -837,12 +839,14 @@ def custom_mixed_precision(network, *, white_list=None, black_list=None, dtype=m
         - Repeatedly calling mixed-precision interfaces, such as `custom_mixed_precision` and `auto_mixed_precision`,
           can result in a larger network hierarchy and slower performance.
         - If interfaces like `Model` and `build_train_network` is used to train the network which is converted by
-          mixed-precision interfaces such as `custom_mixed_precision` and `auto_mixed_precision`, `amp_level`
+          mixed-precision interfaces such as `custom_mixed_precision` and `auto_mixed_precision`, `amp_level` or `level`
           need to be configured to ``O0`` to avoid the duplicated accuracy conversion.
         - Primitives for blacklist is not support yet.
 
     Args:
         network (Cell): Definition of the network.
+
+    Keyword Args:
         white_list (list[Primitive, Cell], optional): White list of custom mixed precision. Defaults: ``None`` , means
             white list is not used.
         black_list (list[Cell], optional): Black list of custom mixed precision. Defaults: ``None`` , means

mindspore/train/callback/__init__.py

@@ -36,9 +36,9 @@ from mindspore.train.callback._reduce_lr_on_plateau import ReduceLROnPlateau
 from mindspore.train.callback._on_request_exit import OnRequestExit
 from mindspore.train.callback._backup_and_restore import BackupAndRestore
 from mindspore.train.callback._flops_collector import FlopsUtilizationCollector
-from mindspore.train.callback._tft_register import TFTRegister
+from mindspore.train.callback._train_fault_tolerance import TrainFaultTolerance
 
 __all__ = ["Callback", "LossMonitor", "TimeMonitor", "ModelCheckpoint", "FlopsUtilizationCollector",
            "SummaryCollector", "CheckpointConfig", "RunContext", "LearningRateScheduler", "SummaryLandscape",
            "History", "LambdaCallback", "ReduceLROnPlateau", "EarlyStopping", "OnRequestExit", "BackupAndRestore",
-           "TFTRegister"]
+           "TrainFaultTolerance"]

mindspore/train/callback/_callback.py

@@ -121,10 +121,7 @@ class Callback:
     When creating a custom Callback, model context information can be obtained in Callback
     methods by calling `RunContext.original_args()`, which is a dictionary varivable
     recording current attributes. Users can add custimized attributes to the information.
-    Training process can also be stopped by calling `request_stop` method. For details
-    of custom Callback, please check
-    `Callback tutorial <https://www.mindspore.cn/docs/en/master/model_train/train_process/model/
-    callback.html#customized-callback-mechanism>`_.
+    Training process can also be stopped by calling `request_stop` method.
 
     Examples:
         >>> import numpy as np
@@ -491,9 +488,7 @@ class RunContext:
 
     Callback objects not only can obtain the Model context information by calling by
     `RunContext.original_args()` and add extra attributes to the information, but also can stop the
-    training process by calling `request_stop` method. For details of custom Callback,
-    please check
-    `Callback Mechanism <https://www.mindspore.cn/docs/en/master/model_train/train_process/model/callback.html>`_.
+    training process by calling `request_stop` method.
 
     `RunContext.original_args()` holds the model context information as a dictionary variable, and
     different attributes of the dictionary are stored in training or eval process. Details are as follows:
@@ -572,10 +567,6 @@ class RunContext:
 
         Returns:
            Dict, an object that holds the original arguments of model.
-
-        Tutorial Examples:
-            - `Callback Mechanism - Customized Callback Mechanism
-              <https://mindspore.cn/docs/en/master/model_train/train_process/model/callback.html#customized-callback-mechanism>`_
         """
         return self._original_args
 
@@ -585,11 +576,6 @@ class RunContext:
 
         Callbacks can use this function to request stop of iterations.
         model.train() checks whether this is called or not.
-
-        Tutorial Examples:
-            - `Callback Mechanism - Customized Training Termination Time
-              <https://mindspore.cn/docs/en/master/model_train/train_process/model/callback.html#
-              customized-training-termination-time>`_
         """
         self._stop_requested = True
 

mindspore/train/callback/_checkpoint.py

@@ -18,25 +18,22 @@ from __future__ import absolute_import
 import os
 import stat
 import time
-import threading
 
 import mindspore.context as context
 from mindspore import log as logger
 from mindspore import nn
 from mindspore import _checkparam as Validator
 from mindspore.train._utils import _make_directory
-from mindspore.train.serialization import save_checkpoint, _save_graph
+from mindspore.train.serialization import save_checkpoint, _save_graph, _wait_async_process_save_ckpt, \
+    _wait_async_thread_save_ckpt, _check_async_save
 from mindspore.parallel._cell_wrapper import destroy_allgather_cell
 from mindspore.parallel._recovery_context import _set_recovery_context, _get_recovery_context
-from mindspore.parallel._auto_parallel_context import _get_auto_parallel_context
-from mindspore.parallel._utils import _get_device_num
-from mindspore.communication.management import get_rank
-from mindspore.train._utils import get_parameter_redundancy, remove_param_redundancy
-from mindspore.train.callback._callback import Callback, set_cur_net
+from mindspore.communication.management import get_rank, get_group_size
+from mindspore.train._utils import get_parameter_redundancy, remove_param_redundancy, _get_pp_size_from_redundancy_map
+from mindspore.train.callback._callback import Callback
 from mindspore.common.tensor import Tensor
 from mindspore.common.parameter import Parameter
 from mindspore.common.generator import Generator
-from mindspore.common.api import _cell_graph_executor
 from mindspore._c_expression import collect_host_info, get_clock_syscnt
 
 _cur_dir = os.getcwd()
@@ -44,15 +41,6 @@ SAVE_DIR = _cur_dir
 _info_list = ["epoch_num", "step_num"]
 
 
-def _wait_async_save_ckpt(async_save=False):
-    """Waiting for asynchronous saving of ckpt to complete."""
-    if async_save:
-        thread_list = threading.enumerate()
-        for thread in thread_list:
-            if thread.getName() == "asyn_save_ckpt":
-                thread.join()
-
-
 def _get_dp_tp_from_redundancy(redundancy_tuple):
     """From redundancy get dp and tp"""
     dp = []
@@ -76,6 +64,15 @@ def _get_dp_tp_from_layout(parameter_redundancy_dict):
     return dp, tp
 
 
+def _wait_async_save_ckpt(async_save=False):
+    """Waiting for asynchronous saving of ckpt to complete."""
+    if async_save:
+        if async_save == "process":
+            _wait_async_process_save_ckpt()
+        else:
+            _wait_async_thread_save_ckpt()
+
+
 def _chg_ckpt_file_name_if_same_exist(directory, prefix, exception=False):
     """Check if there is a file with the same name."""
     if callable(prefix) or callable(directory):
@@ -87,7 +84,7 @@ def _chg_ckpt_file_name_if_same_exist(directory, prefix, exception=False):
         name_ext = os.path.splitext(filename)
         if exception and filename[-16:] != "_breakpoint.ckpt":
             continue
-        if not exception and (name_ext[-1] != ".ckpt" or filename[-16:] == "_breakpoint.ckpt"):
+        if not exception and (name_ext[-1] not in (".ckpt", ".safetensors") or filename[-16:] == "_breakpoint.ckpt"):
             continue
         # find same prefix file
         if filename.find(prefix) == 0 and not filename[pre_len].isalpha():
@@ -106,10 +103,10 @@ def _chg_ckpt_file_name_if_same_exist(directory, prefix, exception=False):
     return prefix
 
 
-def _check_format_and_other_params(format, enc_key, enc_mode, crc_check=False, async_save=False, exception_save=False,
+def _check_format_and_other_params(format, enc_key, enc_mode, crc_check=False, exception_save=False,
                                    map_param_inc=False, global_step_num=None):
-    param_not_default = (enc_key is not None or enc_mode != "AES-GCM" or crc_check or async_save
-                         or exception_save or map_param_inc or global_step_num is not None)
+    param_not_default = (enc_key is not None or enc_mode != "AES-GCM" or crc_check or exception_save or map_param_inc
+                         or global_step_num is not None)
     if format == "safetensors" and param_not_default:
         raise ValueError("For 'save_checkpoint', when format is 'safetensors', other param must be default.")
 
@@ -139,7 +136,10 @@ class CheckpointConfig:
         integrated_save (bool): Whether to merge and save the split Tensor in the automatic parallel scenario.
             Integrated save function is only supported in automatic parallel scene, not supported
             in manual parallel. Default: ``True`` .
-        async_save (bool): Whether asynchronous execution saves the checkpoint to a file. Default: ``False`` .
+        async_save (Union[bool, str], optional):Whether to use asynchronous saving of the checkpoint file or
+                                    safetensors file, if True, the asynchronous thread is used by default. If the type
+                                    is string, the method of asynchronous saving, it can be "process" or "thread".
+                                    Default: ``False`` .
         saved_network (Cell): Network to be saved in checkpoint file. If the saved_network has no relation
             with the network in training, the initial value of saved_network will be saved. Default: ``None`` .
         append_info (list): The information save to checkpoint file. Support "epoch_num", "step_num" and
@@ -247,7 +247,7 @@ class CheckpointConfig:
                 self._keep_checkpoint_max = 1
 
         self._integrated_save = Validator.check_bool(integrated_save)
-        self._async_save = Validator.check_bool(async_save)
+        self._async_save = _check_async_save(async_save)
         self._saved_network = saved_network
         self._append_dict = self._handle_append_info(append_info)
         self._enc_key = Validator.check_isinstance('enc_key', enc_key, (type(None), bytes))
@@ -258,8 +258,7 @@ class CheckpointConfig:
         self.enable_redundance = kwargs.get('enable_redundance', False)
         self.remove_redundancy = Validator.check_isinstance('remove_redundancy', remove_redundancy, bool)
 
-        _check_format_and_other_params(format, enc_key, enc_mode, crc_check, async_save, exception_save,
-                                       self._map_param_inc)
+        _check_format_and_other_params(format, enc_key, enc_mode, crc_check, exception_save, self._map_param_inc)
 
     @property
     def save_checkpoint_steps(self):
@@ -313,10 +312,10 @@ class CheckpointConfig:
     @property
     def async_save(self):
         """
-        Get the value of whether asynchronous execution saves the checkpoint to a file.
+        Get the value of whether or how asynchronous execution saves the checkpoint to a file.
 
         Returns:
-            bool, whether asynchronous execution saves the checkpoint to a file.
+            (bool, str), whether or how asynchronous execution saves the checkpoint to a file.
         """
         return self._async_save
 
@@ -449,8 +448,9 @@ class ModelCheckpoint(Callback):
     Note:
         In the distributed training scenario, please specify different directories for each training process
         to save the checkpoint file. Otherwise, the training may fail.
-        If this callback is used in the `model` function, the checkpoint file will saved
-        parameters of the optimizer by default.
+        If this callback is used in the
+        `Model <https://www.mindspore.cn/docs/en/master/api_python/train/mindspore.train.Model.html>`_ function,
+        the checkpoint file will saved parameters of the optimizer by default.
 
     Args:
         prefix (Union[str, callable object]): The prefix name or callable object to generate name of checkpoint files.
@@ -511,7 +511,7 @@ class ModelCheckpoint(Callback):
         if callable(prefix):
             self._prefix_func = prefix
 
-        if _get_recovery_context("enable_recovery"):
+        if context.get_context("device_target") == "GPU" and _get_recovery_context("enable_recovery"):
             _set_recovery_context(ckpt_path=self._directory)
 
         if config is None:
@@ -538,6 +538,8 @@ class ModelCheckpoint(Callback):
         self._graph_saved = False
         self._need_flush_from_cache = True
         self._map_param_inc = self._config.map_param_inc
+        self._d2h_async = os.environ.get("MS_ENABLE_CKPT_D2H_ASYNC") == "1"
+        self._run_mode = context.get_context("mode")
 
     def step_end(self, run_context):
         """
@@ -551,19 +553,17 @@ class ModelCheckpoint(Callback):
             from aiturbo.checkpoint import aiturbo_mindspore as aiturbo
             ckpt_storage_path = self._directory
             rank_id = get_rank()
-            stage_num = _get_auto_parallel_context("pipeline_stages")
-            stage_rank_num = _get_device_num() // stage_num
+            device_num = get_group_size()
             param_layout = cb_params.train_network.parameter_layout_dict
             if not param_layout:
-                layout = {"stage_num": stage_num, "stage_rank_num": stage_rank_num, "stage_layout": None}
+                layout = {"stage_num": 1, "stage_rank_num": device_num, "stage_layout": None}
                 aiturbo.init(ckpt_storage_path, rank_id, layout, None, False, None)
             else:
-                device_num = _get_device_num()
-                chunk_size = device_num // stage_num
-                initial_rank = (rank_id // chunk_size) * chunk_size
-                param_redundancy_dict = get_parameter_redundancy(param_layout, initial_rank)
+                param_redundancy_dict = get_parameter_redundancy(param_layout)
+                pp_size = _get_pp_size_from_redundancy_map(param_redundancy_dict)
+                stage_num = device_num // pp_size
                 dp, _ = _get_dp_tp_from_layout(param_redundancy_dict)
-                layout = {"stage_num": stage_num, "stage_rank_num": stage_rank_num,
+                layout = {"stage_num": stage_num, "stage_rank_num": pp_size,
                           "stage_layout": param_redundancy_dict}
                 single_params = remove_param_redundancy(param_redundancy_dict)
                 single_params = {device_id: list(params) for device_id, params in single_params.items()}
@@ -632,6 +632,13 @@ class ModelCheckpoint(Callback):
         if "step_num" in self._append_dict:
             self._append_dict["step_num"] = self._append_step_num + step_num
 
+    def _update_save_step(self, cb_params):
+        """update step if used async d2h copy"""
+        step_num_in_epoch = int((cb_params.cur_step_num - 1) % cb_params.batch_num + 1)
+        if self._d2h_async and self._run_mode == context.GRAPH_MODE:
+            step_num_in_epoch -= 1
+        return step_num_in_epoch
+
     def _save_ckpt(self, cb_params, force_to_save=False):
         """Save checkpoint files."""
         if cb_params.cur_step_num == self._last_triggered_step:
@@ -642,10 +649,12 @@ class ModelCheckpoint(Callback):
             self._flush_from_cache(cb_params)
 
         save_ckpt = self._check_save_ckpt(cb_params, force_to_save)
-        step_num_in_epoch = int((cb_params.cur_step_num - 1) % cb_params.batch_num + 1)
+        step_num_in_epoch = self._update_save_step(cb_params)
 
         if save_ckpt:
+
             _wait_async_save_ckpt(self._config.async_save)
+
             if self._prefix_func:
                 cur_ckpoint_file = self._prefix + f".{self._config.format}"
             else:
@@ -670,12 +679,6 @@ class ModelCheckpoint(Callback):
             self._last_time_for_keep = time.time()
             self._last_triggered_step = cb_params.cur_step_num
 
-            # TODO(MS_DISABLE_REF_MODE): Delete when remove MS_DISABLE_REF_MODE env.
-            if context.get_context("enable_ge") and os.getenv('MS_DISABLE_REF_MODE') \
-                    and context.get_context("mode") == context.GRAPH_MODE:
-                set_cur_net(cb_params.train_network)
-                cb_params.train_network.add_flags(ge_sync_data=True)
-                _cell_graph_executor(cb_params.train_network, phase='save')
             self._append_dict_content(cb_params.cur_epoch_num, cb_params.cur_step_num)
             network = self._config.saved_network if self._config.saved_network is not None else cb_params.train_network
             if os.getenv("AITURBO") == "1":
@@ -684,18 +687,13 @@ class ModelCheckpoint(Callback):
                                 crc_check=self._config.crc_check, incremental=self._map_param_inc,
                                 global_step_num=cb_params.cur_step_num)
             elif self._config.remove_redundancy:
-                parallel_mode = context.get_auto_parallel_context("parallel_mode")
-                if parallel_mode == "stand_alone":
+                if get_group_size() == 1:
                     raise TypeError(f"The deduplication feature for saving checkpoint can only be used "
-                                    f"in parallel scenarios, but got {parallel_mode}.")
+                                    f"in parallel scenarios, but got 'stand_alone'.")
                 param_layout = network.parameter_layout_dict
                 rank_id = get_rank()
                 if param_layout:
-                    device_num = _get_device_num()
-                    stage_num = _get_auto_parallel_context("pipeline_stages")
-                    chunk_size = device_num // stage_num
-                    initial_rank = (rank_id // chunk_size) * chunk_size
-                    param_redundancy_dict = get_parameter_redundancy(param_layout, initial_rank)
+                    param_redundancy_dict = get_parameter_redundancy(param_layout)
                     single_params = remove_param_redundancy(param_redundancy_dict)
                     save_param_names = single_params.get(rank_id)
                     param_layout_set = set(param_layout.keys())
@@ -704,14 +702,14 @@ class ModelCheckpoint(Callback):
                             f"For remove_redundancy save checkpoint, the saved parameters are non-redundant.")
 
                     def choice_func(x):
-                        return x not in param_layout_set or x in save_param_names
+                        return x not in param_layout_set or (save_param_names is not None and x in save_param_names)
                 else:
                     param_redundancy_dict = get_parameter_redundancy(network)
                     single_params = remove_param_redundancy(param_redundancy_dict)
                     save_param_names = single_params.get(rank_id)
 
                     def choice_func(x):
-                        return x in save_param_names
+                        return save_param_names is not None and x in save_param_names
                 save_checkpoint(network, cur_file, False, self._config.async_save,
                                 self._append_dict, self._config.enc_key, self._config.enc_mode,
                                 crc_check=self._config.crc_check, format=self._config.format,

mindspore/train/callback/_cluster_monitor.py

@@ -24,9 +24,8 @@ from threading import RLock
 from mindspore.train.callback._callback import Callback
 from mindspore.communication.management import get_rank, get_local_rank
 from mindspore import log as logger
-from mindspore.parallel._auto_parallel_context import _get_auto_parallel_context
 from mindspore.parallel._utils import _get_device_num
-from mindspore.train._utils import get_parameter_redundancy
+from mindspore.train._utils import get_parameter_redundancy, _get_pp_size_from_redundancy_map
 
 _perf_mutex = RLock()
 
@@ -42,7 +41,7 @@ def _get_dp_tp_from_redundancy(redundancy_tuple):
     return dp, tp
 
 
-def _get_dp_tp_from_layout(parameter_layout_dict, initial_rank=0):
+def _get_dp_tp_from_layout(parameter_layout_dict, initial_rank=None):
     """From layout dict get dp and tp"""
     tp = []
     dp = []
@@ -132,21 +131,9 @@ class ClusterMonitor(Callback):
         self.full_path = self.log_path + self.log_name
 
         self.write_dp_tp_flag = True
-        self.initial_rank = 0
 
     def begin(self, run_context):
         _remove_pre_log()
-        pp_num = _get_auto_parallel_context("pipeline_stages")
-        device_num = _get_device_num()
-
-        original_list = list(range(device_num))
-        chunk_size = device_num // pp_num
-        split_pp_lists = []
-        for i in range(0, device_num, chunk_size):
-            end_index = i + chunk_size if i + chunk_size <= device_num else device_num
-            split_pp_lists.append(original_list[i:end_index])
-
-        self.initial_rank = (self.global_rank // chunk_size) * chunk_size
         with _perf_mutex:
             dir_path = os.path.dirname(self.full_path)
             if not os.path.exists(dir_path):
@@ -157,8 +144,6 @@ class ClusterMonitor(Callback):
             with open(self.full_path, 'w') as file:
                 log_message = f'UUID:{self.uuid_value}\nFRAMEWORK:{self.frame_work}\nGLOBAL RANKID:{self.global_rank}\n'
                 file.write(log_message)
-                for _, split_pp_list in enumerate(split_pp_lists):
-                    file.write(f'PP:{split_pp_list}\n')
             os.chmod(self.full_path, stat.S_IRUSR)
 
     def step_begin(self, run_context):
@@ -183,10 +168,21 @@ class ClusterMonitor(Callback):
         if self.enabled and self.enabled_dtp_group and self.write_dp_tp_flag:
             cb_params = run_context.original_args()
             param_layout_dict = cb_params.train_network.parameter_layout_dict
-            dp, tp = _get_dp_tp_from_layout(param_layout_dict, self.initial_rank)
+            device_num = _get_device_num()
+            original_list = list(range(device_num))
+            param_redundancy_dict = get_parameter_redundancy(param_layout_dict)
+            pp_size = _get_pp_size_from_redundancy_map(param_redundancy_dict)
+            split_pp_lists = []
+            for i in range(0, device_num, pp_size):
+                end_index = i + pp_size if i + pp_size <= device_num else device_num
+                split_pp_lists.append(original_list[i:end_index])
+            dp, tp = _get_dp_tp_from_layout(param_layout_dict)
+
             with _perf_mutex:
                 os.chmod(self.full_path, stat.S_IWUSR)
                 with open(self.full_path, 'a') as file:
+                    for _, split_pp_list in enumerate(split_pp_lists):
+                        file.write(f'PP:{split_pp_list}\n')
                     for dp_value in dp:
                         file.write(f'dp:{dp_value}\n')
                     for tp_value in tp: