mindspore 2.6.0__cp311-cp311-win_amd64.whl → 2.7.0rc1__cp311-cp311-win_amd64.whl

mindspore/safeguard/rewrite_obfuscation.py

@@ -81,7 +81,7 @@ def _transform_target_modules(target_modules):
         obfuscate_layers = target_modules[2].split(':')
         if obfuscate_layers[1] != 'all':
             max_layers = int(obfuscate_layers[1])
-        layers = [i for i in range(0, max_layers)]
+        layers = list(range(0, max_layers))
         path_new = path.replace("blocks", "blocks/${layer}")
         network_obf_template['insert_ops'][0]['input_y'] = "obf_metadata_${layer}"
         weight_obf_template['weight_obf_ops'][0]['input_y'] = "obf_metadata_${layer}"
@@ -95,8 +95,8 @@ def _transform_target_modules(target_modules):
     obf_config['obf_metadata_config'].append(obf_medatadata)
 
     for name in target_list:
-        target_weight = path_new + '/' + name + '/weight'
-        target_bias = path_new + '/' + name + '/bias'
+        target_weight = '/'.join([path_new, name, 'weight'])
+        target_bias = '/'.join([path_new, name, 'bias'])
         weight_obf = weight_obf_template.copy()
         weight_obf['target'] = target_weight
         bias_obf = weight_obf_template.copy()
@@ -185,7 +185,7 @@ def obfuscate_ckpt(network, ckpt_files, target_modules=None, obf_config=None, sa
     def _gen_obf_metadata(config):
         name = config.get('name')
         if name is None:
-            return False
+            return
         save_metadata = config.get('save_metadata', False)
         metadata_op_name = config.get('metadata_op')
         layers = config.get('layers')
@@ -213,7 +213,6 @@ def obfuscate_ckpt(network, ckpt_files, target_modules=None, obf_config=None, sa
                         saved_obf_tensor = metadata_op(saved_obf_tensor)
                     if saved_obf_tensor is not None:
                         saved_metadata[obf_name] = saved_obf_tensor.asnumpy()
-        return True
 
     if not isinstance(network, nn.Cell):
         raise TypeError("network must be nn.Cell, but got {}.".format(type(network)))
@@ -283,13 +282,13 @@ def _obfuscate_single_ckpt(ckpt_name, obf_metadata, obf_config, saved_path):
     def _obfuscate_param(param, obf_metadata, obf_ops, layer=0):
         param_dtype = F.dtype(param)
         obf_param = param
-        for i in range(len(obf_ops)):
-            op_name = obf_ops[i].get('name')
+        for obf_op in obf_ops:
+            op_name = obf_op.get('name')
             if not isinstance(op_name, str):
                 raise TypeError('{} should be str type, but got {}'.format(op_name, type(op_name)))
             if op_name == 'mul':
                 input_x = obf_param
-                input_y_name = _get_op_input_name(obf_ops[i], 'input_y', layer)
+                input_y_name = _get_op_input_name(obf_op, 'input_y', layer)
                 input_y = obf_metadata.get(input_y_name)
                 if input_x is None or input_y is None:
                     log.error("input_x or input_y is None")
@@ -297,22 +296,22 @@ def _obfuscate_single_ckpt(ckpt_name, obf_metadata, obf_config, saved_path):
                 input_y = F.cast(input_y, param_dtype)
                 obf_param = ops.mul(input_x, input_y)
             elif op_name == 'permuate':
-                input_x_name = _get_op_input_name(obf_ops[i], 'input_x', layer)
+                input_x_name = _get_op_input_name(obf_op, 'input_x', layer)
                 p = obf_metadata.get(input_x_name, None)
                 if p is None or obf_param is None:
                     log.error("input_x or param is None")
                     return None
                 obf_param = obf_param[p]
             elif op_name == 'matmul':
-                input_x_name = _get_op_input_name(obf_ops[i], 'input_x', layer)
-                input_y_name = _get_op_input_name(obf_ops[i], 'input_y', layer)
+                input_x_name = _get_op_input_name(obf_op, 'input_x', layer)
+                input_y_name = _get_op_input_name(obf_op, 'input_y', layer)
                 input_x = _get_op_input(input_x_name, obf_param)
                 input_y = _get_op_input(input_y_name, obf_param)
                 if input_x is None or input_y is None:
                     log.error("the input_x or input_y of op: {} is None.".format(op_name))
                     return None
-                input_x = ops.transpose(input_x, (1, 0)) if obf_ops[i].get('transpose_a', False) else input_x
-                input_y = ops.transpose(input_y, (1, 0)) if obf_ops[i].get('transpose_b', False) else input_y
+                input_x = ops.transpose(input_x, (1, 0)) if obf_op.get('transpose_a', False) else input_x
+                input_y = ops.transpose(input_y, (1, 0)) if obf_op.get('transpose_b', False) else input_y
                 obf_param = ops.matmul(F.cast(input_x, param_dtype), F.cast(input_y, param_dtype))
             else:
                 log.error("unsupported op, op must be matmul or permuate or mul, but got {}."

mindspore/train/_utils.py

@@ -582,7 +582,8 @@ def _progress_bar(iterable, total=None):
         print_progress_bar(i)
 
 
-def _load_and_transform(path, name_map, load_func, transform_func):
+def _load_and_transform(path, name_map, load_func, transform_func=None):
+    """use load_func to load and use transform_func to convert"""
     if load_func is not None:
         param_dict = load_func(path)
     else:
@@ -590,5 +591,8 @@ def _load_and_transform(path, name_map, load_func, transform_func):
     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)
+        if transform_func is not None:
+            transform_dict[new_name] = transform_func(v, new_name)
+        else:
+            transform_dict[new_name] = v
     return transform_dict

mindspore/train/amp.py

@@ -69,6 +69,9 @@ AMP_BLACK_LIST = [
 AMP_AUTO_WHITE_LIST = [
     P.Conv2D,
     P.Conv3D,
+    gen.Conv2DExt,
+    gen.Conv3DExt,
+    gen.ConvTranspose2D,
     P.Conv2DTranspose,
     P.Conv3DTranspose,
     gen.Convolution,
@@ -80,6 +83,10 @@ AMP_AUTO_WHITE_LIST = [
     P.Einsum,
     gen.Dense,
     gen.Addmm,
+    gen.Addbmm,
+    gen.Addmv,
+    gen.Baddbmm,
+    gen.Mv,
 ]
 
 AMP_AUTO_BLACK_LIST = [
@@ -90,8 +97,10 @@ AMP_AUTO_BLACK_LIST = [
     P.Erfinv,
     P.Exp,
     P.Expm1,
-    P.Log,
-    P.Log1p,
+    gen.Log,
+    gen.Log10,
+    gen.Log1p,
+    gen.Log2,
     P.Reciprocal,
     P.Rsqrt,
     P.Sinh,
@@ -103,6 +112,7 @@ AMP_AUTO_BLACK_LIST = [
     P.BatchNorm,
     gen.BatchNormExt,
     gen.GroupNorm,
+    gen.Norm,
     P.KLDivLoss,
     P.SmoothL1Loss,
     P.MultilabelMarginLoss,
@@ -113,7 +123,19 @@ AMP_AUTO_BLACK_LIST = [
     P.Pdist,
     P.Cdist,
     P.Renorm,
+    gen.ReduceProd,
+    gen.Softmax,
+    gen.LogSoftmax,
+    gen.LogSoftmaxExt,
+    gen.CumProd,
+    gen.CumSum,
+    gen.CumsumExt,
+    gen.ProdExt,
+    gen.SumExt,
+    gen.L1LossExt,
     gen.MSELossExt,
+    gen.NLLLoss,
+    gen.NLLLoss2d,
 ]
 
 # Indicates which inputs of primitives need to be converted
@@ -358,7 +380,7 @@ def _auto_black_list(network, black_list, dtype):
     return network
 
 
-class amp_decorator:
+class AmpDecorator:
     """
     Auto mixed precision decorator.
     Type of lists: List[Tuple[str, List[int]]]
@@ -384,7 +406,7 @@ def _set_amp_decorator(obj, amp_level, amp_dtype, white_list, black_list):
     if inspect.isfunction(obj) or inspect.ismethod(obj):
         @functools.wraps(obj)
         def wrapper(*args, **kwargs):
-            with amp_decorator(amp_level, amp_dtype, white_list, black_list):
+            with AmpDecorator(amp_level, amp_dtype, white_list, black_list):
                 return obj(*args, **kwargs)
         return wrapper
     if isinstance(obj, nn.Cell):
@@ -423,17 +445,18 @@ def auto_mixed_precision(network, amp_level="O0", dtype=mstype.float16):
 
     Operators in `auto_whitelist` are:
 
-    ``Conv2D``, ``Conv3D``, ``Conv2DTranspose``, ``Conv3DTranspose``, ``Convolution``, ``MatMul``, ``MatMulExt``,
-    ``BatchMatMul``, ``BatchMatMulExt``, ``PReLU``, ``Einsum``, ``Dense``, ``Addmm``
+    ``Conv2D``, ``Conv2DExt``, ``Conv3D``, ``Conv3DExt``, ``Conv2DTranspose``, ``ConvTranspose2D``,
+    ``Conv3DTranspose``, ``Convolution``, ``MatMul``, ``MatMulExt``, ``BatchMatMul``, ``BatchMatMulExt``, ``PReLU``,
+    ``Einsum``, ``Dense``, ``Addmm``, ``Addbmm``, ``Addmv``, ``Baddbmm``, ``Mv``
 
     Operators in `auto_blacklist` are:
 
-    ``Pow``, ``ACos``, ``Asin``, ``Cosh``, ``Erfinv``, ``Exp``, ``Expm1``, ``Log``, ``Log1p``, ``Reciprocal``,
-    ``Rsqrt``, ``Sinh``, ``Tan``, ``Softplus``, ``SoftplusExt``, ``LayerNorm``, ``LayerNormExt``, ``BatchNorm``,
-    ``BatchNormExt``, ``GroupNorm``, ``KLDivLoss``, ``SmoothL1Loss``, ``MultilabelMarginLoss``, ``SoftMarginLoss``,
-    ``TripletMarginLoss``, ``MultiMarginLoss``, ``BCEWithLogitsLoss``, ``Pdist``, ``Cdist``, ``Renorm``,
-    ``ReduceProd``, ``Softmax``, ``LogSoftmax``, ``CumProd``, ``CumSum``, ``CumsumExt``, ``ProdExt``, ``SumExt``,
-    ``Norm``, ``MSELossExt``
+    ``Pow``, ``ACos``, ``Asin``, ``Cosh``, ``Erfinv``, ``Exp``, ``Expm1``, ``Log``, ``Log10``, ``Log1p``, ``Log2``,
+    ``Reciprocal``, ``Rsqrt``, ``Sinh``, ``Tan``, ``Softplus``, ``SoftplusExt``, ``LayerNorm``, ``LayerNormExt``,
+    ``BatchNorm``, ``BatchNormExt``, ``GroupNorm``, ``KLDivLoss``, ``SmoothL1Loss``, ``MultilabelMarginLoss``,
+    ``SoftMarginLoss``, ``TripletMarginLoss``, ``MultiMarginLoss``, ``BCEWithLogitsLoss``, ``Pdist``, ``Cdist``,
+    ``Renorm``, ``ReduceProd``, ``Softmax``, ``LogSoftmax``, ``LogSoftmaxExt``, ``CumProd``, ``CumSum``,
+    ``CumsumExt``, ``ProdExt``, ``SumExt``, ``Norm``, ``L1LossExt``, ``MSELossExt``, ``NLLLoss``, ``NLLLoss2d``
 
     Operators in `promote_list` are:
 
@@ -638,7 +661,7 @@ def _add_loss_network(network, loss_fn, cast_model_type):
 
 
 def _is_grad_accumulation(mcell):
-    if mcell.cls_name == "GradAccumulationCell" or mcell.cls_name == "GradAccumulation":
+    if mcell.cls_name in {"GradAccumulationCell", "GradAccumulation"}:
         return True
     for cell in mcell.cells():
         if _is_grad_accumulation(cell):
@@ -675,23 +698,23 @@ def build_train_network(network, optimizer, loss_fn=None, level='O0', boost_leve
     Build the mixed precision training cell automatically.
 
     Note:
-        - After using `custom_mixed_precision` or `auto_mixed_precision` for precision conversion, it is not supported
-          to perform the precision conversion again. If  `build_train_network` is used to train a converted network,
-          `level` need to be configured to ``O0`` to avoid the duplicated accuracy conversion.
+        After using `custom_mixed_precision` or `auto_mixed_precision` for precision conversion, it is not supported
+        to perform the precision conversion again. If  `build_train_network` is used to train a converted network,
+        `level` need to be configured to ``O0`` to avoid the duplicated accuracy conversion.
 
     Args:
         network (Cell): Definition of the network.
         optimizer (:class:`mindspore.nn.Optimizer`): Define the optimizer to update the Parameter.
-        loss_fn (Union[None, Cell]): Define the loss function. If None, the `network` should have the loss inside.
-            Default: ``None`` .
-        level (str): Supports ['O0', 'O1', 'O2', 'O3', 'auto']. Default: ``'O0'`` .
+        loss_fn (Union[None, Cell], optional): Define the loss function. If None,
+            the `network` should have the loss inside. Default: ``None`` .
+        level (str, optional): Supports ['O0', 'O1', 'O2', 'O3', 'auto']. Default: ``'O0'`` .
 
             For details on amp level, refer to :func:`mindspore.amp.auto_mixed_precision`.
 
             Property of `keep_batchnorm_fp32`, `cast_model_type` and `loss_scale_manager` determined by `level`
             setting may be overwritten by settings in `kwargs`.
 
-        boost_level (str): Option for argument `level` in `mindspore.boost` , level for boost mode
+        boost_level (str, optional): Option for argument `level` in `mindspore.boost` , level for boost mode
             training. Supports ['O0', 'O1', 'O2']. Default: ``'O0'`` .
 
             - 'O0': Do not change.

mindspore/train/callback/__init__.py

@@ -15,6 +15,11 @@
 """Callback related classes and functions."""
 from __future__ import absolute_import
 
+__all__ = ["Callback", "LossMonitor", "TimeMonitor", "ModelCheckpoint", "FlopsUtilizationCollector",
+           "SummaryCollector", "CheckpointConfig", "RunContext", "LearningRateScheduler", "SummaryLandscape",
+           "History", "LambdaCallback", "ReduceLROnPlateau", "EarlyStopping", "OnRequestExit", "BackupAndRestore",
+           "TrainFaultTolerance"]
+
 from mindspore.train.callback._callback import Callback
 from mindspore.train.callback._callback import CallbackManager as _CallbackManager
 from mindspore.train.callback._callback import InternalCallbackParam as _InternalCallbackParam
@@ -37,8 +42,3 @@ 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._train_fault_tolerance import TrainFaultTolerance
-
-__all__ = ["Callback", "LossMonitor", "TimeMonitor", "ModelCheckpoint", "FlopsUtilizationCollector",
-           "SummaryCollector", "CheckpointConfig", "RunContext", "LearningRateScheduler", "SummaryLandscape",
-           "History", "LambdaCallback", "ReduceLROnPlateau", "EarlyStopping", "OnRequestExit", "BackupAndRestore",
-           "TrainFaultTolerance"]

mindspore/train/callback/_checkpoint.py

@@ -411,8 +411,6 @@ class CheckpointConfig:
             handle_append_info["epoch_num"] = 0
         if "step_num" in append_info:
             handle_append_info["step_num"] = 0
-        if "random_op" in append_info:
-            handle_append_info["random_op"] = 0
         dict_num = 0
         for element in append_info:
             if not isinstance(element, str) and not isinstance(element, dict):
@@ -588,8 +586,6 @@ class ModelCheckpoint(Callback):
         # save graph (only once)
         if not self._graph_saved:
             graph_file_name = os.path.join(self._directory, self._prefix + '-graph.meta')
-            if os.path.isfile(graph_file_name) and context.get_context("mode") == context.GRAPH_MODE:
-                os.remove(graph_file_name)
             _save_graph(cb_params.train_network, graph_file_name)
             self._graph_saved = True
         self._save_ckpt(cb_params)
@@ -713,12 +709,13 @@ class ModelCheckpoint(Callback):
                 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,
-                                incremental=self._map_param_inc, choice_func=choice_func)
+                                incremental=self._map_param_inc, choice_func=choice_func,
+                                remove_redundancy=self._config.remove_redundancy)
             else:
                 save_checkpoint(network, cur_file, self._config.integrated_save, 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,
-                                incremental=self._map_param_inc)
+                                incremental=self._map_param_inc, remove_redundancy=self._config.remove_redundancy)
 
             self._latest_ckpt_file_name = cur_file
 

mindspore/train/callback/_flops_collector.py

@@ -53,7 +53,7 @@ class FlopsUtilizationCollector(Callback):
     The FlopsUtilizationCollector interface counts the model utilization information MFU
     and the hardware utilization information HFU.
     Currently, the API counts only the forward and backward flops of MatMul,
-    BatchMatMul, FlashAttentionScore, and Conv2D operators.
+    BatchMatMul, flash_attention_score, and Conv2D operators.
     Only used in graph mode with static shape.
 
     Args:

mindspore/train/callback/_landscape.py

@@ -404,7 +404,6 @@ class SummaryLandscape:
     def _set_context(device_id):
         """Set context."""
         context.set_context(device_id=device_id)
-        context.set_context(mode=context.GRAPH_MODE)
 
     def _create_landscape_by_pca(self, epochs, proz, landscape_size, device_ids=None, callback_fn=None, executor=None):
         """Create landscape by PCA."""

mindspore/train/callback/_train_fault_tolerance.py

@@ -25,8 +25,9 @@ from mindspore.communication import get_rank, get_group_size
 from mindspore import log as logger
 from mindspore.train.serialization import _get_cur_rank_dp
 from mindspore._c_expression import _repair_device, _stop_device, _tft_sem_post, _tft_sem_enable
-from mindspore._c_expression import _rebuild_world_group, _rebuild_sub_group, _finalize_comm
+from mindspore._c_expression import _rebuild_world_group, _rebuild_sub_group, _finalize_comm, _clean_rootinfo
 from mindspore._c_expression import clean_tdt_channel
+from mindspore._c_expression import _pre_launch_send_recv
 from mindspore._c_expression import send_recv, reset_params
 from mindspore._c_expression import CollectiveManager
 from mindspore._c_expression import _get_uce_process_strategy, _get_uce_mem_info
@@ -35,6 +36,7 @@ from mindspore.ops.operations.manually_defined._inner import TensorReport
 import mindspore
 import mindspore.common.dtype as mstype
 from mindspore.parallel._recovery_context import _set_recovery_context
+from mindspore import runtime
 
 
 def _get_ckpt_dir(step, ckpt_save_path, is_tmp_file):
@@ -80,7 +82,7 @@ def _save_checkpoint_on_failure(step, save_info, args, cb_ctx):
         append_dict["loss_scale"] = outputs[2]
 
     ckpt_file = f"ttp_rank_{str(cur_rank)}-{str(cur_epoch_num)}_{str(step_num_in_epoch)}.ckpt"
-    cur_ckpt_dir = _get_ckpt_dir(step, ckpt_save_path, True) + "/rank_" + str(cur_rank)
+    cur_ckpt_dir = os.path.join(_get_ckpt_dir(step, ckpt_save_path, True), "rank_" + str(cur_rank))
     os.makedirs(cur_ckpt_dir, exist_ok=True)
     cur_file = os.path.join(cur_ckpt_dir, ckpt_file)
     save_checkpoint(cb_params.train_network, cur_file,
@@ -110,7 +112,7 @@ def _tft_exit_cb(ctx):
 
 def _tft_repair_callback(step, need_rebuild, error_ranks, repair_info, args, cb_ctx):
     """ Callback used for TFT repair function."""
-    logger.warning("Enter _tft_repair_callback repair type: {}".format(repair_info["repair_type"]))
+    logger.warning(f"Enter _tft_repair_callback repair type: {repair_info['repair_type']}")
     if (repair_info["repair_type"] in (cb_ctx.tft.RepairType.RT_UCE_HIGHLEVEL.value,
                                        cb_ctx.tft.RepairType.RT_UCE_LOWLEVEL.value)):
         logger.warning("Enter _tft_repair_callback uce REPARI_DEVICE device_id : {}".format(cb_ctx.device_id))
@@ -138,7 +140,7 @@ def _tft_repair_callback(step, need_rebuild, error_ranks, repair_info, args, cb_
 
 def _tft_clean_callback(is_uce_error, args, ctx):
     """ Callback used for TFT clean function."""
-    logger.warning("Enter _tft_clean_callback")
+    logger.warning(f"Enter _tft_clean_callback, device id:{ctx.device_id}")
     ret = 0
     if is_uce_error:
         _get_uce_mem_info(ctx.device_id)
@@ -154,29 +156,36 @@ def _tft_clean_callback(is_uce_error, args, ctx):
     logger.warning("Enter _tft_clean_callback resume_hccl_comm")
     CollectiveManager.get_instance().resume_hccl_comm()
     logger.warning("Finish _tft_clean_callback, ret: {}".format(ret))
+    if ctx.tft.tft_get_repair_type() == "recover":
+        logger.warning(f"Destroy hcom")
+        _finalize_comm()
+        logger.warning(f"Destroy hcom end")
     return ret
 
 
 def _tft_stop_callback(args, cb_ctx):
     """ Callback used for TFT stop function."""
-    logger.warning("Enter _tft_stop_callback device_id: {}".format(cb_ctx.device_id))
+    logger.warning(f"Enter _tft_stop_callback device_id: {cb_ctx.device_id}")
     _stop_device(cb_ctx.device_id)
+    cb_ctx.stop_been_called = True
     if (not cb_ctx.is_uce_rank) and (not cb_ctx._is_params_consistent()):  # pylint: disable=W0212
         raise RuntimeError("Can't stop device, because training parameters are left in inconsistent state!")
     cb_ctx.is_uce_rank = False
     if cb_ctx.tft.tft_get_repair_type() == "recover":
         logger.warning(f"Reset limit step")
         cb_ctx.tft.tft_reset_limit_step()
-    logger.info("Finish _tft_stop_callback")
+    logger.warning("Finish _tft_stop_callback")
 
 
 def _tft_rebuild_sub_groups(fault_ranks, args, ctx):
     """Callback used for TFT Rebuild Group function."""
-    logger.warning(f"Enter _tft_rebuild_sub_groups, device id: ".format(ctx.device_id))
-    _finalize_comm()
+    logger.warning(f"Enter _tft_rebuild_sub_groups, device id: {ctx.device_id}")
     _rebuild_world_group()
     _rebuild_sub_group()
     _set_recovery_context(is_arf=True)
+    logger.warning(f"try to pre launch send recv before real launch")
+    _pre_launch_send_recv(context.get_context('device_id'))
+    logger.warning(f"Pre launch send recv before real launch end")
     logger.warning("Enter _tft_rebuild_sub_groups ok ")
 
 
@@ -299,6 +308,12 @@ class TrainFaultTolerance(Callback):
 
     def __init__(self, ckpt_save_path=None, **kwargs):
         super(TrainFaultTolerance, self).__init__()
+        logger.info(f"MS_ENABLE_TFT: {os.getenv('MS_ENABLE_TFT', '')}")
+        if self._only_enable_tsp():
+            self.tft = _tft_handler.get_tft()
+            self._check_init()
+            self.tft.tft_register_stream_sync_handler(runtime.synchronize, self)
+            return
         self.save_cb = kwargs.get("ckpt_save_fn", None)
         self.ckpt_save_path = ckpt_save_path
         if self.save_cb is None and self.ckpt_save_path is None:
@@ -308,19 +323,24 @@ class TrainFaultTolerance(Callback):
         self.device_id = context.get_context("device_id")
         self.cur_step_num = 0
         self.cur_epoch_num = 0
+        self.clean_unique_id = False
         # For TREError(Training Result Error) scene, parameter `ckpt_load_fn` must be provided to load checkpoint
         # from file for resuming training, the `ckpt_load_fn` is a function, prototype of which is:
         # `def load_checkpoint() -> tuple(dict, bool)`, the return value is a tuple containing 2 values,
         # i.e. (param_dict, remove_redundancy)
         self.ckpt_load_func = kwargs.get("ckpt_load_fn", None)
-        self.tft = _tft_handler.get_tft()
         if self._only_enable_tre():
             return
+        self.tft = _tft_handler.get_tft()
         self._check_init()
+        if self._only_enable_tre_and_tsp():
+            self.tft.tft_register_stream_sync_handler(runtime.synchronize, self)
+            return
         self.global_step = None
         self.learning_rate = None
         self.has_init_replica = False
         self.is_uce_rank = False
+        self.stop_been_called = False
 
         self.assign = mindspore.ops.Assign()
         self.g_one = Parameter(Tensor([1], dtype=mstype.int32))
@@ -336,6 +356,22 @@ class TrainFaultTolerance(Callback):
             return False
         return "TRE:1" in env_enable
 
+    def _only_enable_tsp(self):
+        """Check if only configured MS_ENABLE_TFT='{TSP:1}'"""
+        env_enable = os.getenv("MS_ENABLE_TFT", "")
+        non_tsp_flags = ["TTP:1", "UCE:1", "ARF:1", "TRE:1"]
+        if any(flag in env_enable for flag in non_tsp_flags):
+            return False
+        return "TSP:1" in env_enable
+
+    def _only_enable_tre_and_tsp(self):
+        """Check if only configured MS_ENABLE_TFT='{TRE:1, TSP:1}'"""
+        env_enable = os.getenv("MS_ENABLE_TFT", "")
+        other_flags = ["TTP:1", "UCE:1", "ARF:1"]
+        if any(flag in env_enable for flag in other_flags):
+            return False
+        return "TRE:1" in env_enable and "TSP:1" in env_enable
+
     def _check_init(self):
         """Check if the mindio-ttp had inited"""
         if self.tft is None:
@@ -427,6 +463,8 @@ class TrainFaultTolerance(Callback):
         self.tft.tft_register_clean_handler(_tft_clean_callback, self)
         self.tft.tft_register_repair_handler(_tft_repair_callback, self)
         self.tft.tft_register_rebuild_group_handler(_tft_rebuild_sub_groups, self)
+        if "TSP:1" in os.getenv("MS_ENABLE_TFT", ""):
+            self.tft.tft_register_stream_sync_handler(runtime.synchronize, self)
 
     def _reset_acc_grads(self):
         accu_grad_params = map(lambda e: e[1],
@@ -436,6 +474,12 @@ class TrainFaultTolerance(Callback):
         if reset_params(accu_grad_list) != 0:
             raise ValueError("Call reset_params failed.")
 
+    def _clear_unique_id(self):
+        """Clean unique id on first train step end"""
+        if not self.clean_unique_id and ("ARF:1" in os.getenv("MS_ENABLE_TFT", "")):
+            _clean_rootinfo()
+            self.clean_unique_id = True
+
     def on_train_step_end(self, run_context):
         """
         Report status to MindIO TFT after every step finished.
@@ -446,13 +490,19 @@ class TrainFaultTolerance(Callback):
         """
         if self._only_enable_tre():
             return
-        if self.has_init_replica is False:
-            self.has_init_replica = True
-            self._set_tft_optimizer_replica(run_context)
+
         cb_params = run_context.original_args()
         logger.info("START Set optimizer finish step status to TFT. step: {}".format(cb_params.cur_step_num))
         self.cur_step_num = cb_params.cur_step_num
         self.cur_epoch_num = cb_params.cur_epoch_num
+        if self._only_enable_tsp() or self._only_enable_tre_and_tsp():
+            logger.info("Go into tft_pause_train.")
+            self.tft.tft_pause_train(self.cur_step_num)
+            return
+
+        if self.has_init_replica is False:
+            self.has_init_replica = True
+            self._set_tft_optimizer_replica(run_context)
         if cb_params.optimizer is not None:
             self.global_step = cb_params.optimizer.global_step.clone()
             self.assign(cb_params.optimizer.tft_g_one_flag, self.g_one)
@@ -462,7 +512,13 @@ class TrainFaultTolerance(Callback):
         else:
             raise ValueError("TFT feature need optimizer or network's optimizer!")
         self.tft.tft_end_updating_os(cb_params.cur_step_num + self.initial_step)
+        if cb_params.is_arf:
+            self.clean_unique_id = False
+        self._clear_unique_id()
         logger.info("END Set optimizer finish step status to TFT.")
+        if "TSP:1" in os.getenv("MS_ENABLE_TFT", ""):
+            logger.info("Go into tft_pause_train.")
+            self.tft.tft_pause_train(self.cur_step_num)
 
     def on_train_begin(self, run_context):
         """
@@ -472,6 +528,8 @@ class TrainFaultTolerance(Callback):
             run_context (RunContext): Context of the train running. Refer to
                                       :class:`mindspore.train.RunContext` for detail.
         """
+        if self._only_enable_tsp():
+            return
         cb_params = run_context.original_args()
         if self._only_enable_tre():
             self.cb_params = cb_params
@@ -491,6 +549,6 @@ class TrainFaultTolerance(Callback):
             run_context (RunContext): Context of the train running. Refer to
                                       :class:`mindspore.train.RunContext` for detail.
         """
-        if self._only_enable_tre():
+        if self._only_enable_tre() or self._only_enable_tsp() or self._only_enable_tre_and_tsp():
             return
         _tft_handler.unregister_tft()

mindspore/train/data_sink.py

@@ -18,7 +18,7 @@ import mindspore.ops as ops
 from mindspore import context
 from mindspore.common.dtype import pytype_to_dtype
 from mindspore.common.api import jit
-from mindspore.train._utils import _exec_datagraph, _get_types_and_shapes
+from mindspore.train._utils import _exec_datagraph, _get_types_and_shapes, enable_data_broadcast
 from mindspore.train.dataset_helper import _has_dynamic_shape, _check_inputs
 import mindspore.dataset as ds
 from mindspore._c_expression import _set_dataset_mode_config
@@ -41,6 +41,15 @@ def _init_sink_dataset(dataset, sink_size, input_signature, create_info):
     is_info_queue = (create_info and sink_size == 1 and dataset_size != 1 and
                      input_signature is None and not dynamic_shape and
                      context.get_context('device_target') == 'Ascend')
+
+    # 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()
+
+    if use_pipeline_parallel or data_broadcast:
+        is_info_queue = False
+
     transfer_dataset = _exec_datagraph(dataset, sink_size, create_data_info_queue=is_info_queue)
     dataset.__transfer_dataset__ = transfer_dataset
 
@@ -214,7 +223,7 @@ def data_sink(fn, dataset, sink_size=1, jit_config=None, input_signature=None):
     loop = sink_size
     create_info = True
     if jit_config is None:
-        create_info = (loop == 1)
+        create_info = loop == 1
         loop = 1
     ori_next_op, is_info_queue = _init_sink_dataset(dataset, loop, input_signature, create_info)
 

mindspore/train/dataset_helper.py

@@ -564,6 +564,15 @@ class _DatasetIter:
                 self.sink_size = dataset.__loop_size__
             create_data_info_queue = (
                 sink_size == 1 and self.sink_count == 1 and dataset.get_dataset_size() != 1)
+
+            # 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()
+
+            if use_pipeline_parallel or data_broadcast:
+                create_data_info_queue = False
+
             dataset.__transfer_dataset__ = _exec_datagraph(dataset, self.sink_size,
                                                            create_data_info_queue=create_data_info_queue)