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

mindspore/train/model.py

@@ -28,7 +28,7 @@ import numpy as np
 
 import mindspore
 from mindspore import log as logger
-from mindspore.train.serialization import save_checkpoint, load_checkpoint
+from mindspore.train.serialization import save_checkpoint
 from mindspore.train.callback._checkpoint import ModelCheckpoint, _chg_ckpt_file_name_if_same_exist
 from mindspore.common.tensor import Tensor
 from mindspore.train.metrics import get_metrics, get_metric_fn
@@ -40,16 +40,12 @@ from mindspore.train.callback import __all__ as internal_cb_names
 from mindspore.train.callback._cluster_monitor import ClusterMonitor
 from mindspore import context
 from mindspore.parallel._utils import _get_parallel_mode, _get_device_num, _get_parameter_broadcast, \
-    _device_number_check, _parameter_broadcast_check, _parallel_predict_check, \
-    _reset_op_id_with_offset
-from mindspore.parallel._ps_context import _is_role_worker, _is_role_pserver, _is_ps_mode, \
-    _cache_enable, _enable_distributed_mindrt
+    _device_number_check, _parameter_broadcast_check, _parallel_predict_check
 from mindspore.train.metrics import Loss
 from mindspore.log import vlog_print
 from mindspore import nn
 from mindspore.boost import AutoBoost
 from mindspore.context import ParallelMode
-from mindspore.parallel._recovery_context import _set_recovery_context, _get_recovery_context
 from mindspore.train.dataset_helper import DatasetHelper, connect_network_with_dataset
 from mindspore.common.api import _pynative_executor, ARG_SPECIFIED, TOTAL_ARG_LEN
 from mindspore.dataset.core.config import get_debug_mode
@@ -57,7 +53,8 @@ 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, _clean_rootinfo
+from mindspore._c_expression import clean_tdt_channel, _clean_rootinfo, check_is_arf, set_is_arf
+from mindspore._c_expression import _get_snapshot_params, _is_snapshot_valid
 
 from mindspore.parallel._utils import _init_auto_parallel_context, _clear_auto_parallel_context
 from .serialization import load_param_into_net
@@ -156,18 +153,14 @@ def _handle_exception_info(obj, uce_env, tft, e):
         tft.tft_report_error(tft.ReportState.RS_UCE.value)
     elif "HCCEError" in e_str:
         logger.warning("uce wrapper caught HCCEError")
-        if  obj.stop_been_called:
-            logger.warning("Received HCCEError after force stop been called, so report force stopped error to MindIO.")
-            tft.tft_report_error(tft.ReportState.RS_NORMAL.value)
-        else:
-            tft.tft_report_error(tft.ReportState.RS_HCCL_FAILED.value)
+        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
         tft.tft_report_error(force_stop_err)
     elif "ARF FINISH" in e_str:
         logger.warning(f"ARF FINISH")
-        _set_recovery_context(is_arf=True)
+        set_is_arf(True)
         tft.tft_report_error(tft.ReportState.RS_PREREPAIR_FINISH.value)
     else:
         logger.error("uce wrapper caught other RuntimeError, enter MindIO TTP process.", exc_info=True)
@@ -179,7 +172,12 @@ def _handle_training_result_error(model, tft_obj):
     """
     Handle training result error for resuming training.
     """
-    ckpt_load_fn = tft_obj.ckpt_load_func
+    def load_snapshot_params():
+        param_dict = {}
+        for name, tensor in _get_snapshot_params().items():
+            param_dict[name] = mindspore.Parameter(tensor, name=name)
+        return (param_dict, False)
+    ckpt_load_fn = load_snapshot_params if _is_snapshot_valid() else tft_obj.ckpt_load_func
     train_network = tft_obj.cb_params.train_network
     logger.warning("Process training result error start.")
     # 1. Clear tdt channel
@@ -238,6 +236,20 @@ def _update_ckpt_callback_info(resume_train_step, **kwargs):
         ckpt_obj._append_step_num = resume_train_step
 
 
+def _get_tft_obj(**kwargs):
+    """
+    Get TrainFaultTolerance from kwargs of callback
+    """
+    obj = None
+    if kwargs.get('callbacks') and isinstance(kwargs.get('callbacks'), TrainFaultTolerance):
+        obj = kwargs.get('callbacks')
+    if kwargs.get('callbacks') and isinstance(kwargs.get('callbacks'), list):
+        for item in kwargs.get('callbacks'):
+            if isinstance(item, TrainFaultTolerance):
+                obj = item
+    return obj
+
+
 def _handle_tft(func):
     """
     Decorator function, which starts uce handle process when an exception occurs during training.
@@ -245,17 +257,11 @@ def _handle_tft(func):
 
     @wraps(func)
     def wrapper(self, *args, **kwargs):
-        obj = None
-        if kwargs.get('callbacks') and isinstance(kwargs.get('callbacks'), TrainFaultTolerance):
-            obj = kwargs.get('callbacks')
-        if kwargs.get('callbacks') and isinstance(kwargs.get('callbacks'), list):
-            for item in kwargs.get('callbacks'):
-                if isinstance(item, TrainFaultTolerance):
-                    obj = item
-        if obj:
+        obj = _get_tft_obj(**kwargs)
+        if obj and not TrainFaultTolerance._only_enable_ckpt_d2h_async():
             tft_env = os.getenv("MS_ENABLE_TFT", "")
             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
+            tre_env = "TRE:1" in tft_env or "TRE:2" in tft_env
             while True:
                 try:
                     return func(self, *args, **kwargs)
@@ -270,7 +276,6 @@ def _handle_tft(func):
                         ret = obj.tft.tft_wait_next_action()
                         if ret == obj.tft.Action.EXIT.value:
                             raise e
-                        obj.stop_been_called = False
                         repair_step = obj.tft.tft_get_repair_step()
                         logger.warning(
                             "uce wrapper caught repair finish REPAIR STEP: {} batch_num:{}".format(repair_step,
@@ -308,9 +313,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")
@@ -564,9 +566,7 @@ class Model:
         self._current_epoch_num = 0
         self._current_step_num = 0
         self.epoch_iter = 0
-        self.enable_recovery = False
         self._backbone_is_train = True
-        self.need_load_ckpt = False
         self._lite_full_predictor = None
         self._lite_incremental_predictor = None
         self._mindspore_lite = None
@@ -739,10 +739,7 @@ class Model:
         metrics = dict()
         # There's no need for server to execute eval, just give fake metrics.
         for key, value in self._metric_fns.items():
-            if not _is_role_pserver():
-                metrics[key] = value.eval()
-            else:
-                metrics[key] = 1
+            metrics[key] = value.eval()
         return metrics
 
     def _get_scaling_sens(self):
@@ -776,7 +773,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._need_reset_data) and is_train:
+        if self._need_reset_data and is_train:
             _set_training_dataset(dataset_helper)
 
         network.set_train(is_train)
@@ -818,9 +815,7 @@ class Model:
         :param cb_params: callback params
         :return: none
         """
-        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 TrainFaultTolerance._enable_snapshot() and 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},"
@@ -888,8 +883,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.")
@@ -1026,13 +1021,10 @@ class Model:
         callbacks = cb_params.list_callback
         cb_params.train_dataset_element = None
         cb_params.network = self._network
-        # Embedding cache server only run one step.
-        if _is_role_pserver() and _cache_enable():
-            epoch = 1
         cb_params.last_save_ckpt_step = None
         cb_params.latest_ckpt_file = None
         cb_params.loss_scale_mananger = self._loss_scale_manager
-        cb_params.is_arf = _get_recovery_context("is_arf")
+        cb_params.is_arf = check_is_arf()
         cb_params.initial_step = self._initial_step
 
         # build callback list
@@ -1094,12 +1086,6 @@ class Model:
             dataset_helper = train_dataset._dataset_helper
 
         self.epoch_iter = 0
-        self._check_enable_recovery()
-        # Used to check whether need perform recovery for process which is restarted.
-        self._check_need_load_ckpt(cb_params, dataset_size, sink_size)
-        # Check whether this process is embedding cache server.
-        is_embedding_cache_server = _is_role_pserver() and _cache_enable()
-
         while self.epoch_iter < (epoch - initial_epoch):
             cb_params.cur_epoch_num = self.epoch_iter + 1 + initial_epoch
             self._current_epoch_num = cb_params.cur_epoch_num
@@ -1115,11 +1101,6 @@ class Model:
             cb_params.train_network = train_network
             cb_params.dataset_helper = dataset_helper
 
-            # Perform recovery for process which is restarted.
-            self._reset_training_step_for_abnormal_process(cb_params, dataset_helper)
-            # Perform recovery for process which is not restarted.
-            self._reset_training_step_for_normal_process(cb_params, dataset_helper)
-
             # For data sink dataset_helper only iter once, other wise iter epoch_size times.
             for inputs in dataset_helper:
                 if is_graph:
@@ -1134,36 +1115,17 @@ class Model:
                 outputs = train_network(*inputs)
                 cb_params.net_outputs = outputs
 
-                # In disaster recovery scenarios, need not to execute callbacks if this step executes failed.
-                need_exec_callback_step_end = not (self.enable_recovery and _get_recovery_context("need_reset"))
-                if need_exec_callback_step_end:
-                    list_callback.on_train_step_end(run_context)
+                list_callback.on_train_step_end(run_context)
+
                 if cb_params.is_arf:
                     cb_params.is_arf = False
-                    _set_recovery_context(is_arf=False)
+                    set_is_arf(False)
                 _clean_rootinfo()
 
-                # Embedding cache server only run one step.
-                if is_embedding_cache_server:
-                    break
-
             dataset_helper.continue_send()
 
-            # When it's distributed training and using MindRT,
-            # the node id should be reset to start from 0.
-            # This is to avoid the timeout when finding the actor route tables in 'train' and 'eval' case(or 'fit').
-            if _enable_distributed_mindrt():
-                _reset_op_id_with_offset()
-
             self._eval_during_train(valid_infos, cb_params, list_callback)
-
-            # In disaster recovery scenarios, need not to execute callbacks if this epoch executes failed.
-            # Embedding cache server need not do epoch end callback, this process only run one step.
-            need_exec_callback_epoch_end = not ((self.enable_recovery and _get_recovery_context("need_reset"))
-                                                or is_embedding_cache_server)
-
-            if need_exec_callback_epoch_end:
-                list_callback.on_train_epoch_end(run_context)
+            list_callback.on_train_epoch_end(run_context)
             if "metrics" in cb_params or "eval_results" in cb_params:
                 cb_params.pop("metrics", None)
                 cb_params.pop("eval_results", None)
@@ -1172,12 +1134,7 @@ class Model:
             if should_stop:
                 break
 
-            need_reset_to_beginning = self.enable_recovery and _get_recovery_context("need_reset") \
-                                      and not _get_recovery_context("latest_ckpt_file")
             self.epoch_iter += 1
-            if need_reset_to_beginning:
-                self.epoch_iter = 0
-                cb_params.cur_step_num = 0
 
         dataset_helper.stop_send()
         dataset_helper.release()
@@ -1211,95 +1168,6 @@ class Model:
             cb_params.dataset_sink_mode = train_dataset_sink_mode
             cb_params.net_outputs = train_net_outputs
 
-    def _check_enable_recovery(self):
-        """
-        Check whether enable recovery and execution mode consistency.
-        """
-
-        enable_recovery = _get_recovery_context("enable_recovery") and context.get_context("device_target") == "GPU"
-        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):
-        """
-        Check whether need to load checkpoint after abnormal process restart.
-
-        Args:
-            cb_params (_InternalCallbackParam): Callback parameters.
-            dataset_size (int): The number of batches in a dataset.
-            sink_size (int): Control the amount of data in each sink. Default: -1.
-        """
-        if context.get_context("device_target") != "GPU":
-            return
-        if not self.enable_recovery:
-            self.need_load_ckpt = False
-
-        cb_params.latest_ckpt_file = _get_recovery_context("latest_ckpt_file")
-        if cb_params.latest_ckpt_file:
-            recovery_epoch_num = _get_recovery_context("latest_ckpt_epoch")
-            recovery_step_num = _get_recovery_context("latest_ckpt_step")
-            dataset_sink_size = sink_size if sink_size > 0 else dataset_size
-            cb_params.cur_step_num = (recovery_epoch_num - 1) * dataset_sink_size + recovery_step_num
-            cb_params.last_save_ckpt_step = cb_params.cur_step_num
-            self.epoch_iter = recovery_epoch_num
-            self.need_load_ckpt = True
-        else:
-            self.need_load_ckpt = False
-
-    def _reset_training_step_for_abnormal_process(self, cb_params, dataset_helper):
-        """
-        Execute recovery for abnormal exit process when restart.
-
-        Args:
-            cb_params (_InternalCallbackParam): Callback parameters.
-        """
-
-        if self.need_load_ckpt:
-            try:
-                load_checkpoint(cb_params.latest_ckpt_file, cb_params.train_network)
-            except BaseException as e:
-                os.remove(cb_params.latest_ckpt_file)
-                raise RuntimeError(e.__str__() + ", load ckpt failed and remove the ckpt: " \
-                                   + cb_params.latest_ckpt_file) from e
-            _reset_training_dataset(cb_params.cur_step_num, dataset_helper.iter.dataset.get_dataset_size())
-            self.need_load_ckpt = False
-
-    def _reset_training_step_for_normal_process(self, cb_params, dataset_helper):
-        """
-        Execute recovery for normal process when there is process exit abnormally.
-
-        Args:
-            cb_params (_InternalCallbackParam): Callback parameters.
-            dataset_helper (DatasetHelper): A class to process the MindData dataset,
-                it provides the type, shape and queue name of the dataset to wrap the `GetNext`.
-        """
-
-        if self.enable_recovery and _get_recovery_context("need_reset"):
-            cb_params.latest_ckpt_file = _get_recovery_context("latest_ckpt_file")
-            if cb_params.latest_ckpt_file:
-                try:
-                    load_checkpoint(cb_params.latest_ckpt_file, cb_params.train_network)
-                except BaseException as e:
-                    os.remove(cb_params.latest_ckpt_file)
-                    raise RuntimeError(e.__str__() + ", load ckpt failed and remove the ckpt: "\
-                         + cb_params.latest_ckpt_file) from e
-
-                recovery_epoch_num = _get_recovery_context("latest_ckpt_epoch")
-                recovery_step_num = _get_recovery_context("latest_ckpt_step")
-                cb_params.cur_step_num = (recovery_epoch_num - 1) * dataset_helper.sink_size() + recovery_step_num
-                self.epoch_iter = recovery_epoch_num
-                cb_params.cur_epoch_num = self.epoch_iter + 1
-                cb_params.last_save_ckpt_step = cb_params.cur_step_num
-                _reset_training_dataset(cb_params.cur_step_num, dataset_helper.iter.dataset.get_dataset_size())
-            else:
-                _reset_training_dataset(0, dataset_helper.iter.dataset.get_dataset_size())
-
-            _set_recovery_context(need_reset=False)
-
     def _train_process(self, epoch, train_dataset, list_callback=None, cb_params=None, initial_epoch=0,
                        valid_infos=None):
         """
@@ -1324,7 +1192,6 @@ class Model:
         cb_params.dataset_sink_mode = False
         run_context = RunContext(cb_params)
         list_callback.on_train_begin(run_context)
-        is_embedding_cache_server = _is_role_pserver() and _cache_enable()
 
         for i in range(initial_epoch, epoch):
             cb_params.cur_epoch_num = i + 1
@@ -1355,21 +1222,12 @@ class Model:
                 list_callback.on_train_step_end(run_context)
                 if cb_params.is_arf:
                     cb_params.is_arf = False
-                    _set_recovery_context(is_arf=False)
+                    set_is_arf(False)
                 _clean_rootinfo()
-                # Embedding cache server only run one step.
-                if is_embedding_cache_server:
-                    break
                 should_stop = run_context.get_stop_requested()
                 if should_stop:
                     break
 
-            # When it's distributed training and using MindRT,
-            # the node id should be reset to start from 0.
-            # This is to avoid the timeout when finding the actor route tables in 'train' and 'eval' case(or 'fit').
-            if _enable_distributed_mindrt():
-                _reset_op_id_with_offset()
-
             self._eval_during_train(valid_infos, cb_params, list_callback)
 
             train_dataset.reset()
@@ -1377,9 +1235,7 @@ class Model:
             # if param is cache enable, flush data from cache to host before epoch end
             self._flush_from_cache(cb_params)
 
-            # Embedding cache server need not do epoch end callback, this process only run one step.
-            if not is_embedding_cache_server:
-                list_callback.on_train_epoch_end(run_context)
+            list_callback.on_train_epoch_end(run_context)
             if "metrics" in cb_params or "eval_results" in cb_params:
                 cb_params.pop("metrics", None)
                 cb_params.pop("eval_results", None)
@@ -1456,10 +1312,6 @@ class Model:
         """
         _init_auto_parallel_context(self._network)
         _check_tft()
-        device_target = context.get_context("device_target")
-        if _is_ps_mode() and not _cache_enable() and (device_target in ["Ascend", "CPU"]) and dataset_sink_mode:
-            logger.info("For PS mode, reset datasink mode to False when using Ascend or CPU backend.")
-            dataset_sink_mode = False
 
         Validator.check_bool(dataset_sink_mode)
         if isinstance(self._train_network, nn.GraphCell) and dataset_sink_mode:
@@ -1471,11 +1323,6 @@ class Model:
                              "the value of epoch in train {} separately."
                              .format(train_dataset._warmup_epoch, epoch))
 
-        # Parameter server and embedding cache mode check.
-        if _is_ps_mode():
-            if not dataset_sink_mode and _cache_enable():
-                raise ValueError("Embedding cache mode should run with 'dataset_sink_mode=True'.")
-
         self._check_sink_mode_for_ds_debug_mode(dataset_sink_mode)
 
         Validator.check_is_int(sink_size)
@@ -1506,12 +1353,6 @@ class Model:
                     sink_size=sink_size,
                     initial_epoch=initial_epoch)
 
-        # When it's distributed training and using MindRT,
-        # the node id should be reset to start from 0.
-        # This is to avoid the timeout when finding the actor route tables in 'train' and 'eval' case(or 'fit').
-        if _enable_distributed_mindrt():
-            _reset_op_id_with_offset()
-
         _clear_auto_parallel_context(self._network)
 
     @staticmethod
@@ -1609,10 +1450,6 @@ class Model:
             >>> model.fit(2, train_dataset, valid_dataset)
         """
         _init_auto_parallel_context(self._network)
-        device_target = context.get_context("device_target")
-        if _is_ps_mode() and not _cache_enable() and (device_target in ["Ascend", "CPU"]) and dataset_sink_mode:
-            logger.info("For PS mode, reset datasink mode to False when using Ascend or CPU backend.")
-            dataset_sink_mode = False
 
         dataset_sink_mode = Validator.check_bool(dataset_sink_mode)
         valid_dataset_sink_mode = Validator.check_bool(valid_dataset_sink_mode)
@@ -1906,13 +1743,6 @@ class Model:
 
         self._clear_metrics()
 
-        # Embedding cache server as a storage service, no need to execute eval.
-        is_embedding_cache_server = _is_role_pserver() and _cache_enable()
-        if is_embedding_cache_server:
-            metrics = self._get_metrics()
-            cb_params.metrics = metrics
-            return metrics
-
         if context.get_context("device_target") == "CPU" and dataset_sink_mode:
             dataset_sink_mode = False
             logger.info("CPU cannot support dataset sink mode currently."
@@ -1924,13 +1754,7 @@ class Model:
             else:
                 eval_result = self._eval_process(valid_dataset, list_callback, cb_params)
 
-        # When it's distributed training and using MindRT,
-        # the node id should be reset to start from 0.
-        # This is to avoid the timeout when finding the actor route tables in 'train' and 'eval' case(or 'fit').
-        if _enable_distributed_mindrt():
-            _reset_op_id_with_offset()
         _clear_auto_parallel_context(self._network)
-
         return eval_result
 
     def _predict_lite(self, *predict_data, config=None):
@@ -2181,13 +2005,6 @@ class Model:
         result = self._predict_network(*predict_data)
 
         check_output_data(result)
-
-        # When it's distributed training and using MindRT,
-        # the node id should be reset to start from 0.
-        # This is to avoid the timeout when finding the actor route tables in 'train' and 'eval' case(or 'fit').
-        if _enable_distributed_mindrt():
-            _reset_op_id_with_offset()
-
         return result
 
     def _infer_train_check(self, train_dataset, dataset_sink_mode, sink_size):
@@ -2199,9 +2016,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()))
@@ -2361,9 +2175,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()