mindspore 2.5.0__cp310-cp310-win_amd64.whl → 2.6.0__cp310-cp310-win_amd64.whl

mindspore/train/callback/_summary_collector.py

@@ -107,22 +107,18 @@ class SummaryCollector(Callback):
               The first output will be treated as the loss and it will be averaged. Default: ``True`` .
             - collect_graph (bool): Whether to collect the computational graph. Currently, only
               training computational graph is collected. Default: ``True`` .
-            - collect_train_lineage (bool): Whether to collect lineage data for the training phase,
-              this field will be displayed on the `lineage page \
-              <https://www.mindspore.cn/mindinsight/docs/en/master/lineage_and_scalars_comparison.html>`_
-              of MindInsight. Default: ``True`` .
-            - collect_eval_lineage (bool): Whether to collect lineage data for the evaluation phase,
-              this field will be displayed on the `lineage page
-              <https://www.mindspore.cn/mindinsight/docs/en/master/lineage_and_scalars_comparison.html>`_
-              of MindInsight. Default: ``True`` .
+            - collect_train_lineage (bool): Whether to collect lineage data for the training phase.
+              Default: ``True`` .
+            - collect_eval_lineage (bool): Whether to collect lineage data for the evaluation phase.
+              Default: ``True`` .
             - collect_input_data (bool): Whether to collect dataset for each training.
               Currently only image data is supported.
               If there are multiple columns of data in the dataset, the first column should be image data.
               Default: ``True`` .
             - collect_dataset_graph (bool): Whether to collect dataset graph for the training phase.
               Default: ``True`` .
-            - histogram_regular (Union[str, None]): Collect weight and bias for parameter distribution page
-              and displayed in MindInsight. This field allows regular strings to control which parameters to collect.
+            - histogram_regular (Union[str, None]): Collect weight and bias for parameter distribution page.
+              This field allows regular strings to control which parameters to collect.
               It is not recommended to collect too many parameters at once, as it can affect performance.
               Note that if you collect too many parameters and run out of memory, the training will fail.
               Default: ``None`` , it means only the first five parameters are collected.
@@ -153,8 +149,7 @@ class SummaryCollector(Callback):
             True: it means that after specified data is set, non-specified data is collected as the default behavior.
             False: it means that after specified data is set, only the specified data is collected,
             and the others are not collected. Default: ``True`` .
-        custom_lineage_data (Union[dict, None]): Allows you to customize the data and present it on the MingInsight
-            `lineage page <https://www.mindspore.cn/mindinsight/docs/en/master/lineage_and_scalars_comparison.html>`_ .
+        custom_lineage_data (Union[dict, None]): Allows you to customize the data.
             In the custom data, the type of the key supports str, and the type of value supports str, int
             and float. Default: ``None`` , it means there is no custom data.
         collect_tensor_freq (Optional[int]): The same semantics as the `collect_freq`, but controls TensorSummary only.
@@ -168,7 +163,7 @@ class SummaryCollector(Callback):
             affect the number of steps TensorSummary will be collected.
             Default: ``None`` , which means to follow the behavior as described above.
         max_file_size (Optional[int]): The maximum size in bytes of each file that can be written to the disk.
-            For example, to write not larger than 4GB, specify `max_file_size=4*1024**3`.
+            For example, to write not larger than 4GB, specify `max_file_size=4*1024*3`.
             Default: ``None`` , which means no limit.
         export_options (Union[None, dict]): Perform custom operations on the export data.
             Note that the size of export files is not limited by the max_file_size.

mindspore/train/callback/_time_monitor.py

@@ -28,7 +28,8 @@ class TimeMonitor(Callback):
     Args:
         data_size (int): How many steps are the intervals between print information each time.
             if the program get `batch_num` during training, `data_size` will be set to `batch_num`,
-            otherwise `data_size` will be used. Default: ``None`` .
+            otherwise `data_size` will be used. If the program does not get `batch_num` during training,
+            meanwhile `data_size` does not set, the program will report an error. Default: ``None`` .
 
         data_time (bool): Whether to show the average time of fetching data in Host.
             Note that data fetch and network compute are processed sequentially in non dataset sink mode, while

mindspore/train/callback/{_tft_register.py → _train_fault_tolerance.py}

@@ -15,24 +15,27 @@
 """Checkpoint related classes and functions."""
 
 import os
+from mindspore.utils import _tft_handler
 from mindspore.train.serialization import save_checkpoint
-from mindspore.parallel._utils import _get_device_num
-from mindspore import _checkparam as Validator
 from mindspore.train.callback._callback import Callback
-from mindspore import context
+from mindspore import context, ops
 from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
 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 clean_tdt_channel
 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
-from mindspore._c_expression import Tensor as Tensor_
+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
+
 
 def _get_ckpt_dir(step, ckpt_save_path, is_tmp_file):
     """ Common func to generate ckpt dir name."""
@@ -40,30 +43,38 @@ def _get_ckpt_dir(step, ckpt_save_path, is_tmp_file):
     mid_dir = f"tft_saved_checkpoints-step_{str(step)}{tmp}"
     return os.path.join(ckpt_save_path, mid_dir)
 
+
 def _save_checkpoint_on_failure(step, save_info, args, cb_ctx):
     """ Callback used for TFT save ckpt function when errors occur."""
     logger.info("Enter _save_checkpoint_on_failure function")
-    if not cb_ctx._is_params_consistent():    # pylint: disable=W0212
+    if not cb_ctx._is_params_consistent():  # pylint: disable=W0212
         raise RuntimeError("Can't save parameters, because they are left in inconsistent state!")
+    cb_params = args
+    # we record the current step and epoch num in on_train_step_end, so we can just reset it here
+    cb_params.cur_step_num = cb_ctx.cur_step_num
+    cb_params.cur_epoch_num = cb_ctx.cur_epoch_num
+    if cb_params.optimizer is not None:
+        cb_params.optimizer.global_step = cb_ctx.global_step
+    if hasattr(cb_params.network, 'optimizer') and cb_params.network.optimizer is not None:
+        cb_params.network.optimizer.global_step = cb_ctx.global_step
+    append_dict = {}
+    append_dict["__exception_save__"] = True
+    # if user has provided a custom save callback, use it
+    if cb_ctx.save_cb:
+        cb_ctx.save_cb(cb_params, append_dict)
+        logger.info("Finish _save_checkpoint_on_failure function")
+        return
 
+    # if user has not provided a custom save callback, use default save logic
     ckpt_save_path = cb_ctx.ckpt_save_path
-    cb_params = args
     cur_rank = get_rank()
-    cur_step_num = cb_params.cur_step_num
+    step_num_in_epoch = int((cb_params.cur_step_num - 1) % cb_params.batch_num + 1)
     cur_epoch_num = cb_params.cur_epoch_num
-    batch_num = cb_params.batch_num
-    if cur_step_num > step:
-        cur_epoch_num = (step - 1) // batch_num + 1
-    step_num_in_epoch = int((step - 1) % batch_num + 1)
-
-    append_dict = {}
     append_dict["epoch_num"] = cur_epoch_num
-    append_dict["step_num"] = step
+    append_dict["step_num"] = cb_params.cur_step_num
     append_dict["cur_rank"] = cur_rank
-    append_dict["batch_num"] = batch_num
-    append_dict["__exception_save__"] = True
-
-    append_dict["global_step"] = Parameter([cb_ctx.global_step])
+    append_dict["batch_num"] = cb_params.batch_num
+    append_dict["global_step"] = cb_ctx.global_step
     outputs = cb_params.net_outputs
     if isinstance(outputs, (tuple, list)) and len(outputs) >= 3:
         append_dict["loss_scale"] = outputs[2]
@@ -76,49 +87,63 @@ def _save_checkpoint_on_failure(step, save_info, args, cb_ctx):
                     integrated_save=False, append_dict=append_dict)
     logger.info("Finish _save_checkpoint_on_failure function")
 
+
 def _rename_save_result(step, cb_ctx):
     """ Callback used for TFT rename function after ckpt save callback was finished and successful."""
     logger.info("Enter _rename_save_result function")
+    if cb_ctx.save_cb:
+        logger.info("User's save callback is provided, skip rename")
+        return
     tmp_dir = _get_ckpt_dir(step, cb_ctx.ckpt_save_path, True)
     fin_dir = _get_ckpt_dir(step, cb_ctx.ckpt_save_path, False)
 
     os.rename(tmp_dir, fin_dir)
     logger.info("Finish _rename_save_result function")
 
+
 def _tft_exit_cb(ctx):
+    """Callback used for TFT exit function."""
     logger.error("Enter mindio ttp exit process, which means other ranks occur exception, check other ranks' logs!")
     _tft_sem_post()
-    os._exit(1)   # pylint: disable=W0212
+    os._exit(1)  # pylint: disable=W0212
 
 
 def _tft_repair_callback(step, need_rebuild, error_ranks, repair_info, args, cb_ctx):
     """ Callback used for TFT repair function."""
-    logger.info("Enter _tft_repair_callback repair type: {}".format(repair_info["repair_type"]))
-    if(repair_info["repair_type"] == cb_ctx.tft.RepairType.RT_UCE_HIGHLEVEL.value\
-or repair_info["repair_type"] == cb_ctx.tft.RepairType.RT_UCE_LOWLEVEL.value):
-        logger.info("Enter _tft_repair_callback uce REPARI_DEVICE device_id : {}".format(cb_ctx.device_id))
+    logger.warning("Enter _tft_repair_callback repair type: {}".format(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))
         _repair_device(cb_ctx.device_id)
 
-    if(repair_info["repair_type"] == cb_ctx.tft.RepairType.RT_UCE_HIGHLEVEL.value\
-       or repair_info["repair_type"] == cb_ctx.tft.RepairType.RT_SEND.value):
-        logger.info("Enter _tft_repair_callback SEND_RECV repair type: \
-{}, src_rank:{}, dst_rank: {}".format(repair_info["repair_type"], repair_info["src"], repair_info["dst"]))
+    if (repair_info["repair_type"] in (cb_ctx.tft.RepairType.RT_UCE_HIGHLEVEL.value,
+                                       cb_ctx.tft.RepairType.RT_SEND.value,
+                                       cb_ctx.tft.RepairType.RT_RECV_REPAIR.value)):
+        logger.warning("Enter _tft_repair_callback SEND_RECV repair type:{}, src_rank:{}, dst_rank: {}".format(
+            repair_info["repair_type"], repair_info["src"], repair_info["dst"]))
         cb_params = args
-        src_rank = repair_info["src"][0]
-        dst_rank = repair_info["dst"][0]
-        if send_recv(cb_params.train_network.trainable_params(), src_rank, dst_rank) != 0:
-            raise ValueError("Call send_recv failed.")
-    logger.info("Finish _tft_repair_callback")
+        if repair_info["repair_type"] == cb_ctx.tft.RepairType.RT_SEND.value:
+            for i in range(len(repair_info["src"])):
+                src_rank = repair_info["src"][i]
+                dst_rank = repair_info["dst"][i]
+                if send_recv(cb_params.train_network.trainable_params(), src_rank, dst_rank) != 0:
+                    raise ValueError("Call send_recv failed.")
+        else:
+            src_rank = repair_info["src"][0]
+            dst_rank = repair_info["dst"][0]
+            if send_recv(cb_params.train_network.trainable_params(), src_rank, dst_rank) != 0:
+                raise ValueError("Call send_recv failed.")
+    logger.warning("Finish _tft_repair_callback")
 
 
 def _tft_clean_callback(is_uce_error, args, ctx):
     """ Callback used for TFT clean function."""
-    logger.info("Enter _tft_clean_callback")
+    logger.warning("Enter _tft_clean_callback")
     ret = 0
     if is_uce_error:
         _get_uce_mem_info(ctx.device_id)
         err_strategy = _get_uce_process_strategy()
-        logger.info("_tft_clean_callback err_strategy: {}".format(err_strategy))
+        logger.warning("_tft_clean_callback err_strategy: {}".format(err_strategy))
         if err_strategy == "RS_UCE_HIGHLEVEL":
             ret = 0
         elif err_strategy == "RS_UCE_LOWLEVEL":
@@ -126,37 +151,49 @@ def _tft_clean_callback(is_uce_error, args, ctx):
         else:
             ret = 1
     clean_tdt_channel()
-    logger.info("Enter _tft_clean_callback resume_hccl_comm")
+    logger.warning("Enter _tft_clean_callback resume_hccl_comm")
     CollectiveManager.get_instance().resume_hccl_comm()
-    logger.info("Finish _tft_clean_callback, ret: {}".format(ret))
+    logger.warning("Finish _tft_clean_callback, ret: {}".format(ret))
     return ret
 
 
 def _tft_stop_callback(args, cb_ctx):
     """ Callback used for TFT stop function."""
-    logger.info("Enter _tft_stop_callback device_id: {}".format(cb_ctx.device_id))
+    logger.warning("Enter _tft_stop_callback device_id: {}".format(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
+    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")
 
 
-class TFTRegister(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()
+    _rebuild_world_group()
+    _rebuild_sub_group()
+    _set_recovery_context(is_arf=True)
+    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>`_.
-    This callback will execute TFT operations during training process, such as TFT init, report and exception handle.
+    `MindIO TFT <https://www.hiascend.com/document/detail/zh/mindx-dl/60rc2/mindio/mindiottp/mindiottp001.html>`_
+    and will execute TFT operations during training process, such as TFT init, report and exception handle.
 
     Note:
         Required for Ascend graph mode only. And sink size must be less than or equal to 1.
 
     Args:
-        ctrl_rank_id (int): TFT controller's running rank_id, used for init TFT controller.
-        ctrl_ip (str): TFT controller's ip address, used for init TFT controller.
-        ctrl_port (int): TFT controller's ip port, used for init TFT controller and processor.
-        ckpt_save_path (str): Checkpoint save directory when failure occurs, checkpoint file will save to directory
-           named ttp_saved_checkpoints-step_{cur_step_num} under this directory.
+        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.
 
     Raises:
         Exception: TFT init failed.
@@ -168,7 +205,7 @@ class TFTRegister(Callback):
 
             It's recommended to use the msrun startup method.
             Please see the `msrun start up
-            <https://www.mindspore.cn/docs/en/master/model_train/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/tutorials/en/master/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 4 devices.
@@ -181,7 +218,7 @@ class TFTRegister(Callback):
         >>> from mindspore import nn, ops, Parameter, train
         >>> from mindspore.communication import init, get_rank
         >>> from mindspore.common.initializer import initializer, HeUniform
-        >>> from mindspore.train import Model, TFTRegister
+        >>> from mindspore.train import Model, TrainFaultTolerance
         >>> from mindspore import dataset as ds
         >>> ms.set_context(mode=ms.GRAPH_MODE, jit_level='O2')
         >>> ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.SEMI_AUTO_PARALLEL, pipeline_stages=2)
@@ -252,43 +289,68 @@ class TFTRegister(Callback):
         >>> optimizer_wrapper = nn.OptTFTWrapper(optimizer)
         >>> loss_fn = nn.CrossEntropyLoss()
         >>>
-        >>> net_with_loss = nn.PipelineCell(nn.WithLossCell(net, loss_fn), 4)
+        >>> net_with_loss = nn.Pipeline(nn.WithLossCell(net, loss_fn), 4)
         >>> net_with_loss.set_train()
         >>> model = Model(net_with_loss, optimizer=optimizer_wrapper)
-        >>> tft_cb = TFTRegister(0, "192.168.0.1", 2000, "./tft_checkpoint/")
+        >>> tft_cb = TrainFaultTolerance()
         >>> loss_cb = train.LossMonitor(1)
         >>> model.train(1, dataset, callbacks=[tft_cb, loss_cb])
     """
 
-    def __init__(self, ctrl_rank_id, ctrl_ip, ctrl_port, ckpt_save_path):
-        super(TFTRegister, self).__init__()
-
-        tft_env = os.getenv("MS_ENABLE_TFT", "")
-        if ("TTP:1" not in tft_env) and ("UCE:1" not in tft_env):
-            raise ValueError("MindIO TFT regitster need custom switch on[MS_ENABLE_TFT='{TTP:1,UCE:1}']!")
-        mode = context.get_context("mode")
-        device_target = context.get_context("device_target")
-        if device_target != "Ascend" or mode != context.GRAPH_MODE:
-            raise ValueError("MindIO adataper only support on Ascend device with GRAPH Mode!")
-
-        # let it raise errors if not install mindio_tft package
-        from mindio_ttp import framework_ttp as tft
-        self.tft = tft
-        self.global_step = 0
-        Validator.check_non_negative_int(ctrl_port)
-        self.has_init_replica = False
-        self.is_uce_rank = False
-        self._controller_ip = ctrl_ip
-        self._controller_rank_id = ctrl_rank_id
-        self._controller_port = ctrl_port
+    def __init__(self, ckpt_save_path=None, **kwargs):
+        super(TrainFaultTolerance, self).__init__()
+        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._init_tft()
-        self.ckpt_save_path = ckpt_save_path
+        self.cur_step_num = 0
+        self.cur_epoch_num = 0
+        # 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._check_init()
+        self.global_step = None
+        self.learning_rate = None
+        self.has_init_replica = False
+        self.is_uce_rank = False
+
         self.assign = mindspore.ops.Assign()
         self.g_one = Parameter(Tensor([1], dtype=mstype.int32))
         self.s1 = mindspore.hal.Stream()
         _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 _check_init(self):
+        """Check if the mindio-ttp had inited"""
+        if self.tft is None:
+            tft_env = os.getenv("MS_ENABLE_TFT", "")
+            if "ARF:1" in tft_env:
+                raise ValueError("Must init by _tft_handler.init(config=params) if use ARF.")
+            logger.warning(f"TFT handle not init, try to init")
+            _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}")
 
     def _is_params_consistent(self):
         for key, param in self.cb_params.train_network.parameters_and_names():
@@ -300,7 +362,7 @@ class TFTRegister(Callback):
         return False
 
     def _set_tft_optimizer_replica(self, run_context):
-        """ set Mindio TFT optimizer replica info, used internal. """
+        """ Set Mindio TFT optimizer replica info, used internal. """
         cur_rank = get_rank()
         cb_params = run_context.original_args()
         train_network = cb_params.train_network
@@ -322,33 +384,49 @@ class TFTRegister(Callback):
         ]
         self.tft.tft_set_optimizer_replica(cur_rank, replica_info)
 
-    def _init_tft(self):
-        """ Init Mindio TFT, used internal. """
-        logger.info("Begin to init tft.")
+    @classmethod
+    def get_optimizer_wrapper(cls, origin_opt_cls):
+        """
+        Optimizer wrapper func when using tft.
+
+        Args:
+            origin_opt_cls (Class): origin optimizer class.
+        """
+
+        class TFTOptSubCls(origin_opt_cls):
+            """
+            Optimizer wrapper class when using tft.
+            """
+
+            def __init__(self, *args, **kwargs):
+                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.depend = ops.Depend()
+                self.allreduce_sum = ops.AllReduce()
+                self.allreduce_sum.add_prim_attr("tft_report_before", True)
+                self.tft_g_one_flag = Parameter(Tensor([1], dtype=mstype.int32))
+
+            def construct(self, gradients, **kwargs):
+                tft_g_one_flag = self.depend(self.tft_g_one_flag, gradients)
+                self.tft_g_one_flag = self.allreduce_sum(tft_g_one_flag)
+                grads = self.depend(gradients, self.report("tft_report", self.tft_g_one_flag))
+                opt_ret = super(TFTOptSubCls, self).construct(grads, **kwargs)
+                self.report_end("tft_report", self.tft_g_one_flag)
+                return opt_ret
+
+        return TFTOptSubCls
+
+    def _tft_register(self):
+        """Register callback functions."""
         self.tft.tft_register_save_ckpt_handler(_save_checkpoint_on_failure, self)
         self.tft.tft_register_rename_handler(_rename_save_result, self)
         self.tft.tft_register_exit_handler(_tft_exit_cb, self)
         self.tft.tft_register_stop_handler(_tft_stop_callback, self)
         self.tft.tft_register_clean_handler(_tft_clean_callback, self)
         self.tft.tft_register_repair_handler(_tft_repair_callback, self)
-
-        world_size = _get_device_num()
-        cur_rank = get_rank()
-        enable_local_copy = False
-        enable_arf = False
-        enable_tls = False
-        tls_key_dir = ""
-
-        if cur_rank == self._controller_rank_id:
-            logger.info(f"Begin to start tft controller on rank_id:{cur_rank}")
-            self.tft.tft_init_controller(cur_rank, world_size, enable_local_copy, enable_arf)
-            self.tft.tft_start_controller(self._controller_ip, self._controller_port, enable_tls, tls_key_dir)
-            logger.info("Finish start tft controller.")
-
-        logger.info("Begin to start tft processor.")
-        self.tft.tft_init_processor(cur_rank, world_size, enable_local_copy, enable_tls, tls_key_dir)
-        self.tft.tft_start_processor(self._controller_ip, self._controller_port)
-        logger.info("Finished start tft processor.")
+        self.tft.tft_register_rebuild_group_handler(_tft_rebuild_sub_groups, self)
 
     def _reset_acc_grads(self):
         accu_grad_params = map(lambda e: e[1],
@@ -360,29 +438,44 @@ class TFTRegister(Callback):
 
     def on_train_step_end(self, run_context):
         """
-        And report status to MindIO TFT after every step finished.
+        Report status to MindIO TFT after every step finished.
 
         Args:
             run_context (RunContext): Context of the train running. Refer to
                                       :class:`mindspore.train.RunContext` for detail.
         """
+        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 cb_params.optimizer is not None:
-            self.global_step = int(cb_params.optimizer.global_step.data)
+            self.global_step = cb_params.optimizer.global_step.clone()
             self.assign(cb_params.optimizer.tft_g_one_flag, self.g_one)
-        else:
-            self.global_step = int(cb_params.network.optimizer.global_step.data)
+        elif hasattr(cb_params.network, 'optimizer') and cb_params.network.optimizer is not None:
+            self.global_step = cb_params.network.optimizer.global_step.clone()
             self.assign(cb_params.network.optimizer.tft_g_one_flag, self.g_one)
-        self.tft.tft_end_updating_os(cb_params.cur_step_num)
+        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)
         logger.info("END Set optimizer finish step status to TFT.")
 
-
     def on_train_begin(self, run_context):
+        """
+        Register train params to MindIO TFT on train beginning.
+
+        Args:
+            run_context (RunContext): Context of the train running. Refer to
+                                      :class:`mindspore.train.RunContext` for detail.
+        """
         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.")
@@ -391,7 +484,13 @@ class TFTRegister(Callback):
         self.cb_params = cb_params
 
     def end(self, run_context):
-        cur_rank = get_rank()
-        if cur_rank == self._controller_rank_id:
-            self.tft.tft_destroy_controller()
-        self.tft.tft_destroy_processor()
+        """
+        Unregister MindIO TFT on train end.
+
+        Args:
+            run_context (RunContext): Context of the train running. Refer to
+                                      :class:`mindspore.train.RunContext` for detail.
+        """
+        if self._only_enable_tre():
+            return
+        _tft_handler.unregister_tft()

mindspore/train/data_sink.py

@@ -98,6 +98,29 @@ def _get_next_op(dataset, ori_next_op, is_info_queue):
     return next_op, (key, dataset_shapes, dataset_types)
 
 
+def _get_jit_func(sink_fun, jit_config):
+    """
+    Get the jit function.
+    """
+    jit_config_dict = jit_config.jit_config_dict
+    jit_level = jit_config_dict['jit_level']
+    if jit_level == "":
+        jit_level = "O0"
+    backend = ""
+    if jit_level == "O2":
+        jit_level = "O0"
+        backend = "GE"
+    if "backend" in jit_config_dict:
+        backend = jit_config_dict["backend"]
+    fullgraph = False
+    if jit_config_dict['jit_syntax_level'] == "STRICT":
+        fullgraph = True
+    exc_mode = jit_config_dict['exc_mode']
+    infer_boost = jit_config_dict['infer_boost']
+    return jit(sink_fun, jit_level=jit_level, backend=backend, fullgraph=fullgraph, exc_mode=exc_mode,
+               infer_boost=infer_boost)
+
+
 def _get_sink_fun(sink_fun, key_info, is_info_queue, dataset, jit_config):
     """
     get the sink function.
@@ -107,7 +130,7 @@ def _get_sink_fun(sink_fun, key_info, is_info_queue, dataset, jit_config):
             if jit_config is None:
                 dst_sink_fun = sink_fun
             else:
-                dst_sink_fun = jit(sink_fun, jit_config=jit_config)
+                dst_sink_fun = _get_jit_func(sink_fun, jit_config)
             dataset.__sink_fun__ = dst_sink_fun
 
         return dataset.__sink_fun__
@@ -119,7 +142,7 @@ def _get_sink_fun(sink_fun, key_info, is_info_queue, dataset, jit_config):
         if jit_config is None:
             dst_sink_fun = sink_fun
         else:
-            dst_sink_fun = jit(sink_fun, jit_config=jit_config)
+            dst_sink_fun = _get_jit_func(sink_fun, jit_config)
         dataset.__sink_aux__.sink_funcs[key] = dst_sink_fun
 
     return dst_sink_fun

mindspore/train/dataset_helper.py

@@ -214,8 +214,7 @@ def _get_dataset_aux(dataset):
 
 def connect_network_with_dataset(network, dataset_helper):
     """
-    Connect the `network` with dataset in `dataset_helper`. Only supported in `sink mode
-    <https://mindspore.cn/docs/en/master/model_train/train_process/train_optimize.html>`_,
+    Connect the `network` with dataset in `dataset_helper`. Only supported in sink mode,
     (dataset_sink_mode=True).
 
     Args:
@@ -335,11 +334,11 @@ class DatasetHelper:
         dataset_sink_mode (bool): If the value is True, GetNext is employed to fetch the data at device through the
                                   dataset pipeline, otherwise fetch the data at host by iterating through the dataset.
                                   Default: ``True``.
-        sink_size (int): Control the amount of data in each sink.
+        sink_size (int): Control the amount of data in each sink. Must be -1 or positive.
                           If sink_size=-1, sink the complete dataset for each epoch.
                           If sink_size>0, sink sink_size data for each epoch.
-                          Default: -1.
-        epoch_num (int): The number of passes of the entire dataset to be sent. Default: 1.
+                          Default: ``-1``.
+        epoch_num (int): The number of passes of the entire dataset to be sent. Default: ``1``.
 
     Examples:
         >>> import numpy as np