mindspore 2.7.0rc1__cp310-cp310-win_amd64.whl → 2.7.1__cp310-cp310-win_amd64.whl

mindspore/tools/stress_detect.py

@@ -0,0 +1,63 @@
+# Copyright 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.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Stress detect."""
+from mindspore import _c_expression
+from mindspore import log as logger
+from mindspore.communication import init, create_group, get_rank
+from mindspore.communication import get_local_rank_size
+
+
+def stress_detect(detect_type="aic"):
+    """
+    Used to detect whether there are faults in hardware accuracy or communication between links.
+    The common usage scenario is to initiate a new thread or call this interface through a Callback function
+    at each step or when saving checkpoints, to check whether hardware malfunctions could affect accuracy.
+
+    Args:
+        detect_type (str, optional): The type of stress test to perform. There are two options available: ``'aic'`` and
+            ``'hccs'``, which perform AiCore and HCCS link stress tests on the device, respectively. Default: "aic".
+
+    Returns:
+        int, the return value represents the error type. 0 indicates normal. 1 indicates failure to start some or
+        all test cases. 2 indicates a hardware failure, and it is recommended to replace the device.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> from mindspore.tools import stress_detect
+        >>> ret = stress_detect()
+        >>> print(ret)
+        0
+    """
+    if detect_type not in ["aic", "hccs"]:
+        logger.error(f"For stress detect, detection type must be 'aic' or 'hccs'."
+                     f"But got {detect_type}. Exiting stress detect.")
+        return 1
+
+    if detect_type == "aic":
+        return _c_expression.stress_detect("aic")
+
+    init()
+    local_ranks = []
+    local_rank_size = get_local_rank_size()
+    node_num = get_rank() // local_rank_size
+    for i in range(local_rank_size):
+        local_ranks.append(local_rank_size * node_num + i)
+    if get_rank() in local_ranks:
+        group = f"new_group_{node_num}"
+        create_group(group, local_ranks)
+
+    return _c_expression.stress_detect(group)

mindspore/train/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Huawei Technologies Co., Ltd
+# Copyright 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.
@@ -25,8 +25,8 @@ from mindspore.train import amp
 from mindspore.train.amp import build_train_network
 from mindspore.train.loss_scale_manager import LossScaleManager, FixedLossScaleManager, DynamicLossScaleManager
 from mindspore.train.serialization import save_checkpoint, load_checkpoint, load_param_into_net, export, \
-    load, parse_print, async_ckpt_thread_status, convert_model, export_split_mindir, \
-    load_checkpoint_async, check_checkpoint, get_ckpt_path_with_strategy, ckpt_to_safetensors, safetensors_to_ckpt, \
+    load, async_ckpt_thread_status, export_split_mindir, \
+    load_checkpoint_async, get_ckpt_path_with_strategy, ckpt_to_safetensors, safetensors_to_ckpt, \
     build_searched_strategy, merge_sliced_parameter, load_distributed_checkpoint, restore_group_info_list
 from mindspore.train.callback import Callback, LossMonitor, TimeMonitor, ModelCheckpoint, SummaryCollector, \
     CheckpointConfig, RunContext, LearningRateScheduler, SummaryLandscape, FlopsUtilizationCollector, \
@@ -37,9 +37,9 @@ from mindspore.train.metrics import *
 from mindspore.train.data_sink import data_sink
 
 __all__ = ["Model", "DatasetHelper", "connect_network_with_dataset", "build_train_network", "LossScaleManager",
-           "FixedLossScaleManager", "DynamicLossScaleManager", "save_checkpoint", "load_checkpoint", "check_checkpoint",
-           "load_param_into_net", "export", "load", "export_split_mindir", "parse_print", "async_ckpt_thread_status",
-           "convert_model", "data_sink", "load_checkpoint_async", "get_ckpt_path_with_strategy", "ckpt_to_safetensors",
+           "FixedLossScaleManager", "DynamicLossScaleManager", "save_checkpoint", "load_checkpoint",
+           "load_param_into_net", "export", "load", "export_split_mindir", "async_ckpt_thread_status",
+           "data_sink", "load_checkpoint_async", "get_ckpt_path_with_strategy", "ckpt_to_safetensors",
            "safetensors_to_ckpt", "build_searched_strategy", "merge_sliced_parameter", "load_distributed_checkpoint",
            "restore_group_info_list"]
 __all__.extend(callback.__all__)

mindspore/train/_utils.py

@@ -26,7 +26,7 @@ import numpy as np
 from mindspore.common.tensor import 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.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
@@ -54,7 +54,7 @@ def _convert_type(types):
     """
     ms_types = []
     for np_type in types:
-        ms_type = pytype_to_dtype(np_type)
+        ms_type = _pytype_to_dtype(np_type)  # pylint:disable=protected-access
         ms_types.append(ms_type)
     return ms_types
 
@@ -131,7 +131,7 @@ def _construct_tensor_list(types, shapes, batch_expand_num=1):
                 new_shape += (item * batch_expand_num,)
             else:
                 new_shape += (item,)
-        tensor = Tensor(np.zeros(new_shape, dtype_to_nptype(type_)), dtype=type_)
+        tensor = Tensor(np.zeros(new_shape, _dtype_to_nptype(type_)), dtype=type_)  # pylint:disable=protected-access
         tensor.virtual_flag = True
         tensor_list.append(tensor)
     return tensor_list
@@ -344,15 +344,7 @@ def _get_layout_opt_shard(layout_obj, param_redundancy_dict):
     """Layout ckpt append opt shard."""
     for key, value in layout_obj.items():
         if value[5]:
-            world_groups = ("hccl_world_group", "nccl_world_group", "mccl_world_group")
-            if value[5] in world_groups:
-                opt_para_num = get_group_size()
-            elif "-" in value[5]:
-                opt_para_str = value[5].split("-")[0]
-                opt_para_num = int(opt_para_str)
-            else:
-                raise ValueError(f"For get_parameter_redundancy, the format of the parallel communication domain for "
-                                 f"the optimizer is incorrect.")
+            opt_para_num = get_group_size(value[5])
             param_redundancy_ranks = param_redundancy_dict.get(key)
             res = []
             for param_ranks in param_redundancy_ranks:
@@ -582,17 +574,12 @@ def _progress_bar(iterable, total=None):
         print_progress_bar(i)
 
 
-def _load_and_transform(path, name_map, load_func, transform_func=None):
+def _load_and_transform(path, name_map, load_func):
     """use load_func to load and use transform_func to convert"""
-    if load_func is not None:
-        param_dict = load_func(path)
-    else:
-        param_dict = path
+    param_dict = load_func(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
-        if transform_func is not None:
-            transform_dict[new_name] = transform_func(v, new_name)
-        else:
-            transform_dict[new_name] = v
+        transform_dict[new_name] = v
     return transform_dict

mindspore/train/amp.py

@@ -463,9 +463,6 @@ def auto_mixed_precision(network, amp_level="O0", dtype=mstype.float16):
     ``Addcdiv``, ``Addcmul``, ``Cross``, ``_PyboostCrossPrim``, ``Dot``, ``GridSampler2D``, ``GridSampler3D``,
     ``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>`_ .
-
     Note:
         - Repeatedly calling mixed-precision interfaces, such as `custom_mixed_precision` and `auto_mixed_precision`,
           can result in a larger network hierarchy and slower performance.
@@ -821,8 +818,10 @@ def get_white_list():
          <class 'mindspore.ops.operations.nn_ops.Conv2DTranspose'>,
          <class 'mindspore.ops.operations.nn_ops.Conv3DTranspose'>,
          <class 'mindspore.ops.operations.nn_ops.Conv2DBackpropInput'>,
-         <class 'mindspore.ops.operations.math_ops.MatMul'>, <class 'mindspore.ops.operations.math_ops.BatchMatMul'>,
-         <class 'mindspore.ops.operations.nn_ops.PReLU'>, <class 'mindspore.ops.operations.nn_ops.ReLU'>,
+         <class 'mindspore.ops.auto_generate.gen_ops_prim.MatMul'>,
+         <class 'mindspore.ops.auto_generate.gen_ops_prim.BatchMatMul'>,
+         <class 'mindspore.ops.auto_generate.gen_ops_prim.PReLU'>,
+         <class 'mindspore.ops.auto_generate.gen_ops_prim.ReLU'>,
          <class 'mindspore.ops.operations.math_ops.Ger'>]
     """
     white_list = AMP_WHITE_LIST.copy()
@@ -874,8 +873,8 @@ def custom_mixed_precision(network, *, white_list=None, black_list=None, dtype=m
             white list is not used.
         black_list (list[Cell], optional): Black list of custom mixed precision. Defaults: ``None`` , means
             black list is not used.
-        dtype (Type): The type used in lower precision calculations, can be ``mstype.float16`` or ``mstype.bfloat16`` ,
-            default: ``mstype.float16`` .
+        dtype (Type, optional): The type used in lower precision calculations, can be ``mstype.float16`` or
+            ``mstype.bfloat16`` , default: ``mstype.float16`` .
 
     Returns:
         network (Cell), A network supporting mixed precision.

mindspore/train/callback/_callback.py

@@ -60,7 +60,8 @@ def _fill_param_into_net(net, parameter_list):
         if np_val.shape == (1,):
             parameter_dict[param_name] = Parameter(np_val, name=param_name)
         elif np_val.shape == ():
-            parameter_dict[param_name] = Parameter(Tensor(np_val.tolist(), mstype.pytype_to_dtype(np_val.dtype)),
+            # pylint:disable=protected-access
+            parameter_dict[param_name] = Parameter(Tensor(np_val.tolist(), mstype._pytype_to_dtype(np_val.dtype)),
                                                    name=param_name)
         else:
             parameter_dict[param_name] = Parameter(Tensor(np_val), name=param_name)

mindspore/train/callback/_checkpoint.py

@@ -27,7 +27,6 @@ from mindspore.train._utils import _make_directory
 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.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
@@ -509,9 +508,6 @@ class ModelCheckpoint(Callback):
         if callable(prefix):
             self._prefix_func = prefix
 
-        if context.get_context("device_target") == "GPU" and _get_recovery_context("enable_recovery"):
-            _set_recovery_context(ckpt_path=self._directory)
-
         if config is None:
             self._config = CheckpointConfig()
         else:
@@ -577,11 +573,6 @@ class ModelCheckpoint(Callback):
             self._directory = self._directory_func(cb_params)
             _make_directory(self._directory)
         collect_host_info("Callback", "ModelCheckpoint", "step_end", start_time=get_clock_syscnt(), level=1)
-        # In disaster recovery scenario, the training process may be rolled back to the last step where
-        # the ckpt was successfully saved, so the _last_triggered_step should be updated.
-        if _get_recovery_context("enable_recovery") and cb_params.last_save_ckpt_step is not None:
-            self._last_triggered_step = cb_params.last_save_ckpt_step
-            cb_params.last_save_ckpt_step = None
 
         # save graph (only once)
         if not self._graph_saved:
@@ -628,13 +619,6 @@ 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:
@@ -645,7 +629,7 @@ class ModelCheckpoint(Callback):
             self._flush_from_cache(cb_params)
 
         save_ckpt = self._check_save_ckpt(cb_params, force_to_save)
-        step_num_in_epoch = self._update_save_step(cb_params)
+        step_num_in_epoch = int((cb_params.cur_step_num - 1) % cb_params.batch_num + 1)
 
         if save_ckpt:
 

mindspore/train/callback/_flops_collector.py

@@ -31,7 +31,6 @@ from mindspore.communication.management import (create_group, get_group_size,
 from mindspore.parallel._auto_parallel_context import auto_parallel_context
 from mindspore.ops import operations as P
 from mindspore.common import Tensor
-from mindspore import context
 import mindspore.nn as nn
 
 
@@ -152,16 +151,21 @@ class FlopsUtilizationCollector(Callback):
         """
         Check whether FlopsUtilizationCollector is working in the current environment
         """
-        if context.get_context("mode") != context.GRAPH_MODE:
-            if self.verbose:
-                raise ValueError("FlopsUtilizationCollector now only support graph mode.")
-            logger.info("FlopsUtilizationCollector now only support graph mode.")
-            return False
         cb_params = run_context.original_args()
         if cb_params.mode == 'train':
             network = cb_params.train_network
+            if not network.compiled:
+                if self.verbose:
+                    raise ValueError("FlopsUtilizationCollector now only support graph mode.")
+                logger.info("FlopsUtilizationCollector now only support graph mode.")
+                return False
         elif cb_params.mode == 'eval':
             network = cb_params.eval_network
+            if not network.compiled:
+                if self.verbose:
+                    raise ValueError("FlopsUtilizationCollector now only support graph mode.")
+                logger.info("FlopsUtilizationCollector now only support graph mode.")
+                return False
         else:
             if self.verbose:
                 raise ValueError('FlopsUtilizationCollector only support train and eval mode!')

mindspore/train/callback/_train_fault_tolerance.py

@@ -28,15 +28,15 @@ from mindspore._c_expression import _repair_device, _stop_device, _tft_sem_post,
 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 send_recv, reset_params, direct_copy_to_host
+from mindspore._c_expression import _reg_snapshot_params, _reset_snapshot_state, _clear_snapshot_saving_flag
 from mindspore._c_expression import CollectiveManager
 from mindspore._c_expression import _get_uce_process_strategy, _get_uce_mem_info
-from mindspore._c_expression import TensorPy as Tensor_
 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
+from mindspore._c_expression import set_is_arf
 
 
 def _get_ckpt_dir(step, ckpt_save_path, is_tmp_file):
@@ -157,6 +157,7 @@ def _tft_clean_callback(is_uce_error, args, ctx):
     CollectiveManager.get_instance().resume_hccl_comm()
     logger.warning("Finish _tft_clean_callback, ret: {}".format(ret))
     if ctx.tft.tft_get_repair_type() == "recover":
+        _reset_snapshot_state()
         logger.warning(f"Destroy hcom")
         _finalize_comm()
         logger.warning(f"Destroy hcom end")
@@ -166,11 +167,10 @@ def _tft_clean_callback(is_uce_error, args, ctx):
 def _tft_stop_callback(args, cb_ctx):
     """ Callback used for TFT stop function."""
     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
+    _stop_device(cb_ctx.device_id)
     if cb_ctx.tft.tft_get_repair_type() == "recover":
         logger.warning(f"Reset limit step")
         cb_ctx.tft.tft_reset_limit_step()
@@ -182,7 +182,7 @@ def _tft_rebuild_sub_groups(fault_ranks, args, ctx):
     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)
+    set_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")
@@ -192,7 +192,7 @@ def _tft_rebuild_sub_groups(fault_ranks, args, ctx):
 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:
@@ -202,7 +202,10 @@ class TrainFaultTolerance(Callback):
         ckpt_save_path (str): Checkpoint save directory when failure occurs. When saved,
             a new directory named 'ttp_saved_checkpoints-step_{cur_step_num}'
             is created in that directory. Default: ``None``.
-        kwargs (dict): Other dictionary type parameters.
+        kwargs (dict): Other dictionary type parameters. When argument `ckpt_save_path` is ``None``, `kwargs` must
+            provide a parameter named `ckpt_save_fn`, which points to a function used to save checkpoint. The
+            prototype of `ckpt_save_fn` is ``def save_ckpt(cb_params, append_dict)``. When both `ckpt_save_path`
+            and `ckpt_save_fn` are provided, `ckpt_save_fn` is used in priority.
 
     Raises:
         Exception: TFT init failed.
@@ -329,7 +332,7 @@ class TrainFaultTolerance(Callback):
         # `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():
+        if self._only_enable_tre() or self._only_enable_ckpt_d2h_async():
             return
         self.tft = _tft_handler.get_tft()
         self._check_init()
@@ -340,11 +343,9 @@ class TrainFaultTolerance(Callback):
         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))
-        self.s1 = mindspore.hal.Stream()
+        self.g_one = Tensor([1], dtype=mstype.int32)
         _tft_sem_enable()
         self._tft_register()
 
@@ -354,7 +355,21 @@ class TrainFaultTolerance(Callback):
         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
+        return "TRE:1" in env_enable or "TRE:2" in env_enable
+
+    @staticmethod
+    def _only_enable_ckpt_d2h_async():
+        """Check whether only set MS_ENABLE_CKPT_D2H_ASYNC=1 without setting MS_ENABLE_TFT"""
+        if os.getenv("MS_ENABLE_TFT", "") != "":
+            return False
+        return os.getenv("MS_ENABLE_CKPT_D2H_ASYNC") == "1"
+
+    @staticmethod
+    def _enable_snapshot():
+        """Check whether parameter snapshot enabled"""
+        enable_step_tre = "TRE:2" in os.getenv("MS_ENABLE_TFT", "")
+        enable_ckpt_d2h_async = os.getenv("MS_ENABLE_CKPT_D2H_ASYNC") == "1"
+        return enable_step_tre or enable_ckpt_d2h_async
 
     def _only_enable_tsp(self):
         """Check if only configured MS_ENABLE_TFT='{TSP:1}'"""
@@ -382,18 +397,14 @@ 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():
             if "tft_g_one_flag" in key:
-                with mindspore.hal.StreamCtx(self.s1):
-                    tft_g_one_flag = Tensor(Tensor_.move_to(param, "CPU", False))
-                self.s1.synchronize()
+                tft_g_one_flag = direct_copy_to_host(param)
                 return int(tft_g_one_flag) == 1
         return False
 
@@ -438,7 +449,7 @@ class TrainFaultTolerance(Callback):
                 super(TFTOptSubCls, self).__init__(*args, **kwargs)
                 self.report = TensorReport()
                 self.report_end = TensorReport()
-                self.report_end.add_prim_attr("side_effect_mem", True).add_prim_attr("optimizer_end", True)
+                self.report_end.add_prim_attr("optimizer_end", True)
                 self.depend = ops.Depend()
                 self.allreduce_sum = ops.AllReduce()
                 self.allreduce_sum.add_prim_attr("tft_report_before", True)
@@ -452,7 +463,27 @@ class TrainFaultTolerance(Callback):
                 self.report_end("tft_report", self.tft_g_one_flag)
                 return opt_ret
 
-        return TFTOptSubCls
+        class TFTOptSnapShotCls(origin_opt_cls):
+            """
+            Optimizer wrapper class when using tft.
+            """
+
+            def __init__(self, *args, **kwargs):
+                super(TFTOptSnapShotCls, self).__init__(*args, **kwargs)
+                self.report = TensorReport()
+                self.report.add_prim_attr("side_effect_mem", True).add_prim_attr("snapshot", True)
+                self.dummy_input = Tensor([1], dtype=mstype.int32)
+
+            def construct(self, gradients, **kwargs):
+                """Add fake op TensorReport to insert wait event for copying parameters"""
+                self.report("tft_report", self.dummy_input)
+                opt_ret = super(TFTOptSnapShotCls, self).construct(gradients, **kwargs)
+                return opt_ret
+
+        env_tft = os.getenv('MS_ENABLE_TFT', '')
+        features = ['TTP:1', 'UCE:1', 'ARF:1']
+        need_redundancy = any([env_tft.find(feat) >= 0 for feat in features])
+        return TFTOptSubCls if need_redundancy else TFTOptSnapShotCls
 
     def _tft_register(self):
         """Register callback functions."""
@@ -480,6 +511,17 @@ class TrainFaultTolerance(Callback):
             _clean_rootinfo()
             self.clean_unique_id = True
 
+    def on_train_step_begin(self, run_context):
+        """
+        Clear saving snapshot state at each step begin.
+
+        Args:
+            run_context (RunContext): Context of the train running. Refer to
+                                      :class:`mindspore.train.RunContext` for detail.
+        """
+        if self._enable_snapshot():
+            _clear_snapshot_saving_flag()
+
     def on_train_step_end(self, run_context):
         """
         Report status to MindIO TFT after every step finished.
@@ -488,7 +530,7 @@ 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_ckpt_d2h_async():
             return
 
         cb_params = run_context.original_args()
@@ -528,10 +570,15 @@ class TrainFaultTolerance(Callback):
             run_context (RunContext): Context of the train running. Refer to
                                       :class:`mindspore.train.RunContext` for detail.
         """
+        cb_params = run_context.original_args()
+        if self._enable_snapshot():
+            param_dict = {}
+            for param in cb_params.train_network.trainable_params():
+                param_dict[param.name] = param
+            _reg_snapshot_params(param_dict)
         if self._only_enable_tsp():
             return
-        cb_params = run_context.original_args()
-        if self._only_enable_tre():
+        if self._only_enable_tre() or self._only_enable_ckpt_d2h_async():
             self.cb_params = cb_params
             return
         sink_size = cb_params.get("sink_size", 0)

mindspore/train/data_sink.py

@@ -16,7 +16,7 @@
 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, enable_data_broadcast
 from mindspore.train.dataset_helper import _has_dynamic_shape, _check_inputs
@@ -61,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)
@@ -94,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)
@@ -238,12 +238,8 @@ def data_sink(fn, dataset, sink_size=1, jit_config=None, input_signature=None):
 
         real_sink_fun = _get_sink_fun(sink_fun, key_info, is_info_queue, dataset, jit_config)
 
-        loop = sink_size
-        if jit_config is not None and context.get_context('mode') == context.GRAPH_MODE:
-            loop = 1
-
         out = None
-        for _ in range(loop):
+        for _ in range(sink_size):
             out = real_sink_fun()
 
         return out

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)