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

mindspore/train/callback/_train_fault_tolerance.py

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

mindspore/train/data_sink.py

@@ -16,9 +16,9 @@
 from functools import wraps
 import mindspore.ops as ops
 from mindspore import context
-from mindspore.common.dtype import pytype_to_dtype
+from mindspore.common.dtype import _pytype_to_dtype
 from mindspore.common.api import jit
-from mindspore.train._utils import _exec_datagraph, _get_types_and_shapes
+from mindspore.train._utils import _exec_datagraph, _get_types_and_shapes, enable_data_broadcast
 from mindspore.train.dataset_helper import _has_dynamic_shape, _check_inputs
 import mindspore.dataset as ds
 from mindspore._c_expression import _set_dataset_mode_config
@@ -41,6 +41,15 @@ def _init_sink_dataset(dataset, sink_size, input_signature, create_info):
     is_info_queue = (create_info and sink_size == 1 and dataset_size != 1 and
                      input_signature is None and not dynamic_shape and
                      context.get_context('device_target') == 'Ascend')
+
+    # Don't enable dynamic shape(multi-subgraph) feature in pp/data_broadcast mode,
+    # otherwise get_data_info will stuck since some rank do not consume data.
+    use_pipeline_parallel = (context.get_auto_parallel_context("pipeline_stages") > 1)
+    data_broadcast = enable_data_broadcast()
+
+    if use_pipeline_parallel or data_broadcast:
+        is_info_queue = False
+
     transfer_dataset = _exec_datagraph(dataset, sink_size, create_data_info_queue=is_info_queue)
     dataset.__transfer_dataset__ = transfer_dataset
 
@@ -52,7 +61,7 @@ def _init_sink_dataset(dataset, sink_size, input_signature, create_info):
         _check_inputs(input_signature, dataset_shapes, dataset_types)
 
     queue_name = transfer_dataset.queue_name
-    if _need_to_full() and context.get_context('mode') == context.GRAPH_MODE:
+    if _need_to_full():
         device_num = _get_device_num() // _get_pipeline_stages()
         dataset_shapes = _to_full_shapes(dataset_shapes, device_num)
     next_op = ops.GetNext(dataset_types, dataset_shapes, len(dataset_types), queue_name)
@@ -85,12 +94,12 @@ def _get_next_op(dataset, ori_next_op, is_info_queue):
 
     queue_name = dataset.__transfer_dataset__.queue_name
     dataset_types, dataset_shapes = dataset.__transfer_dataset__.get_data_info()
-    dataset_types = [pytype_to_dtype(x) for x in dataset_types]
+    dataset_types = [_pytype_to_dtype(x) for x in dataset_types]  # pylint:disable=protected-access
     key = str(dataset_types) + str(dataset_shapes)
     if key in dataset.__sink_aux__.next_ops:
         next_op = dataset.__sink_aux__.next_ops[key]
     else:
-        if _need_to_full() and context.get_context('mode') == context.GRAPH_MODE:
+        if _need_to_full():
             device_num = _get_device_num() // _get_pipeline_stages()
             dataset_shapes = _to_full_shapes(dataset_shapes, device_num)
         next_op = ops.GetNext(dataset_types, dataset_shapes, len(dataset_types), queue_name)
@@ -214,7 +223,7 @@ def data_sink(fn, dataset, sink_size=1, jit_config=None, input_signature=None):
     loop = sink_size
     create_info = True
     if jit_config is None:
-        create_info = (loop == 1)
+        create_info = loop == 1
         loop = 1
     ori_next_op, is_info_queue = _init_sink_dataset(dataset, loop, input_signature, create_info)
 

mindspore/train/dataset_helper.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Huawei Technologies Co., Ltd
+# Copyright 2020-2025 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -20,8 +20,8 @@ import copy
 
 from mindspore import _checkparam as Validator
 from mindspore import log as logger
-from mindspore.common._auto_dynamic import is_auto_dynamic, convert_new_shapes
-from mindspore.common.dtype import pytype_to_dtype
+from mindspore.common.dynamic_shape._auto_dynamic import is_auto_dynamic, convert_new_shapes
+from mindspore.common.dtype import _pytype_to_dtype
 from mindspore.common.api import _cell_graph_executor, _is_args_fullmode, ARG_SPECIFIED
 from mindspore.common._utils import is_shape_unknown
 from mindspore.dataset.core import config as dataset_config
@@ -34,7 +34,7 @@ from mindspore.parallel._utils import _get_device_num, _get_global_rank, _need_t
     _origin_shapes, _dynamic_shape_for_dataset
 from mindspore.parallel._ps_context import _is_role_sched
 from mindspore.ops import operations as P
-from mindspore.common.auto_dynamic_shape import _auto_dynamic_shape
+from mindspore.common.dynamic_shape.auto_dynamic_shape import _auto_dynamic_shape
 
 
 def _send_data(dataset, epoch_num):
@@ -275,7 +275,7 @@ def connect_network_with_dataset(network, dataset_helper):
         # Need to do full_batch for shapes which also do in the _DatasetIterMSLoopSink
         if _need_to_full():
             dataset_shapes = _to_full_shapes(dataset_shapes, _get_device_num() // _get_pipeline_stages())
-        dataset_types = [pytype_to_dtype(x) for x in dataset_types]
+        dataset_types = [_pytype_to_dtype(x) for x in dataset_types]  # pylint:disable=protected-access
         if not is_dynamic:
             dataset_shapes = _auto_dynamic_shape.auto_dynamic_generate_compile_args(dataset_shapes, True)
         key = str(dataset_types) + str(dataset_shapes)
@@ -564,6 +564,15 @@ class _DatasetIter:
                 self.sink_size = dataset.__loop_size__
             create_data_info_queue = (
                 sink_size == 1 and self.sink_count == 1 and dataset.get_dataset_size() != 1)
+
+            # Don't enable dynamic shape(multi-subgraph) feature in pp/data_broadcast mode,
+            # otherwise get_data_info will stuck since some rank do not consume data.
+            use_pipeline_parallel = (context.get_auto_parallel_context("pipeline_stages") > 1)
+            data_broadcast = enable_data_broadcast()
+
+            if use_pipeline_parallel or data_broadcast:
+                create_data_info_queue = False
+
             dataset.__transfer_dataset__ = _exec_datagraph(dataset, self.sink_size,
                                                            create_data_info_queue=create_data_info_queue)
 

mindspore/train/model.py

@@ -57,7 +57,7 @@ from mindspore.dataset.engine.datasets import _set_training_dataset, _reset_trai
 from mindspore.train import amp
 from mindspore._c_expression import _framework_profiler_step_start, _framework_profiler_step_end
 from mindspore._c_expression import _get_optimzer_timestamps
-from mindspore._c_expression import clean_tdt_channel
+from mindspore._c_expression import clean_tdt_channel, _clean_rootinfo
 
 from mindspore.parallel._utils import _init_auto_parallel_context, _clear_auto_parallel_context
 from .serialization import load_param_into_net
@@ -154,6 +154,9 @@ def _handle_exception_info(obj, uce_env, tft, e):
                 tft.tft_report_error(tft.ReportState.RS_UNKNOWN.value)
                 raise e
         tft.tft_report_error(tft.ReportState.RS_UCE.value)
+    elif "HCCEError" in e_str:
+        logger.warning("uce wrapper caught HCCEError")
+        tft.tft_report_error(tft.ReportState.RS_HCCL_FAILED.value)
     elif "ForceStopError" in e_str:
         logger.warning("uce wrapper caught RuntimeError ForceStopError")
         force_stop_err = tft.ReportState.RS_NORMAL.value
@@ -246,9 +249,8 @@ def _handle_tft(func):
                 if isinstance(item, TrainFaultTolerance):
                     obj = item
         if obj:
-            tft = obj.tft
             tft_env = os.getenv("MS_ENABLE_TFT", "")
-            uce_env = "UCE:1" in tft_env or "ARF:1" in tft_env
+            uce_env = "UCE:1" in tft_env or "ARF:1" in tft_env or "HCCE:1" in tft_env
             tre_env = "TRE:1" in tft_env
             while True:
                 try:
@@ -260,11 +262,11 @@ def _handle_tft(func):
                         _update_ckpt_callback_info(repair_step, **kwargs)
                         logger.warning(f'Resume training after TREError from step {repair_step}.')
                     else:
-                        _handle_exception_info(obj, uce_env, tft, e)
-                        ret = tft.tft_wait_next_action()
-                        if ret == tft.Action.EXIT.value:
+                        _handle_exception_info(obj, uce_env, obj.tft, e)
+                        ret = obj.tft.tft_wait_next_action()
+                        if ret == obj.tft.Action.EXIT.value:
                             raise e
-                        repair_step = tft.tft_get_repair_step()
+                        repair_step = obj.tft.tft_get_repair_step()
                         logger.warning(
                             "uce wrapper caught repair finish REPAIR STEP: {} batch_num:{}".format(repair_step,
                                                                                                    self.batch_num))
@@ -274,6 +276,7 @@ def _handle_tft(func):
                     cb_initial_step = _calc_cb_initial_step(initial_epoch, initial_step, *args, **kwargs)
                     if not self.enable_tre:
                         kwargs["initial_step"] = cb_initial_step
+                        self._initial_step = 0
                     # reset all accu grads to zero
                     obj._reset_acc_grads()
                     logger.warning(
@@ -281,9 +284,9 @@ def _handle_tft(func):
                                                                                                      cb_initial_step))
                     continue
                 except BaseException as e:
-                    if tft:
+                    if obj.tft:
                         logger.error("uce wrapper caught BaseException error, enter MindIO TTP process.", exc_info=True)
-                        tft.tft_report_error(tft.ReportState.RS_UNKNOWN.value)
+                        obj.tft.tft_report_error(obj.tft.ReportState.RS_UNKNOWN.value)
                     raise e
         else:
             return func(self, *args, **kwargs)
@@ -300,9 +303,6 @@ def _check_tft():
         ascend_target = MSContext.get_instance().get_ascend_soc_version()
         if ascend_target == 'ascend910':
             raise ValueError("TFT is not supported when using ascend910")
-        ms_mode = context.get_context("mode")
-        if ms_mode != mindspore.GRAPH_MODE:
-            raise ValueError("TFT is only supported in GRAPH_MODE")
         jit_level = context.get_context("jit_level")
         if jit_level == "O2" and ("UCE:1" in tft_env or "ARF:1" in tft_env):
             raise ValueError("TFT is not supported when using jit_level == O2")
@@ -443,6 +443,11 @@ def _set_with_processed_inputs(network, inputs):
             "Reset inputs from a process inputs, should be a list/tuple or a dict, but got %s!" % str(inputs))
 
 
+def _check_tft_reset_dataset():
+    env_tft = os.getenv("MS_ENABLE_TFT", "")
+    return any([v in env_tft for v in ["TRE:1", "UCE:1", "HCCE:1", "ARF:1"]])
+
+
 class Model:
     """
     High-Level API for training or inference.
@@ -561,7 +566,9 @@ class Model:
         self._mindspore_lite_model_group_id = id(self) & 0xFFFF
         self.batch_num = -1
         self.enable_tre = "TRE:1" in os.getenv("MS_ENABLE_TFT", "")
+        self.enable_hcce = "HCCE:1" in os.getenv("MS_ENABLE_TFT", "")
         self._initial_step = None
+        self._need_reset_data = _check_tft_reset_dataset()
         _clear_auto_parallel_context(self._network)
 
     def _check_for_graph_cell(self, kwargs):
@@ -761,7 +768,7 @@ class Model:
             logger.info("Begin to connect network with dataset.")
             network = connect_network_with_dataset(network, dataset_helper)
 
-        if (_get_recovery_context("enable_recovery") or self.enable_tre) and is_train:
+        if (_get_recovery_context("enable_recovery") or self._need_reset_data) and is_train:
             _set_training_dataset(dataset_helper)
 
         network.set_train(is_train)
@@ -805,7 +812,7 @@ class Model:
         """
         if os.environ.get("MS_ENABLE_CKPT_D2H_ASYNC") != "1":
             return
-        if (context.get_context("mode") == context.GRAPH_MODE) and (context.get_context("device_target") == "Ascend"):
+        if context.get_context("device_target") == "Ascend":
             cb_params.need_ckpt, cb_params.save_checkpoint_steps, \
             cb_params.last_triggered_step = self._check_need_ckpt(cb_params.list_callback)
             logger.info(f"need_ckpt:{cb_params.need_ckpt},"
@@ -873,8 +880,8 @@ class Model:
             sink_size (int): Control the amount of data in each sink. Default: -1.
             epoch (int): Total number of iterations on the data. Default: 1.
         """
-        if context.get_context("mode") != context.GRAPH_MODE or context.get_context("device_target") != "Ascend":
-            raise RuntimeError('Pre-init process only supports GRAPH MODE and Ascend target currently.')
+        if context.get_context("device_target") != "Ascend":
+            raise RuntimeError('Pre-init process only supports Ascend target currently.')
 
         if not train_dataset and not valid_dataset:
             raise ValueError("The argument 'train_dataset' and 'valid_dataset' can not both be None or empty.")
@@ -1057,7 +1064,7 @@ class Model:
             initial_epoch (int): Epoch at which to start train, it used for resuming a previous training run.
                                  Default: 0.
         """
-        is_graph = (context.get_context("mode") == context.GRAPH_MODE)
+        is_graph = context.get_context("mode") == context.GRAPH_MODE
         dataset_size = train_dataset.get_dataset_size()
         if dataset_size % sink_size != 0:
             logger.info("In dataset_sink mode (dataset_size % sink_size) should equal to 0, "
@@ -1126,6 +1133,7 @@ class Model:
                 if cb_params.is_arf:
                     cb_params.is_arf = False
                     _set_recovery_context(is_arf=False)
+                _clean_rootinfo()
 
                 # Embedding cache server only run one step.
                 if is_embedding_cache_server:
@@ -1204,8 +1212,6 @@ class Model:
         if not enable_recovery:
             self.enable_recovery = False
         else:
-            if context.get_context("mode") != context.GRAPH_MODE:
-                raise RuntimeError("Recovery for training only support graph mode currently.")
             self.enable_recovery = enable_recovery and _is_role_worker()
 
     def _check_need_load_ckpt(self, cb_params, dataset_size, sink_size=-1):
@@ -1340,6 +1346,7 @@ class Model:
                 if cb_params.is_arf:
                     cb_params.is_arf = False
                     _set_recovery_context(is_arf=False)
+                _clean_rootinfo()
                 # Embedding cache server only run one step.
                 if is_embedding_cache_server:
                     break
@@ -2182,9 +2189,6 @@ class Model:
             dataset_sink_mode (bool): Determines whether to pass the data through dataset channel.
             sink_size (int): Control the amount of data in each sink.
         """
-        if context.get_context("mode") != context.GRAPH_MODE:
-            raise RuntimeError("Pre-compile process that generate parameter layout for the train network "
-                               "only supports GRAPH MODE and Ascend target currently.")
         if _get_parallel_mode() not in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL):
             raise RuntimeError("'infer_train_layout' only supports 'semi_auto_parallel' and 'auto_parallel' "
                                "mode, but got {}.".format(_get_parallel_mode()))
@@ -2303,6 +2307,7 @@ class Model:
 
         Examples:
             >>> import numpy as np
+            >>> import mindspore as ms
             >>> import mindspore.nn as nn
             >>> from mindspore import Tensor
             >>> from mindspore.train import Model
@@ -2312,28 +2317,28 @@ class Model:
             >>> from mindspore.parallel.auto_parallel import AutoParallel
             >>>
             >>> class Net(nn.Cell):
-            >>>     def __init__(self):
-            >>>         super(Net, self).__init__()
-            >>>         self.fc1 = nn.Dense(128, 768, activation='relu')
-            >>>         self.fc2 = nn.Dense(128, 768, activation='relu')
-            >>>         self.fc3 = nn.Dense(128, 768, activation='relu')
-            >>>         self.fc4 = nn.Dense(768, 768, activation='relu')
-            >>>         self.relu4 = nn.ReLU()
-            >>>         self.relu5 = nn.ReLU()
-            >>>         self.transpose = P.Transpose()
-            >>>         self.matmul1 = P.MatMul()
-            >>>         self.matmul2 = P.MatMul()
-            >>>
-            >>>     def construct(self, x):
-            >>>         q = self.fc1(x)
-            >>>         k = self.fc2(x)
-            >>>         v = self.fc3(x)
-            >>>         k = self.transpose(k, (1, 0))
-            >>>         c = self.relu4(self.matmul1(q, k))
-            >>>         s = self.relu5(self.matmul2(c, v))
-            >>>         s = self.fc4(s)
-            >>>         return s
-            >>>
+            ...     def __init__(self):
+            ...         super(Net, self).__init__()
+            ...         self.fc1 = nn.Dense(128, 768, activation='relu')
+            ...         self.fc2 = nn.Dense(128, 768, activation='relu')
+            ...         self.fc3 = nn.Dense(128, 768, activation='relu')
+            ...         self.fc4 = nn.Dense(768, 768, activation='relu')
+            ...         self.relu4 = nn.ReLU()
+            ...         self.relu5 = nn.ReLU()
+            ...         self.transpose = P.Transpose()
+            ...         self.matmul1 = P.MatMul()
+            ...         self.matmul2 = P.MatMul()
+            ...
+            ...     def construct(self, x):
+            ...         q = self.fc1(x)
+            ...         k = self.fc2(x)
+            ...         v = self.fc3(x)
+            ...         k = self.transpose(k, (1, 0))
+            ...         c = self.relu4(self.matmul1(q, k))
+            ...         s = self.relu5(self.matmul2(c, v))
+            ...         s = self.fc4(s)
+            ...         return s
+            ...
             >>> ms.set_context(mode=ms.GRAPH_MODE)
             >>> init()
             >>> inputs = Tensor(np.ones([32, 128]).astype(np.float32))
@@ -2343,9 +2348,6 @@ class Model:
             >>> predict_map = model.infer_predict_layout(inputs)
         """
         _init_auto_parallel_context(self._network)
-        if context.get_context("mode") != context.GRAPH_MODE:
-            raise RuntimeError("Pre-compile process that generate parameter layout for the predict network "
-                               "only supports GRAPH MODE and Ascend target currently.")
         if _get_parallel_mode() not in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL):
             raise RuntimeError('Infer predict layout only supports semi auto parallel and auto parallel mode.')
         _parallel_predict_check()