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

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,12 +156,16 @@ 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)
     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!")
@@ -167,23 +173,25 @@ def _tft_stop_callback(args, cb_ctx):
     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 ")
 
 
 class TrainFaultTolerance(Callback):
     """
     This callback is used to enable the TFT feature
-    `MindIO TFT <https://www.hiascend.com/document/detail/zh/mindx-dl/60rc2/mindio/mindiottp/mindiottp001.html>`_
+    `MindIO TFT <https://www.hiascend.com/document/detail/zh/mindx-dl/600/clusterscheduling/ref/mindiottp/mindiotft001.html>`_
     and will execute TFT operations during training process, such as TFT init, report and exception handle.
 
     Note:
@@ -299,27 +307,69 @@ 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:
             raise ValueError("TrainFaultTolerance construct need to set ckpt_save_fn or ckpt_save_path!")
+        self.cb_params = None
+        self.initial_step = kwargs.get("initial_step", 0)
+        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)
+        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.cb_params = None
-        self.initial_step = kwargs.get("initial_step", 0)
-        self.device_id = context.get_context("device_id")
+
         self.assign = mindspore.ops.Assign()
         self.g_one = Parameter(Tensor([1], dtype=mstype.int32))
         self.s1 = mindspore.hal.Stream()
-        self.cur_step_num = 0
-        self.cur_epoch_num = 0
         _tft_sem_enable()
         self._tft_register()
 
+    def _only_enable_tre(self):
+        """Check if only configured MS_ENABLE_TFT='{TRE:1}'"""
+        env_enable = os.getenv("MS_ENABLE_TFT", "")
+        non_tre_flags = ["TTP:1", "UCE:1", "ARF:1"]
+        if any(flag in env_enable for flag in non_tre_flags):
+            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:
@@ -330,11 +380,9 @@ class TrainFaultTolerance(Callback):
             _tft_handler.init(config=None)
             self.tft = _tft_handler.get_tft()
             logger.warning(f"TFT handle init ok.")
-        mode = context.get_context("mode")
         device_target = context.get_context("device_target")
-        if device_target != "Ascend" or mode != context.GRAPH_MODE:
-            raise ValueError(f"MindIO adataper only support on Ascend device with GRAPH Mode!"
-                             f"device:{device_target}, run mode: {mode}")
+        if device_target != "Ascend":
+            raise ValueError(f"MindIO adataper only support on Ascend device but got device {device_target}!")
 
     def _is_params_consistent(self):
         for key, param in self.cb_params.train_network.parameters_and_names():
@@ -411,6 +459,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],
@@ -420,6 +470,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.
@@ -428,13 +484,21 @@ class TrainFaultTolerance(Callback):
             run_context (RunContext): Context of the train running. Refer to
                                       :class:`mindspore.train.RunContext` for detail.
         """
-        if self.has_init_replica is False:
-            self.has_init_replica = True
-            self._set_tft_optimizer_replica(run_context)
+        if self._only_enable_tre():
+            return
+
         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)
@@ -444,7 +508,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):
         """
@@ -454,7 +524,12 @@ 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
+            return
         sink_size = cb_params.get("sink_size", 0)
         if sink_size > 1:
             raise ValueError("TFT feature doesn't support sink_size > 1.")
@@ -470,4 +545,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() or self._only_enable_tsp() or self._only_enable_tre_and_tsp():
+            return
         _tft_handler.unregister_tft()

mindspore/train/data_sink.py

@@ -16,9 +16,9 @@
 from functools import wraps
 import mindspore.ops as ops
 from mindspore import context
-from mindspore.common.dtype import pytype_to_dtype
+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
 
@@ -52,7 +61,7 @@ def _init_sink_dataset(dataset, sink_size, input_signature, create_info):
         _check_inputs(input_signature, dataset_shapes, dataset_types)
 
     queue_name = transfer_dataset.queue_name
-    if _need_to_full() and context.get_context('mode') == context.GRAPH_MODE:
+    if _need_to_full():
         device_num = _get_device_num() // _get_pipeline_stages()
         dataset_shapes = _to_full_shapes(dataset_shapes, device_num)
     next_op = ops.GetNext(dataset_types, dataset_shapes, len(dataset_types), queue_name)
@@ -85,12 +94,12 @@ def _get_next_op(dataset, ori_next_op, is_info_queue):
 
     queue_name = dataset.__transfer_dataset__.queue_name
     dataset_types, dataset_shapes = dataset.__transfer_dataset__.get_data_info()
-    dataset_types = [pytype_to_dtype(x) for x in dataset_types]
+    dataset_types = [_pytype_to_dtype(x) for x in dataset_types]  # pylint:disable=protected-access
     key = str(dataset_types) + str(dataset_shapes)
     if key in dataset.__sink_aux__.next_ops:
         next_op = dataset.__sink_aux__.next_ops[key]
     else:
-        if _need_to_full() and context.get_context('mode') == context.GRAPH_MODE:
+        if _need_to_full():
             device_num = _get_device_num() // _get_pipeline_stages()
             dataset_shapes = _to_full_shapes(dataset_shapes, device_num)
         next_op = ops.GetNext(dataset_types, dataset_shapes, len(dataset_types), queue_name)
@@ -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

@@ -1,4 +1,4 @@
-# Copyright 2020 Huawei Technologies Co., Ltd
+# Copyright 2020-2025 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -20,8 +20,8 @@ import copy
 
 from mindspore import _checkparam as Validator
 from mindspore import log as logger
-from mindspore.common._auto_dynamic import is_auto_dynamic, convert_new_shapes
-from mindspore.common.dtype import pytype_to_dtype
+from mindspore.common.dynamic_shape._auto_dynamic import is_auto_dynamic, convert_new_shapes
+from mindspore.common.dtype import _pytype_to_dtype
 from mindspore.common.api import _cell_graph_executor, _is_args_fullmode, ARG_SPECIFIED
 from mindspore.common._utils import is_shape_unknown
 from mindspore.dataset.core import config as dataset_config
@@ -34,7 +34,7 @@ from mindspore.parallel._utils import _get_device_num, _get_global_rank, _need_t
     _origin_shapes, _dynamic_shape_for_dataset
 from mindspore.parallel._ps_context import _is_role_sched
 from mindspore.ops import operations as P
-from mindspore.common.auto_dynamic_shape import _auto_dynamic_shape
+from mindspore.common.dynamic_shape.auto_dynamic_shape import _auto_dynamic_shape
 
 
 def _send_data(dataset, epoch_num):
@@ -275,7 +275,7 @@ def connect_network_with_dataset(network, dataset_helper):
         # Need to do full_batch for shapes which also do in the _DatasetIterMSLoopSink
         if _need_to_full():
             dataset_shapes = _to_full_shapes(dataset_shapes, _get_device_num() // _get_pipeline_stages())
-        dataset_types = [pytype_to_dtype(x) for x in dataset_types]
+        dataset_types = [_pytype_to_dtype(x) for x in dataset_types]  # pylint:disable=protected-access
         if not is_dynamic:
             dataset_shapes = _auto_dynamic_shape.auto_dynamic_generate_compile_args(dataset_shapes, True)
         key = str(dataset_types) + str(dataset_shapes)
@@ -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)