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

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,13 @@ 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")
+        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)
     elif "ForceStopError" in e_str:
         logger.warning("uce wrapper caught RuntimeError ForceStopError")
         force_stop_err = tft.ReportState.RS_NORMAL.value
@@ -246,9 +253,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 +266,12 @@ 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()
+                        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,
                                                                                                    self.batch_num))
@@ -274,6 +281,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 +289,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)
@@ -443,6 +451,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 +574,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 +776,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)
@@ -1057,7 +1072,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 +1141,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:
@@ -1340,6 +1356,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
@@ -2303,6 +2320,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 +2330,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))