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

mindspore/train/model.py

@@ -57,8 +57,10 @@ 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.parallel._utils import _init_auto_parallel_context, _clear_auto_parallel_context
+from .serialization import load_param_into_net
 
 def _transfer_tensor_to_tuple(inputs):
     """
@@ -130,7 +132,8 @@ def _handle_exception_info(obj, uce_env, tft, e):
     if not uce_env:
         logger.error("uce wrapper caught RuntimeError but uce not enable, enter MindIO TTP process.",
                      exc_info=True)
-        tft.tft_report_error(tft.ReportState.RS_UNKNOWN.value)
+        if tft:
+            tft.tft_report_error(tft.ReportState.RS_UNKNOWN.value)
         raise e
     e_str = str(e)
     logger.warning("uce wrapper caught RuntimeError e_str:{}".format(e_str))
@@ -151,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
@@ -165,6 +171,69 @@ def _handle_exception_info(obj, uce_env, tft, e):
         raise e
 
 
+def _handle_training_result_error(model, tft_obj):
+    """
+    Handle training result error for resuming training.
+    """
+    ckpt_load_fn = tft_obj.ckpt_load_func
+    train_network = tft_obj.cb_params.train_network
+    logger.warning("Process training result error start.")
+    # 1. Clear tdt channel
+    logger.warning("Clean tdt channel.")
+    clean_tdt_channel()
+
+    # 2. Load checkpoint
+    logger.warning("Load checkpoint.")
+    new_param_dict, remove_redundancy = ckpt_load_fn()
+    param_not_load, ckpt_not_load = load_param_into_net(train_network, new_param_dict, True, remove_redundancy)
+    logger.warning(f"param_not_load: {param_not_load}")
+    logger.warning(f"ckpt_not_load: {ckpt_not_load}")
+    resume_epoch = new_param_dict.get('epoch_num')
+    resume_step = new_param_dict.get('step_num')
+    model._initial_step = int(resume_step.asnumpy())
+    logger.warning("Process training result error end.")
+    return (resume_epoch, resume_step)
+
+
+def _calc_cb_initial_step(org_epoch, org_step, *args, **kwargs):
+    """calculate initial step for callback"""
+    train_dataset = args[1]
+    dataset_sink_mode = args[3] if len(args) > 3 else kwargs.get('dataset_sink_mode', True)
+    sink_size = args[4] if len(args) > 4 else kwargs.get('sink_size', -1)
+
+    cb_initial_step = 0
+    if dataset_sink_mode:
+        train_dataset.set_init_step(org_epoch)
+        dataset_size = train_dataset.get_dataset_size()
+        if sink_size != -1:
+            cb_initial_step = org_epoch * sink_size + org_step
+        else:
+            cb_initial_step = org_epoch * dataset_size + org_step
+    else:
+        train_dataset.set_init_step(org_step)
+        cb_initial_step = org_step
+    if hasattr(train_dataset, '_dataset_helper'):
+        dataset_helper = train_dataset._dataset_helper
+        _reset_training_dataset(cb_initial_step, dataset_helper.iter.dataset.get_dataset_size())
+    return cb_initial_step
+
+
+def _update_ckpt_callback_info(resume_train_step, **kwargs):
+    """
+    Update checkpoint callback internal state
+    """
+    ckpt_obj = None
+    if kwargs.get('callbacks') and isinstance(kwargs.get('callbacks'), ModelCheckpoint):
+        ckpt_obj = kwargs.get('callbacks')
+    if kwargs.get('callbacks') and isinstance(kwargs.get('callbacks'), list):
+        for item in kwargs.get('callbacks'):
+            if isinstance(item, ModelCheckpoint):
+                ckpt_obj = item
+    if ckpt_obj is not None:
+        ckpt_obj._last_triggered_step = 0
+        ckpt_obj._append_step_num = resume_train_step
+
+
 def _handle_tft(func):
     """
     Decorator function, which starts uce handle process when an exception occurs during training.
@@ -180,42 +249,34 @@ 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:
                     return func(self, *args, **kwargs)
                 except RuntimeError as e:
-                    _handle_exception_info(obj, uce_env, tft, e)
-                    ret = tft.tft_wait_next_action()
-                    if ret == tft.Action.EXIT.value:
-                        raise e
-                    repair_step = tft.tft_get_repair_step()
-                    logger.warning(
-                        "uce wrapper caught repair finish REPAIR STEP: {} batch_num:{}".format(repair_step,
-                                                                                               self.batch_num))
+                    if tre_env and 'TREError' in str(e):
+                        _, resume_step = _handle_training_result_error(self, obj)
+                        repair_step = int(resume_step.asnumpy())
+                        _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, obj.tft, e)
+                        ret = obj.tft.tft_wait_next_action()
+                        if ret == obj.tft.Action.EXIT.value:
+                            raise e
+                        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))
                     initial_epoch = int(repair_step / self.batch_num)
                     initial_step = repair_step % self.batch_num
                     kwargs["initial_epoch"] = initial_epoch
-
-                    train_dataset = args[1]
-                    dataset_sink_mode = args[3] if len(args) > 3 else kwargs.get('dataset_sink_mode', True)
-                    sink_size = args[4] if len(args) > 4 else kwargs.get('sink_size', -1)
-
-                    cb_initial_step = 0
-                    if dataset_sink_mode:
-                        train_dataset.set_init_step(initial_epoch)
-                        dataset_size = train_dataset.get_dataset_size()
-                        if sink_size != -1:
-                            cb_initial_step = initial_epoch * sink_size + initial_step
-                        else:
-                            cb_initial_step = initial_epoch * dataset_size + initial_step
-                    else:
-                        train_dataset.set_init_step(initial_step)
-                        cb_initial_step = initial_step
-
-                    kwargs["initial_step"] = cb_initial_step
+                    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(
@@ -223,8 +284,9 @@ def _handle_tft(func):
                                                                                                      cb_initial_step))
                     continue
                 except BaseException as e:
-                    logger.error("uce wrapper caught BaseException error, enter MindIO TTP process.", exc_info=True)
-                    tft.tft_report_error(tft.ReportState.RS_UNKNOWN.value)
+                    if obj.tft:
+                        logger.error("uce wrapper caught BaseException error, enter MindIO TTP process.", exc_info=True)
+                        obj.tft.tft_report_error(obj.tft.ReportState.RS_UNKNOWN.value)
                     raise e
         else:
             return func(self, *args, **kwargs)
@@ -241,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")
@@ -384,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.
@@ -501,6 +565,10 @@ class Model:
         self._lite_infer = True  # if backend lite infer fails, set False
         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):
@@ -700,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") 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)
@@ -744,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},"
@@ -812,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.")
@@ -957,6 +1025,7 @@ class Model:
         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.initial_step = self._initial_step
 
         # build callback list
         with _CallbackManager(callbacks) as list_callback:
@@ -995,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, "
@@ -1064,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:
@@ -1142,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):
@@ -1278,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
@@ -2120,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()))
@@ -2241,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
@@ -2250,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))
@@ -2281,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()