mindspore 2.4.1__cp311-cp311-manylinux1_x86_64.whl → 2.5.0__cp311-cp311-manylinux1_x86_64.whl

mindspore/train/callback/_checkpoint.py

@@ -18,14 +18,14 @@ from __future__ import absolute_import
 import os
 import stat
 import time
-import threading
 
 import mindspore.context as context
 from mindspore import log as logger
 from mindspore import nn
 from mindspore import _checkparam as Validator
 from mindspore.train._utils import _make_directory
-from mindspore.train.serialization import save_checkpoint, _save_graph
+from mindspore.train.serialization import save_checkpoint, _save_graph, _wait_async_process_save_ckpt, \
+    _wait_async_thread_save_ckpt, _check_async_save
 from mindspore.parallel._cell_wrapper import destroy_allgather_cell
 from mindspore.parallel._recovery_context import _set_recovery_context, _get_recovery_context
 from mindspore.parallel._auto_parallel_context import _get_auto_parallel_context
@@ -44,15 +44,6 @@ SAVE_DIR = _cur_dir
 _info_list = ["epoch_num", "step_num"]
 
 
-def _wait_async_save_ckpt(async_save=False):
-    """Waiting for asynchronous saving of ckpt to complete."""
-    if async_save:
-        thread_list = threading.enumerate()
-        for thread in thread_list:
-            if thread.getName() == "asyn_save_ckpt":
-                thread.join()
-
-
 def _get_dp_tp_from_redundancy(redundancy_tuple):
     """From redundancy get dp and tp"""
     dp = []
@@ -76,6 +67,15 @@ def _get_dp_tp_from_layout(parameter_redundancy_dict):
     return dp, tp
 
 
+def _wait_async_save_ckpt(async_save=False):
+    """Waiting for asynchronous saving of ckpt to complete."""
+    if async_save:
+        if async_save == "process":
+            _wait_async_process_save_ckpt()
+        else:
+            _wait_async_thread_save_ckpt()
+
+
 def _chg_ckpt_file_name_if_same_exist(directory, prefix, exception=False):
     """Check if there is a file with the same name."""
     if callable(prefix) or callable(directory):
@@ -139,7 +139,10 @@ class CheckpointConfig:
         integrated_save (bool): Whether to merge and save the split Tensor in the automatic parallel scenario.
             Integrated save function is only supported in automatic parallel scene, not supported
             in manual parallel. Default: ``True`` .
-        async_save (bool): Whether asynchronous execution saves the checkpoint to a file. Default: ``False`` .
+        async_save (Union[bool, str]):Whether to use asynchronous saving of the checkpoint file, if True,
+                                    the asynchronous thread is used by default. If the type is string,
+                                    the method of asynchronous saving, it can be "process" or "thread".
+                                    Default: ``False`` .
         saved_network (Cell): Network to be saved in checkpoint file. If the saved_network has no relation
             with the network in training, the initial value of saved_network will be saved. Default: ``None`` .
         append_info (list): The information save to checkpoint file. Support "epoch_num", "step_num" and
@@ -247,7 +250,7 @@ class CheckpointConfig:
                 self._keep_checkpoint_max = 1
 
         self._integrated_save = Validator.check_bool(integrated_save)
-        self._async_save = Validator.check_bool(async_save)
+        self._async_save = _check_async_save(async_save)
         self._saved_network = saved_network
         self._append_dict = self._handle_append_info(append_info)
         self._enc_key = Validator.check_isinstance('enc_key', enc_key, (type(None), bytes))
@@ -313,10 +316,10 @@ class CheckpointConfig:
     @property
     def async_save(self):
         """
-        Get the value of whether asynchronous execution saves the checkpoint to a file.
+        Get the value of whether or how asynchronous execution saves the checkpoint to a file.
 
         Returns:
-            bool, whether asynchronous execution saves the checkpoint to a file.
+            (bool, str), whether or how asynchronous execution saves the checkpoint to a file.
         """
         return self._async_save
 
@@ -538,6 +541,8 @@ class ModelCheckpoint(Callback):
         self._graph_saved = False
         self._need_flush_from_cache = True
         self._map_param_inc = self._config.map_param_inc
+        self._d2h_async = os.environ.get("MS_ENABLE_CKPT_D2H_ASYNC") == "1"
+        self._run_mode = context.get_context("mode")
 
     def step_end(self, run_context):
         """
@@ -632,6 +637,13 @@ class ModelCheckpoint(Callback):
         if "step_num" in self._append_dict:
             self._append_dict["step_num"] = self._append_step_num + step_num
 
+    def _update_save_step(self, cb_params):
+        """update step if used async d2h copy"""
+        step_num_in_epoch = int((cb_params.cur_step_num - 1) % cb_params.batch_num + 1)
+        if self._d2h_async and self._run_mode == context.GRAPH_MODE:
+            step_num_in_epoch -= 1
+        return step_num_in_epoch
+
     def _save_ckpt(self, cb_params, force_to_save=False):
         """Save checkpoint files."""
         if cb_params.cur_step_num == self._last_triggered_step:
@@ -642,10 +654,12 @@ class ModelCheckpoint(Callback):
             self._flush_from_cache(cb_params)
 
         save_ckpt = self._check_save_ckpt(cb_params, force_to_save)
-        step_num_in_epoch = int((cb_params.cur_step_num - 1) % cb_params.batch_num + 1)
+        step_num_in_epoch = self._update_save_step(cb_params)
 
         if save_ckpt:
+
             _wait_async_save_ckpt(self._config.async_save)
+
             if self._prefix_func:
                 cur_ckpoint_file = self._prefix + f".{self._config.format}"
             else:
@@ -704,14 +718,14 @@ class ModelCheckpoint(Callback):
                             f"For remove_redundancy save checkpoint, the saved parameters are non-redundant.")
 
                     def choice_func(x):
-                        return x not in param_layout_set or x in save_param_names
+                        return x not in param_layout_set or (save_param_names is not None and x in save_param_names)
                 else:
                     param_redundancy_dict = get_parameter_redundancy(network)
                     single_params = remove_param_redundancy(param_redundancy_dict)
                     save_param_names = single_params.get(rank_id)
 
                     def choice_func(x):
-                        return x in save_param_names
+                        return save_param_names is not None and x in save_param_names
                 save_checkpoint(network, cur_file, False, self._config.async_save,
                                 self._append_dict, self._config.enc_key, self._config.enc_mode,
                                 crc_check=self._config.crc_check, format=self._config.format,

mindspore/train/callback/_early_stop.py

@@ -198,7 +198,7 @@ class EarlyStopping(Callback):
         """
         Get the monitor value at the end of epoch during training.
 
-        If `mindspore.train.callback.ReduceLROnPlateau` used with `model.train`, no evaluation process
+        If :class:`mindspore.train.callback.ReduceLROnPlateau` used with `model.train`, no evaluation process
         during training, only monitor="loss" is valid; if it used with `model.fit`, evaluation process will be
         performed at the end of epoch, valid monitor is "loss", "eval_loss" and metrics passed to `Model`.
 

mindspore/train/callback/_flops_collector.py

@@ -61,10 +61,13 @@ class FlopsUtilizationCollector(Callback):
         computility (int): The peak flops of each compute card. Default: ``1`` .
         full_flops(bool): Whether to count the full model flops. If set full_flops to False,
             FlopsUtilizationCollector would count the shard model flops in each device. Default: ``True`` .
+        enable_ma_collector(bool): Whether to write flops into the log and provide them to tasks
+            on the cloud for retrieval. Default: ``False`` .
 
     Raises:
         TypeError: If data_size is not positive int.
         TypeError: If full_flops is not bool.
+        TypeError: If enable_ma_collector is not bool.
         AssertionError: If the training mode is not a static graph or not a static shape.
 
     Examples:
@@ -86,7 +89,7 @@ class FlopsUtilizationCollector(Callback):
         Train per step time: 135.572 ms, mfu:0.47% hfu:0.47%
         Train per step time: 1.317 ms, mfu:48.59% hfu:48.59%
     """
-    def __init__(self, data_size, computility=1, full_flops=True):
+    def __init__(self, data_size=None, computility=1, full_flops=True, enable_ma_collector=False):
         super(FlopsUtilizationCollector, self).__init__()
         self.step_time = time.time()
         self.computility = computility
@@ -101,10 +104,14 @@ class FlopsUtilizationCollector(Callback):
         self.mfu_calculated = False
         self.data_size = data_size
         self.time_step_path = ''
-        self.verbose = full_flops
-        self.ma = os.environ.get("ENABLE_FLOPS_UTILIZATION_COLLECTOR") == "1"
-        Validator.check_bool(full_flops, "verbose")
-        Validator.check_positive_int(data_size, "data_size")
+        self.full_flops = full_flops
+        self.verbose = not(computility == 1 and enable_ma_collector)
+        self.ma = enable_ma_collector
+        self.batch_step_size = None
+        Validator.check_bool(full_flops, "full_flops")
+        Validator.check_bool(enable_ma_collector, "enable_ma_collector")
+        if data_size:
+            Validator.check_positive_int(data_size, "data_size")
 
     def step_begin(self, run_context):
         """
@@ -115,6 +122,14 @@ class FlopsUtilizationCollector(Callback):
             run_context (RunContext): Context of the process running. For more details,
                     please refer to :class:`mindspore.train.RunContext`.
         """
+        if self.batch_step_size is None:
+            self.batch_step_size = self.data_size
+            cb_params = run_context.original_args()
+            if hasattr(cb_params, "batch_num"):
+                batch_num = cb_params.batch_num
+                if isinstance(batch_num, int) and batch_num > 0:
+                    self.batch_step_size = cb_params.batch_num
+            Validator.check_positive_int(self.batch_step_size)
         self.step_time = time.time()
 
     def _get_pipeline_group(self):
@@ -134,6 +149,40 @@ class FlopsUtilizationCollector(Callback):
         rank_list_str = "-".join(rank_str_list)
         return rank_list, rank_list_str
 
+    def _check_run_mode_valid(self, run_context):
+        """
+        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
+        elif cb_params.mode == 'eval':
+            network = cb_params.eval_network
+        else:
+            if self.verbose:
+                raise ValueError('FlopsUtilizationCollector only support train and eval mode!')
+            logger.info('FlopsUtilizationCollector only support train and eval mode!')
+            return False
+        try:
+            self.full_model_flops, self.full_hardware_flops, self.shard_model_flops, \
+            self.shard_hardware_flops, is_dynamic_shape = flops_collection(network.current_phase)
+        except Exception as e:
+            if self.verbose:
+                raise ValueError("FlopsUtilizationCollector is not supported because {}.".format(e))
+            logger.info("FlopsUtilizationCollector is not supported because {}.".format(e))
+            return False
+        if is_dynamic_shape:
+            if self.verbose:
+                raise ValueError("FlopsUtilizationCollector now do not support dynamic shape.")
+            logger.info("FlopsUtilizationCollector now do not support dynamic shape.")
+            return False
+        return True
+
     def step_end(self, run_context):
         """
         Print mfu and hfu time at the end of step.
@@ -142,84 +191,67 @@ class FlopsUtilizationCollector(Callback):
            run_context (RunContext): Context of the process running. For more details,
                    please refer to :class:`mindspore.train.RunContext`.
         """
-        if context.get_context("mode") != context.GRAPH_MODE:
-            logger.warning("FlopsUtilizationCollector now only support graph mode.")
-            return
-
         step_seconds = (time.time() - self.step_time) * 1000
         if not self.mfu_calculated:
-            cb_params = run_context.original_args()
-            if cb_params.mode == 'train':
-                network = cb_params.train_network
-            elif cb_params.mode == 'eval':
-                network = cb_params.eval_network
-            else:
-                logger.warning('FlopsUtilizationCollector only support train and eval mode!')
+            if not self._check_run_mode_valid(run_context):
                 return
-            full_model_flops, full_hardware_flops, shard_model_flops, \
-                shard_hardware_flops, is_dynamic_shape = flops_collection(network.current_phase)
-            if is_dynamic_shape:
-                logger.warning("FlopsUtilizationCollector now do not support dynamic shape.")
-                return
-            self.full_mfu = full_model_flops / self.computility
-            self.full_hfu = full_hardware_flops / self.computility
-
-            self.shard_mfu = shard_model_flops / self.computility
-            self.shard_hfu = shard_hardware_flops / self.computility
-            self.full_model_flops = full_model_flops
-            self.full_hardware_flops = full_hardware_flops
-            self.shard_model_flops = shard_model_flops
-            self.shard_hardware_flops = shard_hardware_flops
+            self.full_mfu = self.full_model_flops / self.computility
+            self.full_hfu = self.full_hardware_flops / self.computility
+            self.shard_mfu = self.shard_model_flops / self.computility
+            self.shard_hfu = self.shard_hardware_flops / self.computility
             self.mfu_calculated = True
             shard_mf_dir = os.path.realpath(os.getenv('MA_LOG_DIR', './'))
             if self.ma:
-                flops_path = os.path.join(shard_mf_dir, "flops.txt")
-                self.time_step_path = os.path.join(shard_mf_dir, "time_step.txt")
-                if auto_parallel_context().get_parallel_mode() != "stand_alone":
-                    flops_path = os.path.join(
-                        shard_mf_dir, "flops_rank_" + str(get_rank())) + ".txt"
-                    self.time_step_path = os.path.join(
-                        shard_mf_dir, "time_step_rank_" + str(get_rank())) + ".txt"
+                rank_id = get_rank() if auto_parallel_context().get_parallel_mode() != "stand_alone" else 0
+                flops_path = os.path.join(
+                    shard_mf_dir, "flops_rank_" + str(rank_id)) + ".txt"
+                self.time_step_path = os.path.join(
+                    shard_mf_dir, "time_step_rank_" + str(rank_id)) + ".txt"
                 time_stamp = time.time()
-                model_flops_log = "flops{{type=\"model_flops\"}} {} {}\n".\
-                    format(shard_model_flops, int(round(time_stamp * 1000)))
-                hardware_flops_log = "flops{{type=\"hardware_flops\"}} {} {}\n".\
-                    format(shard_hardware_flops, int(round(time_stamp * 1000)))
+                model_flops_log = "flops{{type=\"model_flops\", rank_id=\"{}\"}} {} {}\n".\
+                    format(str(rank_id), self.shard_model_flops, int(round(time_stamp * 1000)))
+                hardware_flops_log = "flops{{type=\"hardware_flops\", rank_id=\"{}\"}} {} {}\n".\
+                    format(str(rank_id), self.shard_hardware_flops, int(round(time_stamp * 1000)))
                 flags = os.O_WRONLY | os.O_CREAT
                 modes = stat.S_IWUSR | stat.S_IRUSR
                 with os.fdopen(os.open(flops_path, flags, modes), 'w') as f:
                     f.write(model_flops_log)
                     f.write(hardware_flops_log)
             if self.verbose:
-                flops_log = f"Full model flops is {full_model_flops}, Full hardware flops is {full_hardware_flops}, " \
-                            f"Shard model flops is {shard_model_flops}, Shard hardware flops is {shard_hardware_flops}."
+                if self.full_flops:
+                    pipeline_num = auto_parallel_context().get_pipeline_stages()
+                    if pipeline_num > 1:
+                        pipeline_group_list, pipeline_group_name = self._get_pipeline_group()
+                        auto_parallel_context().set_pipeline_stages(1)
+                        hashed = hashlib.md5(
+                            pipeline_group_name.encode()).hexdigest()[:48]
+                        pipeline_group_name = str(hashed)
+                        create_group(pipeline_group_name, pipeline_group_list)
+                        self.full_mfu = AllReduceNet(pipeline_group_name)(
+                            Tensor([self.full_mfu])).asnumpy()[0]
+                        self.full_hfu = AllReduceNet(pipeline_group_name)(
+                            Tensor([self.full_hfu])).asnumpy()[0]
+                        auto_parallel_context().set_pipeline_stages(pipeline_num)
+                    full_model_flops = self.full_mfu * self.computility
+                    full_hardware_flops = self.full_hfu * self.computility
+                    if auto_parallel_context().get_parallel_mode() != "stand_alone":
+                        self.full_mfu = self.full_mfu / get_group_size()
+                        self.full_hfu = self.full_hfu / get_group_size()
+                    flops_log = f"Full model flops is {full_model_flops}, " \
+                                f"Full hardware flops is {full_hardware_flops}, " \
+                                f"Shard model flops is {self.shard_model_flops}, " \
+                                f"Shard hardware flops is {self.shard_hardware_flops}."
+                else:
+                    flops_log = f"Shard model flops is {self.shard_model_flops}, " \
+                                f"Shard hardware flops is {self.shard_hardware_flops}."
                 print(flops_log, flush=True)
-                if auto_parallel_context().get_pipeline_stages() > 1:
-                    pipeline_group_list, pipeline_group_name = self._get_pipeline_group()
-                    auto_parallel_context().set_pipeline_stages(1)
-                    hashed = hashlib.md5(
-                        pipeline_group_name.encode()).hexdigest()[:48]
-                    pipeline_group_name = str(hashed)
-                    create_group(pipeline_group_name, pipeline_group_list)
-                    self.full_mfu = AllReduceNet(pipeline_group_name)(
-                        Tensor([self.full_mfu])).asnumpy()[0]
-                    self.full_hfu = AllReduceNet(pipeline_group_name)(
-                        Tensor([self.full_hfu])).asnumpy()[0]
-                if auto_parallel_context().get_parallel_mode() != "stand_alone":
-                    self.full_mfu = self.full_mfu / get_group_size()
-                    self.full_hfu = self.full_hfu / get_group_size()
-
-        step_size = self.data_size
         cb_params = run_context.original_args()
-        if hasattr(cb_params, "batch_num"):
-            batch_num = cb_params.batch_num
-            if isinstance(batch_num, int) and batch_num > 0:
-                step_size = cb_params.batch_num
-        Validator.check_positive_int(step_size)
         if cb_params.dataset_sink_mode:
-            step_seconds = step_seconds / step_size
+            step_seconds = step_seconds / self.batch_step_size
         time_stamp = time.time()
-        train_log = "time_monitor{{type=\"per_step_time\"}} {} {}".format(step_seconds, int(round(time_stamp * 1000)))
+        rank_id = get_rank() if auto_parallel_context().get_parallel_mode() != "stand_alone" else 0
+        train_log = "time_monitor{{type=\"per_step_time\", rank_id=\"{}\"}} {} {}".format(
+            str(rank_id), step_seconds, int(round(time_stamp * 1000)))
         if self.ma:
             flags = os.O_WRONLY | os.O_CREAT
             modes = stat.S_IWUSR | stat.S_IRUSR
@@ -227,9 +259,13 @@ class FlopsUtilizationCollector(Callback):
                 f.write(train_log + '\n')
         train_log = "{} per step time: {:5.3f} ms".format(
             cb_params.mode.title(), step_seconds)
-        if self.verbose:
-            mfu = 1000 * self.full_mfu / step_seconds
-            hfu = 1000 * self.full_hfu / step_seconds
+        if self.verbose and cb_params.cur_step_num % self.data_size:
+            if self.full_flops:
+                mfu = 1000 * self.full_mfu / step_seconds
+                hfu = 1000 * self.full_hfu / step_seconds
+            else:
+                mfu = 1000 * self.shard_mfu / step_seconds
+                hfu = 1000 * self.shard_hfu / step_seconds
 
             def floored_percentage(index, val, digits):
                 val *= 10 ** (digits + 2)

mindspore/train/callback/_history.py

@@ -31,7 +31,7 @@ class History(Callback):
     outputs will be recorded.
 
     Note:
-        Normally used in `mindspore.train.Model.train` or `mindspore.train.Model.fit`.
+        Normally used in :func:`mindspore.train.Model.train` or :func:`mindspore.train.Model.fit`.
 
     Examples:
         >>> import numpy as np

mindspore/train/callback/_summary_collector.py

@@ -16,6 +16,7 @@
 from __future__ import absolute_import
 
 import os
+import platform
 import stat
 import re
 import json
@@ -43,6 +44,9 @@ from mindspore.train._utils import check_value_type, _make_directory
 from mindspore._c_expression import security
 from mindspore._c_expression import collect_host_info, get_clock_syscnt
 
+if platform.system() == "Linux":
+    import fcntl
+
 HYPER_CONFIG_ENV_NAME = "MINDINSIGHT_HYPER_CONFIG"
 HYPER_CONFIG_LEN_LIMIT = 100000
 
@@ -606,13 +610,32 @@ class SummaryCollector(Callback):
             "landscape_size": landscape_size,
             "create_landscape": create_landscape
         }
+
         meta_path = os.path.join(self._ckpt_dir, 'train_metadata.json')
+        if platform.system() != "Linux":
+            try:
+                with open(meta_path, 'w') as file:
+                    json.dump(data, file)
+                os.chmod(meta_path, stat.S_IRUSR)
+            except OSError as e:
+                logger.error("Write meta data %s failed, detail: %s" % (meta_path, str(e)))
+            return
+
+        lock_file = f"{meta_path}.lock"
         try:
-            with open(meta_path, 'w') as file:
-                json.dump(data, file)
-            os.chmod(meta_path, stat.S_IRUSR)
+            with os.fdopen(os.open(lock_file, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, stat.S_IWUSR), 'w') as f:
+                fcntl.flock(f.fileno(), fcntl.LOCK_EX)
+                if not os.path.isfile(meta_path):
+                    with open(meta_path, 'w') as file:
+                        json.dump(data, file)
+                    os.chmod(meta_path, stat.S_IRUSR)
         except OSError as e:
             logger.error("Write meta data %s failed, detail: %s" % (meta_path, str(e)))
+        try:
+            if os.path.isfile(lock_file):
+                os.remove(lock_file)
+        except OSError:
+            logger.warning("The lock file %s has been removed.", lock_file)
 
     def _save_model_params(self, cur_num, unit, backbone):
         """Save model params."""
@@ -629,12 +652,27 @@ class SummaryCollector(Callback):
 
         ckpt_file_name = f"{type(backbone).__name__}_{cur_num}_{unit}.ckpt"
         file_path = os.path.join(self._ckpt_dir, ckpt_file_name)
+        self._model_params_file_map[str(cur_num)] = file_path
+        if platform.system() != "Linux":
+            try:
+                save_checkpoint(param_list, file_path)
+            except OSError as e:
+                logger.error(str(e))
+            return
+
+        lock_file = f"{file_path}.lock"
         try:
-            save_checkpoint(param_list, file_path)
+            with os.fdopen(os.open(lock_file, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, stat.S_IWUSR), 'w') as f:
+                fcntl.flock(f.fileno(), fcntl.LOCK_EX)
+                if not os.path.isfile(file_path):
+                    save_checkpoint(param_list, file_path)
         except OSError as e:
             logger.error(str(e))
-
-        self._model_params_file_map[str(cur_num)] = file_path
+        try:
+            if os.path.isfile(lock_file):
+                os.remove(lock_file)
+        except OSError:
+            logger.warning("The lock file %s has been removed.", lock_file)
 
     def _save_model_params_for_landscape(self, cb_params):
         """Save model params for landscape."""

mindspore/train/callback/_tft_register.py

@@ -25,9 +25,9 @@ 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
+from mindspore._c_expression import _repair_device, _stop_device, _tft_sem_post, _tft_sem_enable
 from mindspore._c_expression import clean_tdt_channel
-from mindspore._c_expression import 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
 from mindspore._c_expression import Tensor as Tensor_
@@ -90,6 +90,7 @@ def _tft_exit_cb(ctx):
     _tft_sem_post()
     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"]))
@@ -105,11 +106,12 @@ or repair_info["repair_type"] == cb_ctx.tft.RepairType.RT_UCE_LOWLEVEL.value):
         cb_params = args
         src_rank = repair_info["src"][0]
         dst_rank = repair_info["dst"][0]
-        send_recv(cb_params.network.trainable_params(), src_rank, dst_rank)
+        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")
 
 
-def _tft_clean_callback(is_uce_error, ctx):
+def _tft_clean_callback(is_uce_error, args, ctx):
     """ Callback used for TFT clean function."""
     logger.info("Enter _tft_clean_callback")
     ret = 0
@@ -130,12 +132,13 @@ def _tft_clean_callback(is_uce_error, ctx):
     return ret
 
 
-def _tft_stop_callback(cb_ctx):
+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))
     _stop_device(cb_ctx.device_id)
-    if 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
     logger.info("Finish _tft_stop_callback")
 
 
@@ -160,13 +163,23 @@ class TFTRegister(Callback):
         ModuleNotFoundError: Mindio TFT whl package is not installed.
 
     Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            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>`_
+            for more details.
+
+            This example should be run with 4 devices.
+
         >>> import numpy as np
         >>> import os
         >>> import math
         >>> import mindspore as ms
         >>> import mindspore.dataset as ds
         >>> from mindspore import nn, ops, Parameter, train
-        >>> from mindspore.communication import init
+        >>> from mindspore.communication import init, get_rank
         >>> from mindspore.common.initializer import initializer, HeUniform
         >>> from mindspore.train import Model, TFTRegister
         >>> from mindspore import dataset as ds
@@ -175,7 +188,7 @@ class TFTRegister(Callback):
         >>> init()
         >>> ms.set_seed(1)
         >>> ms.set_auto_parallel_context(strategy_ckpt_config={"save_file":
-        >>>                             "./src_pipeline_strategys/src_strategy_{}.ckpt".format(get_rank())})
+        ...                             "./src_pipeline_strategys/src_strategy_{}.ckpt".format(get_rank())})
         >>> class MatMulCell(nn.Cell):
         ...     def __init__(self, param=None, shape=None):
         ...         super().__init__()
@@ -233,7 +246,7 @@ class TFTRegister(Callback):
         ...     dataset = dataset.batch(batch_size)
         ...     return dataset
         >>>
-        >>> data_set = create_dataset(32)
+        >>> dataset = create_dataset(32)
         >>>
         >>> optimizer = nn.SGD(net.trainable_params(), 1e-2)
         >>> optimizer_wrapper = nn.OptTFTWrapper(optimizer)
@@ -241,8 +254,8 @@ class TFTRegister(Callback):
         >>>
         >>> net_with_loss = nn.PipelineCell(nn.WithLossCell(net, loss_fn), 4)
         >>> net_with_loss.set_train()
-        >>> model = Model(net_with_loss, optimizer=optimizer)
-        >>> tft_cb = TFTRegister("192.168.0.1", 2000, "./tft_checkpoint/")
+        >>> model = Model(net_with_loss, optimizer=optimizer_wrapper)
+        >>> tft_cb = TFTRegister(0, "192.168.0.1", 2000, "./tft_checkpoint/")
         >>> loss_cb = train.LossMonitor(1)
         >>> model.train(1, dataset, callbacks=[tft_cb, loss_cb])
     """
@@ -264,6 +277,7 @@ class TFTRegister(Callback):
         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
@@ -274,6 +288,7 @@ class TFTRegister(Callback):
         self.assign = mindspore.ops.Assign()
         self.g_one = Parameter(Tensor([1], dtype=mstype.int32))
         self.s1 = mindspore.hal.Stream()
+        _tft_sem_enable()
 
     def _is_params_consistent(self):
         for key, param in self.cb_params.train_network.parameters_and_names():
@@ -300,7 +315,7 @@ class TFTRegister(Callback):
         replica_info = [
             {
                 "type": 1,
-                "rank_list": dp,
+                "rank_list": list(dp),
                 "replica_cnt": len(dp),
                 "replica_shift": 0
             }
@@ -321,13 +336,12 @@ class TFTRegister(Callback):
         cur_rank = get_rank()
         enable_local_copy = False
         enable_arf = False
-        enable_zit = 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, enable_zit)
+            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.")
 
@@ -336,6 +350,14 @@ class TFTRegister(Callback):
         self.tft.tft_start_processor(self._controller_ip, self._controller_port)
         logger.info("Finished start tft processor.")
 
+    def _reset_acc_grads(self):
+        accu_grad_params = map(lambda e: e[1],
+                               filter(lambda e: e[1].name.startswith('accu_grads'),
+                                      self.cb_params.train_network.parameters_and_names()))
+        accu_grad_list = list(accu_grad_params)
+        if reset_params(accu_grad_list) != 0:
+            raise ValueError("Call reset_params failed.")
+
     def on_train_step_end(self, run_context):
         """
         And report status to MindIO TFT after every step finished.
@@ -349,13 +371,13 @@ class TFTRegister(Callback):
             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.tft.tft_end_updating_os(cb_params.cur_step_num)
         if cb_params.optimizer is not None:
             self.global_step = int(cb_params.optimizer.global_step.data)
             self.assign(cb_params.optimizer.tft_g_one_flag, self.g_one)
         else:
             self.global_step = int(cb_params.network.optimizer.global_step.data)
             self.assign(cb_params.network.optimizer.tft_g_one_flag, self.g_one)
+        self.tft.tft_end_updating_os(cb_params.cur_step_num)
         logger.info("END Set optimizer finish step status to TFT.")