mindspore 2.5.0__cp310-cp310-win_amd64.whl → 2.6.0rc1__cp310-cp310-win_amd64.whl

mindspore/profiler/profiler_action_controller.py

@@ -98,7 +98,7 @@ class ProfilerActionController:
         # Handle special case for ProfilerAction.RECORD_AND_SAVE to ProfilerAction.RECORD_AND_SAVE transition
         if start_state == ProfilerAction.RECORD_AND_SAVE and end_state == ProfilerAction.RECORD_AND_SAVE:
             process_action = [self.prof_interface.stop, self.prof_interface.finalize,
-                              self._trace_ready, self.prof_interface.init,
+                              self._trace_ready, self.prof_interface.clear, self.prof_interface.init,
                               self.prof_interface.start]
         else:
             while start_state != end_state:

mindspore/profiler/profiler_interface.py

@@ -20,7 +20,9 @@ from mindspore.common.api import _pynative_executor
 from mindspore.profiler.common.registry import PROFILERS
 from mindspore.profiler.platform.base_profiler import BaseProfiler
 from mindspore.profiler.common.profiler_context import ProfilerContext
+from mindspore.profiler.common.log import ProfilerLogger
 from mindspore.profiler.common.profiler_path_manager import ProfilerPathManager
+from mindspore.profiler.common.profiler_meta_data import ProfilerMetaData
 
 
 class ProfilerInterface:
@@ -92,6 +94,7 @@ class ProfilerInterface:
             logger.warning("ProfilerInterface finalize failed, profiler has not been initialized.")
             return
 
+        ProfilerMetaData.dump_metadata()
         for profiler in cls.platform_profilers_set:
             profiler.finalize()
             profiler = None
@@ -106,6 +109,7 @@ class ProfilerInterface:
             return
         cls.platform_profilers_set.clear()
         cls.is_initialized = False
+        ProfilerLogger.destroy()
         logger.info("ProfilerInterface clear")
 
     @classmethod

mindspore/profiler/schedule.py

@@ -74,10 +74,19 @@ class Schedule:
 
     Keyword Args:
         wait (int): The number of steps to wait before starting the warm-up phase.
+            must be greater than or equal to 0. If the wait parameter is not set externally,
+            it is set to ``0`` when the schedule class is initialized.
         active (int): The number of steps to record data during the active phase.
-        warmup (int, optional): The number of steps to perform the warm-up phase. Default: ``0``.
-        repeat (int, optional): The number of times to repeat the cycle. Default: ``0``.
-        skip_first (int, optional): The number of steps to skip at the beginning. Default: ``0``.
+            must be greater than or equal to 1. If the active parameter is not set externally,
+            it is set to ``1`` when the schedule class is initialized.
+        warmup (int, optional): The number of steps to perform the warm-up phase.
+            must be greater than or equal to 0. Default value: ``0``.
+        repeat (int, optional): The number of times to repeat the cycle.
+            If repeat is set to 0, the Profiler will determine the repeat value based on the number of times the model
+            is trained, which will generate one more performance data with incomplete collection. The data in the last
+            step is abnormal data that users do not need to pay attention to. Default value: ``0``.
+        skip_first (int, optional): The number of steps to skip at the beginning. Must be greater than or equal to 0.
+            Default value: ``0``
 
     Raises:
         ValueError: When the parameter step is less than 0.
@@ -87,10 +96,10 @@ class Schedule:
 
     Examples:
         >>> import numpy as np
-        >>> import mindspore as ms
+        >>> import mindspore
         >>> import mindspore.dataset as ds
-        >>> from mindspore import context, nn, Profiler
-        >>> from mindspore.profiler import schedule, tensor_board_trace_handler
+        >>> from mindspore import context, nn
+        >>> from mindspore.profiler import ProfilerLevel, AicoreMetrics, ExportType, ProfilerActivity
         >>>
         >>> class Net(nn.Cell):
         ...     def __init__(self):
@@ -108,18 +117,34 @@ class Schedule:
         ...     optimizer = nn.Momentum(test_net.trainable_params(), 1, 0.9)
         ...     loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True)
         ...     data = ds.GeneratorDataset(generator_net(), ["data", "label"])
-        ...     model = ms.train.Model(test_net, loss, optimizer)
+        ...     model = mindspore.train.Model(test_net, loss, optimizer)
         ...     model.train(1, data)
         >>>
         >>> if __name__ == '__main__':
-        ...     context.set_context(mode=ms.PYNATIVE_MODE, device_target="Ascend")
+        ...     # If the device_target is GPU, set the device_target to "GPU"
+        ...     context.set_context(mode=mindspore.GRAPH_MODE)
+        ...     mindspore.set_device("Ascend")
         ...
+        ...     # Init Profiler
+        ...     experimental_config = mindspore.profiler._ExperimentalConfig(
+        ...                                 profiler_level=ProfilerLevel.Level0,
+        ...                                 aic_metrics=AicoreMetrics.AiCoreNone,
+        ...                                 l2_cache=False,
+        ...                                 mstx=False,
+        ...                                 data_simplification=False,
+        ...                                 export_type=[ExportType.Text])
+        ...     steps = 10
         ...     net = Net()
-        ...     STEP_NUM = 15
+        ...     # Note that the Profiler should be initialized before model.train
+        ...     with mindspore.profiler.profile(activities=[ProfilerActivity.CPU, ProfilerActivity.NPU],
+        ...                                     schedule=mindspore.profiler.schedule(wait=1, warmup=1, active=2,
+        ...                                           repeat=1, skip_first=2),
+        ...                                     on_trace_ready=mindspore.profiler.tensorboard_trace_handler("./data"),
+        ...                                     profile_memory=False,
+        ...                                     experimental_config=experimental_config) as prof:
         ...
-        ...     with Profiler(schedule=schedule(wait=1, warmup=1, active=2, repeat=1, skip_first=2),
-        ...                   on_trace_ready=tensor_board_trace_handler) as prof:
-        ...         for i in range(STEP_NUM):
+        ...         # Train Model
+        ...         for step in range(steps):
         ...             train(net)
         ...             prof.step()
     """
@@ -160,25 +185,35 @@ class Schedule:
             return ProfilerAction.WARM_UP
         return ProfilerAction.RECORD if mod_step < num_steps - 1 else ProfilerAction.RECORD_AND_SAVE
 
+    def __repr__(self):
+        return (f"Schedule(wait={self.wait!r}, active={self.active!r}, "
+                f"warmup={self.warmup!r}, repeat={self.repeat!r}, "
+                f"skip_first={self.skip_first!r})")
+
     def _check_params(self):
         """
         Verify all parameters in the schedule,
         and set them to default values if the parameters are not compliant.
         """
-        if not isinstance(self.wait, int) or self.wait < 0:
-            logger.warning("Invalid parameter wait, reset it to 0.")
+        if not isinstance(self.wait, int) or isinstance(self.wait, bool) or self.wait < 0:
+            logger.warning(f"Parameter 'wait' should be of type int, but got "
+                           f"{type(self.wait).__name__}. reset to int 0.")
             self.wait = 0
-        if not isinstance(self.warmup, int) or self.warmup < 0:
-            logger.warning("Invalid parameter warmup, reset it to 0.")
+        if not isinstance(self.warmup, int) or isinstance(self.warmup, bool) or self.warmup < 0:
+            logger.warning(f"Parameter 'warmup' should be of type int, but got "
+                           f"{type(self.warmup).__name__}. reset to int 0.")
             self.warmup = 0
-        if not isinstance(self.active, int) or self.active <= 0:
-            logger.warning("Invalid parameter active, reset it to 1.")
+        if not isinstance(self.active, int) or isinstance(self.active, bool) or self.active <= 0:
+            logger.warning(f"Parameter 'active' should be of type int, but got "
+                           f"{type(self.active).__name__}. reset to int 1.")
             self.active = 1
-        if not isinstance(self.repeat, int) or self.repeat < 0:
-            logger.warning("Invalid parameter repeat, reset it to 0.")
+        if not isinstance(self.repeat, int) or isinstance(self.repeat, bool) or self.repeat < 0:
+            logger.warning(f"Parameter 'repeat' should be of type int, but got "
+                           f"{type(self.repeat).__name__}. reset to int 0.")
             self.repeat = 0
-        if not isinstance(self.skip_first, int) or self.skip_first < 0:
-            logger.warning("Invalid parameter skip_first, reset it to 0.")
+        if not isinstance(self.skip_first, int) or isinstance(self.skip_first, bool) or self.skip_first < 0:
+            logger.warning(f"Parameter 'skip_first' should be of type int, but got "
+                           f"{type(self.skip_first).__name__}. reset to int 0.")
             self.skip_first = 0
         if self.warmup == 0:
             logger.warning("Profiler won't be using warmup, this can skew profiler results")

mindspore/rewrite/api/node.py

@@ -28,9 +28,11 @@ from .scoped_value import ScopedValue
 
 class Node:
     """
-    A node is a data structure that expresses source code statements in a network.
-
-    Each node usually corresponds to a statement in expanded forward evaluation process.
+    A node (Node) can be understood as a basic data structure unit in the computational graph of a neural network,
+    which represents an operation or computational step in the network.
+    Each node usually corresponds to a statement or expression in the source code,
+    which contains the information needed to perform the operation,
+    such as the type of operation, input data, output result, and connection relationships with other nodes.
 
     Nodes can express a ``Cell`` call statement, a ``Primitive`` call statement, an arithmetic operation statement, a
     return statements, etc. of the forward calculation process.
@@ -66,9 +68,9 @@ class Node:
                 source code.
             args (List[ScopedValue]): Indicate input names. Used as args of a call expression of an assign statement in
                 source code. Default: ``None`` , which indicates the `cell` has no args inputs.
-            kwargs (Dict[str, ScopedValue]): Type of key must be `str` and type of value must be `ScopedValue`.
-                Indicate keyword input names. Used as kwargs of a call expression of an assign statement in source
-                code. Default: ``None`` , which indicates the `cell` has no kwargs inputs.
+            kwargs (Dict[str, ScopedValue]): Used as kwargs of a call expression of an assign statement in source
+                code. Indicate keyword input names. Type of key must be `str` and type of value must be `ScopedValue`.
+                Default: ``None`` , which indicates the `cell` has no kwargs inputs.
             name (str): Indicate the name of node. Used as field name in source code. Default is None. Rewrite will
                 generate name from `cell` when name is None. Rewrite will check and ensure the uniqueness of `name`
                 while node being inserted. Default: ``""`` .
@@ -83,7 +85,7 @@ class Node:
             TypeError: If `targets` is not `list`.
             TypeError: If the type of `targets` is not in `[ScopedValue, str]`.
             TypeError: If arg in `args` is not a `ScopedValue`.
-            TypeError: If key of `kwarg` is not a str or value of kwarg in `kwargs` is not a `ScopedValue`.
+            TypeError: If key of `kwargs` is not a str or value of kwarg in `kwargs` is not a `ScopedValue`.
 
         Examples:
             >>> from mindspore.rewrite import SymbolTree, ScopedValue
@@ -125,9 +127,9 @@ class Node:
                 source code.
             args (List[ScopedValue]): Indicate input names. Used as args of a call expression of an assign statement in
                 source code. Default: ``None`` , which indicates the `function` has no args inputs.
-            kwargs (Dict[str, ScopedValue]): Type of key must be `str` and type of value must be `ScopedValue`.
-                Indicate keyword input names. Used as kwargs of a call expression of an assign statement in source
-                code. Default: ``None`` , which indicates the `function` has no kwargs inputs.
+            kwargs (Dict[str, ScopedValue]): Used as kwargs of a call expression of an assign statement in source
+                code. Indicate keyword input names. Type of key must be `str` and type of value must be `ScopedValue`.
+                Default: ``None`` , which indicates the `function` has no kwargs inputs.
 
         Returns:
             An instance of `Node`.
@@ -137,7 +139,7 @@ class Node:
             TypeError: If `targets` is not `list`.
             TypeError: If the type of `targets` is not in `[ScopedValue, str]`.
             TypeError: If arg in `args` is not a `ScopedValue`.
-            TypeError: If key of `kwarg` is not a str or value of kwarg in `kwargs` is not a `ScopedValue`.
+            TypeError: If key of `kwargs` is not a str or value of kwarg in `kwargs` is not a `ScopedValue`.
 
         Examples:
             >>> from mindspore.rewrite import SymbolTree, ScopedValue
@@ -476,14 +478,14 @@ class Node:
             ...         self.relu = nn.ReLU()
             ...
             ...     def construct(self, input):
-            ...         output = self.relu(x=input)
+            ...         output = self.relu(input=input)
             ...         return output
             >>>
             >>> net = ReLUNet()
             >>> stree = SymbolTree.create(net)
             >>> node = stree.get_node("relu")
             >>> print(node.get_kwargs())
-            {'x': input}
+            {'input': input}
         """
         return self._node.get_kwargs()
 

mindspore/rewrite/api/symbol_tree.py

@@ -74,7 +74,7 @@ class SymbolTree:
 
         - :class:`mindspore.nn.SequentialCell`
         - Functions(Excludes Python built-in functions and third-party library functions)
-        - Control flow statements, such as `if` statements
+        - Control flow statements(such as `if` statements)
 
         Note:
             Because the specific execution branch of control flows are still unknown during the rewrite operation

mindspore/run_check/_check_version.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2022 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.
@@ -26,7 +26,6 @@ from packaging import version
 import numpy as np
 from mindspore import log as logger
 from mindspore.log import vlog_print
-from mindspore._c_expression import MSContext, ms_ctx_param
 from ..version import __version__
 
 
@@ -259,7 +258,7 @@ class AscendEnvChecker(EnvChecker):
 
     def __init__(self, library_path):
         self.library_path = library_path
-        self.version = ["7.5", "7.6"]
+        self.version = ["7.6", "7.7"]
 
         # env
         self.path = os.getenv("PATH")
@@ -276,6 +275,7 @@ class AscendEnvChecker(EnvChecker):
         self.python_path_check = "opp/built-in/op_impl/ai_core/tbe"
         self.ld_lib_path_check_fwk = "/lib64"
         self.ascend_opp_path_check = "/op"
+        self.ascend_opp_kernel_path_check = "/opp_kernel"
         self.v = ""
 
     @staticmethod
@@ -342,27 +342,35 @@ class AscendEnvChecker(EnvChecker):
                 logger.warning(f"MindSpore version {mindspore_version} and \"te\" wheel package version {v} does not "
                                "match. For details, refer to the installation guidelines: "
                                "https://www.mindspore.cn/install")
-            from hccl import sys_version as hccl_version
-            v = '.'.join(hccl_version.__sys_version__.split('.')[0:2])
-            if v not in supported_version:
-                attention_warning = True
-                logger.warning(f"MindSpore version {mindspore_version} and \"hccl\" wheel package version {v} does not "
-                               "match. For details, refer to the installation guidelines: "
-                               "https://www.mindspore.cn/install")
         # DO NOT modify exception type to any other, you DO NOT know what kind of exceptions the te will throw.
         # pylint: disable=broad-except
         except Exception as e:
             logger.error(f"CheckFailed: {e}")
-            logger.error("MindSpore relies on whl packages of \"te\" and \"hccl\" in the \"latest\" "
-                         "folder of the Ascend AI software package (Ascend Data Center Solution). Please check whether"
-                         " they are installed correctly or not, refer to the match info on: "
-                         "https://www.mindspore.cn/install")
+            logger.critical("MindSpore relies on whl packages of \"te\" in the \"latest\" folder of the "
+                            "Ascend AI software package (Ascend Data Center Solution). Please check whether they are "
+                            "installed correctly or not, refer to the match info on: https://www.mindspore.cn/install")
         if attention_warning:
             warning_countdown = 3
             for i in range(warning_countdown, 0, -1):
                 logger.warning(f"Please pay attention to the above warning, countdown: {i}")
                 time.sleep(1)
 
+    def check_opp_kernel(self):
+        """
+            opp kernel install check
+        """
+        from mindspore._c_expression import MSContext
+        soc_version = MSContext.get_instance().get_ascend_soc_version()
+        if soc_version == "ascend310":
+            return
+
+        opp_kernel_path = self.ascend_opp_path.replace("opp", "opp_kernel")
+        if not os.path.exists(opp_kernel_path):
+            logger.critical("MindSpore relies on \"Ascend opp_kernel\" folder of the Ascend AI software package ("
+                            "Ascend Data Center Solution). Please check whether they are installed correctly or not, "
+                            "refer to the match info on: https://www.mindspore.cn/install")
+            raise Exception("Ascend opp_kernel is not installed")
+
     def set_env(self):
         curr_path = os.path.realpath(os.path.dirname(__file__))
         cust_aicpu_path = os.path.realpath(os.path.join(curr_path, "../lib/plugin/ascend/custom_aicpu_ops"))
@@ -391,6 +399,8 @@ class AscendEnvChecker(EnvChecker):
         # check te version after set te env
         self.check_deps_version()
 
+        self.check_opp_kernel()
+
     def _check_env(self):
         """ascend dependence path check"""
         if self.path is None or self.path_check not in self.path:
@@ -467,6 +477,7 @@ def check_version_and_env_config():
             logger.warning("Pre-Load Library libgomp.so.1 failed, which might cause TLS memory allocation failure. If "
                            "the failure occurs, please refer to the FAQ for a solution: "
                            "https://www.mindspore.cn/docs/en/master/faq/installation.html.")
+        from mindspore._c_expression import MSContext, ms_ctx_param
         MSContext.get_instance().register_check_env_callback(check_env)
         MSContext.get_instance().register_set_env_callback(set_env)
         MSContext.get_instance().set_device_target_inner(MSContext.get_instance().get_param(ms_ctx_param.device_target))

mindspore/run_check/run_check.py

@@ -47,7 +47,7 @@ def _check_mul():
 def run_check():
     """
     Provide a convenient API to check if the installation is successful or failed.
-    If the version in the check result is not what you expect, use :func:`mindspore.set_context` to set device_target
+    If the version in the check result is not as expected, use :func:`mindspore.set_context` to set device_target
     before run_check().
 
     Examples:

mindspore/runtime/__init__.py

@@ -17,7 +17,7 @@
 The runtime interface.
 """
 
-from mindspore.runtime.executor import launch_blocking, dispatch_threads_num, set_cpu_affinity
+from mindspore.runtime.executor import launch_blocking, dispatch_threads_num, set_cpu_affinity, set_kernel_launch_group
 from mindspore.runtime.memory import set_memory, memory_stats, memory_reserved, max_memory_reserved, empty_cache,\
                                  reset_peak_memory_stats, memory_summary, memory_allocated,\
                                  max_memory_allocated, reset_max_memory_reserved, reset_max_memory_allocated
@@ -27,7 +27,7 @@ from mindspore.runtime.event import Event
 from .executor import launch_blocking
 
 __all__ = [
-    "launch_blocking", "dispatch_threads_num", "set_cpu_affinity",
+    "launch_blocking", "dispatch_threads_num", "set_cpu_affinity", "set_kernel_launch_group",
     "Stream", "communication_stream", "synchronize", "set_cur_stream", "current_stream", "default_stream", "StreamCtx",
     "set_memory", "memory_stats", "memory_reserved", "max_memory_reserved", "empty_cache", "reset_peak_memory_stats",
     "memory_summary", "memory_allocated", "max_memory_allocated", "reset_max_memory_reserved",

mindspore/runtime/executor.py

@@ -24,18 +24,14 @@ from mindspore import log as logger
 
 def launch_blocking():
     """
-    Whether to enable synchronous execution.
-    Default: ``False`` . When the value is set to ``False`` , the operator is executed asynchronously on the
-    device. When an error occurs in the execution of the operator, the specific error script code location
-    cannot be located. When this API is called , the value will be changed to true, the operator is executed
-    synchronously on the device. It will reduce the execution performance of the program. At this time, when
-    an error occurs in the execution of the operator, the location of the error script code can be located
-    according to the call stack of the error.
+    Indicates that synchronizing the execution of the startup device reduces the execution performance of the program.
 
-    Note:
-        - No parameters are required.
-        - By default, operations are executed asynchronously.
-        - Calling this function enables synchronous execution.
+    - In the initial state when this interface is not called, the operator executes asynchronously on the device.
+      In this case, when an error occurs in the execution of the operator,
+      it will not be possible to locate the position of the particular error script code.
+    - When this interface is called, the operator is executed in a synchronized manner on the device.
+      At this point, when an error occurs in the execution of the operator,
+      the location of the erroneous script code can be located based on the error call stack.
 
     Examples:
         >>> import mindspore as ms
@@ -146,3 +142,36 @@ def set_cpu_affinity(enable_affinity, affinity_cpu_list=None):
     else:
         RuntimeConf.get_instance().set_thread_bind_core_configured()
         return
+
+
+@args_type_check(thread_num=int, kernel_group_num=int)
+def set_kernel_launch_group(thread_num=2, kernel_group_num=8):
+    """
+    O0 mode supports operator batch parallel delivery interface, supports enabling
+    parallel delivery, and configures parallel number.
+
+    Args:
+        thread_num (int, optional): The number of concurrent threads, generally not recommended
+            to increase. The `thread_num` and the number of threads configured by the existing interface
+            mindspore.runtime.dispatch_threads_num are independent of each other. Default value is ``2``.
+        kernel_group_num (int, optional): Total number of operator groups,
+            kernel_group_num/thread_num groups per thread. Default value is ``8``.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> ms.runtime.set_kernel_launch_group(thread_num=2, kernel_group_num=8)
+    """
+    if RuntimeConf.get_instance().is_kernel_launch_group_configured():
+        raise RuntimeError("The 'kernel_launch_group' can not be set repeatedly.")
+
+    if thread_num < 1:
+        raise ValueError(f"The value of thread_num should be at least 1, but got {thread_num}")
+
+    if kernel_group_num < 1:
+        raise ValueError(f"The value of kernel_group_num should be at least 1, but got {kernel_group_num}")
+
+    if (kernel_group_num % thread_num) != 0:
+        raise ValueError(f"Invalid parameter value, kernel_group_num: {kernel_group_num} cannot "
+                         f"be evenly divisible by thread_num: {thread_num}")
+
+    return RuntimeConf.get_instance().set_kernel_launch_group(thread_num, kernel_group_num)

mindspore/runtime/memory.py

@@ -23,11 +23,19 @@ from mindspore import log as logger
 import mindspore as ms
 
 _MEMORY_PATTERN = r'[1-9][0-9]*(\.)?[0-9]*GB|0\.[0-9]*GB'
+_RESERVE_PATTERN = r'[0-9][0-9]*(\.)?[0-9]*GB|0\.[0-9]*GB'
 _device_context_mgr = DeviceContextManager.get_instance()
 
 
-@args_type_check(init_size=str, increase_size=str, max_size=str, optimize_level=str)
-def set_memory(init_size="2GB", increase_size="2GB", max_size="1024GB", optimize_level="O0"):
+@args_type_check(
+    init_size=str,
+    increase_size=str,
+    max_size=str,
+    optimize_level=str,
+    huge_page_reserve_size=str,
+)
+def set_memory(init_size="2GB", increase_size="2GB", max_size="1024GB", optimize_level="O0",
+               huge_page_reserve_size="0GB"):
     """
     Set the memory parameters of runtime device memory management that is implemented using a memory pool.
 
@@ -41,6 +49,7 @@ def set_memory(init_size="2GB", increase_size="2GB", max_size="1024GB", optimize
             The actual used memory size is the minimum of the available memory of the device and max_device_memory.
             The format is "xxGB". Default is the maximum available memory of the device, expressed as ``1024GB``.
         optimize_level (str): The memory optimize level. The value must be in ['O0', 'O1']. Default: ``O0`` .
+        huge_page_reserve_size (str): The reserved size of huge page memory. The format is "xxGB". Default: ``0GB``.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -48,7 +57,7 @@ def set_memory(init_size="2GB", increase_size="2GB", max_size="1024GB", optimize
     Examples:
         >>> import mindspore as ms
         >>> ms.set_device("Ascend", 1)
-        >>> ms.runtime.set_memory("10GB", "2GB", "60GB", "O1")
+        >>> ms.runtime.set_memory("10GB", "2GB", "60GB", "O1", "0GB")
     """
     if RuntimeConf.get_instance().is_memory_configured():
         raise RuntimeError("The 'set_memory' can not be set repeatedly.")
@@ -56,9 +65,11 @@ def set_memory(init_size="2GB", increase_size="2GB", max_size="1024GB", optimize
     _check_memory_conf_valid(init_size)
     _check_memory_conf_valid(increase_size)
     _check_memory_conf_valid(max_size)
+    Validator.check_str_by_regular(huge_page_reserve_size, _RESERVE_PATTERN)
     init_value = float(init_size[:-2])
     increase_value = float(increase_size[:-2])
     max_value = float(max_size[:-2])
+    huge_page_reserve_value = float(huge_page_reserve_size[:-2])
 
     memory_optimize_levels = ["O0", "O1"]
     if optimize_level not in memory_optimize_levels:
@@ -68,19 +79,25 @@ def set_memory(init_size="2GB", increase_size="2GB", max_size="1024GB", optimize
     if optimize_level == "O1":
         optimize_value = 1
 
-    return RuntimeConf.get_instance().set_memory(init_value, increase_value, max_value, optimize_value)
+    return RuntimeConf.get_instance().set_memory(
+        init_value,
+        increase_value,
+        max_value,
+        optimize_value,
+        huge_page_reserve_value,
+    )
 
 
 def _check_memory_conf_valid(memory_size):
     """
-    Check whether the configuration memory value format is "xxGB" and can not be "0G".
+    Check whether the configuration memory value format is "xxGB" and can not be "0GB".
     """
     if not Validator.check_str_by_regular(memory_size, _MEMORY_PATTERN):
         raise ValueError("The memory value should be in correct format!"
                          "It must be a string ending with 'GB', in addition to that, it must contain "
                          "only numbers or decimal points, such as \"5GB\" or \"3.5GB\", but got {}."
                          .format(memory_size))
-    if memory_size == "0G" or memory_size == "0.0G":
+    if memory_size == "0GB" or memory_size == "0.0GB":
         raise ValueError("The memory value should not be \"0GB\".")
 
 def _is_initialized(device_target):
@@ -97,7 +114,7 @@ def memory_stats():
     Returns status information queried from the memory pool.
 
     Note:
-        - For the `CPU` backend, a dictionary with empty data is always returned.
+        For the `CPU` backend, a dictionary with empty data is always returned.
 
     Returns:
         dict, the queried memory information.
@@ -196,7 +213,7 @@ def empty_cache():
 
     Note:
         Currently, the MindSpore memory pool does not have the function of releasing memory fragments.
-        This interface is reserved but implemented as an empty method and prompted in log mode.
+        This interface is reserved but implemented as an empty method and prompted in log mode when using.
     """
     logger.warning(f"The empty_cache operation is currently not supported.")
 

mindspore/safeguard/rewrite_obfuscation.py

@@ -129,7 +129,7 @@ def obfuscate_ckpt(network, ckpt_files, target_modules=None, obf_config=None, sa
     Args:
         network (nn.Cell): The original network that need to be obfuscated.
         ckpt_files (str): The directory path of original ckpt files.
-        target_modules (list[str]): The target ops that need to be obfuscated in the network. The first string
+        target_modules (list[str], optional): The target ops that need to be obfuscated in the network. The first string
             represents the network path of the target ops in the original network, which should be in form of
             ``"A/B/C"``. The second string represents the names of multiple target ops in the same path, which
             should be in form of ``"D|E|F"``. For example, the target_modules of GPT2 can be ``['backbone/blocks
@@ -137,10 +137,11 @@ def obfuscate_ckpt(network, ckpt_files, target_modules=None, obf_config=None, sa
             format of 'obfuscate_layers:all' or 'obfuscate_layers:int', which represents the number of layers
             need to be obfuscated of duplicate layers (such as transformer layers or resnet blocks).
             Default: ``None``.
-        obf_config (dict): The configuration of model obfuscation polices. Default: ``None``.
-        saved_path (str): The directory path for saving obfuscated ckpt files. Default: ``'./'``.
-        obfuscate_scale (Union[float, int]): Obfuscate scale of weights. The generated random obf_ratios will be in
-            range of (1 / obfuscate_scale, obfuscate_scale). Default: 100.
+        obf_config (dict, optional): The configuration of model obfuscation polices. Default: ``None``.
+        saved_path (str, optional): The directory path for saving obfuscated ckpt files. Default: ``'./'``.
+        obfuscate_scale (Union[float, int], optional): Obfuscate scale of weights.
+            The generated random obf_ratios will be in
+            range of (1 / obfuscate_scale, obfuscate_scale). Default: ``100``.
 
     Returns:
         dict[str], obf_metadata, which is the necessary data that needs to be load when running obfuscated network.
@@ -371,7 +372,8 @@ def load_obf_params_into_net(network, target_modules=None, obf_ratios=None, obf_
 
     Args:
         network (nn.Cell): The original network that need to be obfuscated.
-        target_modules (list[str]): The target ops that need to be obfuscated in the network. The first string
+        target_modules (list[str], optional): The target ops that need to be obfuscated in the network.
+            The first string
             represents the network path of the target ops in the original network, which should be in form of
             ``"A/B/C"``. The second string represents the names of multiple target ops in the same path, which
             should be in form of ``"D|E|F"``. For example, thr target_modules of GPT2 can be ``['backbone
@@ -379,9 +381,10 @@ def load_obf_params_into_net(network, target_modules=None, obf_ratios=None, obf_
             in the format of 'obfuscate_layers:all' or 'obfuscate_layers:int', which represents the number of
             layers need to be obfuscated of duplicate layers (such as transformer layers or resnet blocks).
             Default: ``None``.
-        obf_ratios (Tensor): The obf ratios generated when execute :func:`mindspore.obfuscate_ckpt`. Default: ``None``.
-        obf_config (dict): The configuration of model obfuscation polices. Default: ``None``.
-        data_parallel_num (int): The data parallel number of parallel training. Default: 1.
+        obf_ratios (Tensor, optional): The obf ratios generated when execute :func:`mindspore.obfuscate_ckpt`.
+            Default: ``None``.
+        obf_config (dict, optional): The configuration of model obfuscation polices. Default: ``None``.
+        data_parallel_num (int, optional): The data parallel number of parallel training. Default: ``1``.
         kwargs (dict): Configuration options dictionary.
 
             - ignored_func_decorators (list[str]): The name list of function decorators in network's python code.

mindspore/train/__init__.py

@@ -25,12 +25,12 @@ from mindspore.train import amp
 from mindspore.train.amp import build_train_network
 from mindspore.train.loss_scale_manager import LossScaleManager, FixedLossScaleManager, DynamicLossScaleManager
 from mindspore.train.serialization import save_checkpoint, load_checkpoint, load_param_into_net, export, \
-    load, parse_print, build_searched_strategy, merge_sliced_parameter, load_distributed_checkpoint, \
-    async_ckpt_thread_status, restore_group_info_list, convert_model, obfuscate_model, export_split_mindir, \
-    load_checkpoint_async, check_checkpoint, get_ckpt_path_with_strategy, ckpt_to_safetensors, safetensors_to_ckpt
+    load, parse_print, async_ckpt_thread_status, convert_model, export_split_mindir, \
+    load_checkpoint_async, check_checkpoint, get_ckpt_path_with_strategy, ckpt_to_safetensors, safetensors_to_ckpt, \
+    build_searched_strategy, merge_sliced_parameter, load_distributed_checkpoint, restore_group_info_list
 from mindspore.train.callback import Callback, LossMonitor, TimeMonitor, ModelCheckpoint, SummaryCollector, \
     CheckpointConfig, RunContext, LearningRateScheduler, SummaryLandscape, FlopsUtilizationCollector, \
-    History, LambdaCallback, ReduceLROnPlateau, EarlyStopping, OnRequestExit, BackupAndRestore, TFTRegister
+    History, LambdaCallback, ReduceLROnPlateau, EarlyStopping, OnRequestExit, BackupAndRestore, TrainFaultTolerance
 from mindspore.train.summary import SummaryRecord
 from mindspore.train.train_thor import ConvertNetUtils, ConvertModelUtils
 from mindspore.train.metrics import *
@@ -38,10 +38,10 @@ from mindspore.train.data_sink import data_sink
 
 __all__ = ["Model", "DatasetHelper", "connect_network_with_dataset", "build_train_network", "LossScaleManager",
            "FixedLossScaleManager", "DynamicLossScaleManager", "save_checkpoint", "load_checkpoint", "check_checkpoint",
-           "load_param_into_net", "export", "load", "export_split_mindir", "parse_print", "build_searched_strategy",
-           "merge_sliced_parameter", "load_distributed_checkpoint", "async_ckpt_thread_status",
-           "restore_group_info_list", "convert_model", "data_sink", "obfuscate_model", "load_checkpoint_async",
-           "get_ckpt_path_with_strategy", "ckpt_to_safetensors", "safetensors_to_ckpt"]
+           "load_param_into_net", "export", "load", "export_split_mindir", "parse_print", "async_ckpt_thread_status",
+           "convert_model", "data_sink", "load_checkpoint_async", "get_ckpt_path_with_strategy", "ckpt_to_safetensors",
+           "safetensors_to_ckpt", "build_searched_strategy", "merge_sliced_parameter", "load_distributed_checkpoint",
+           "restore_group_info_list"]
 __all__.extend(callback.__all__)
 __all__.extend(summary.__all__)
 __all__.extend(train_thor.__all__)