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

mindspore/common/hook_handle.py

@@ -15,6 +15,7 @@
 """The removable handle for cell hook function."""
 from __future__ import absolute_import
 import weakref
+from collections import OrderedDict
 from mindspore._c_expression import TensorPy as Tensor_
 from mindspore._check_jit_forbidden_api import jit_forbidden_register
 
@@ -173,3 +174,62 @@ class HookHandle:
             extra_dict = self.extra_dict_ref()
             if extra_dict is not None and self.handle_id in extra_dict:
                 del extra_dict[self.handle_id]
+
+
+def _check_hook_results(pre_res, new_res, hook_fn):
+    if not isinstance(new_res, tuple):
+        raise RuntimeError(f"hook {hook_fn.__name__} should return a tuple of grad.")
+
+    new_res_len = len(new_res)
+    pre_res_len = len(pre_res)
+    if new_res_len != pre_res_len:
+        raise RuntimeError(
+            f"hook {hook_fn.__name__} returned incorrect length {new_res_len}, expected {pre_res_len}."
+        )
+
+
+class _HookUtils:
+    r"""
+    Internal utility class for hook registration and execution.
+    """
+
+    @staticmethod
+    def register_hook(hook_dict, hook_fn):
+        """
+        Register hook
+
+        Args:
+            hook_dict (dict): hook dict.
+            hook_fn (function): hook function.
+
+        Returns:
+            tuple: Updated hook_dict and HookHandle object.
+        """
+        if hook_dict is None:
+            hook_dict = OrderedDict()
+        handle = HookHandle(hook_dict)
+        hook_dict[handle.handle_id] = hook_fn
+        return hook_dict, handle
+
+    @staticmethod
+    def run_hook(hook_dict, args):
+        """
+        Run all hooks in the hook_dict with the given arguments.
+
+        Args:
+            hook_dict (dict): Dictionary of registered hooks.
+            args (tuple): Arguments to pass to the hook functions.
+
+        Returns:
+            Modified first argument if any hook returns a new value; otherwise, None.
+        """
+        is_modify = False
+        args_list = list(args)
+        # Note: We create a list from hook_dict.values() to ensure safe iteration.
+        for hook_fn in list(hook_dict.values()):
+            res = hook_fn(*args_list)
+            if res is not None:
+                _check_hook_results(args_list[0], res, hook_fn)
+                args_list[0] = res
+                is_modify = True
+        return args_list[0] if is_modify else None

mindspore/common/jit_config.py

@@ -27,7 +27,11 @@ class JitConfig:
               adopt KernelByKernel execution mode.
             - ``"O1"``: Using commonly used optimizations and automatic operator fusion optimizations,
               adopt KernelByKernel execution mode.
-            - ``"O2"``: Ultimate performance optimization, adopt Sink execution mode.
+            - ``"O2"``: Utilizes the GraphEngine, a graph compilation and execution engine within CANN,
+              for Ascend model compilation and execution. Note: O2 only supports GRAPH Mode in Ascend,
+              only supports whole graph sinking or sub graph sinking in pipeline parallel, and does not support
+              dynamic shape scenes. In addition, this mode incurs additional compilation costs and is difficult to
+              debug and tune.
 
         exc_mode (str, optional): Control the execution mode of the model.
             Supports ["auto", "sink", "no_sink"]. Default: ``"auto"`` .

mindspore/common/jit_trace.py

@@ -28,6 +28,7 @@ from mindspore._c_expression import TraceRecorder as tr
 from mindspore._c_expression import JitExecutor_
 from mindspore._c_expression import TensorPy as Tensor, CSRTensor, COOTensor
 from mindspore._c_expression import typing
+from mindspore.common.jit_config import JitConfig
 
 
 class TraceJitContext(JitContext):
@@ -123,19 +124,19 @@ def nested_run(obj, cell, *args):
     return file_names, linenos, res
 
 
-def _jit_trace():
+def _jit_trace(jit_config):
     """Return the wrapped function for trace mode jit."""
     def wrap_func(fn):
         if hasattr(fn, "construct"):
             if isinstance(fn, ms.nn.Cell):
                 # Bound the cell object to get the self arg.
-                return types.MethodType(_jit_trace()(fn.construct.__func__), fn)
+                return types.MethodType(_jit_trace(jit_config)(fn.construct.__func__), fn)
             if isinstance(fn, type) and issubclass(fn, ms.nn.Cell):
-                fn.construct = _jit_trace()(fn.construct)
+                fn.construct = _jit_trace(jit_config)(fn.construct)
             return fn
 
         if isinstance(fn, types.MethodType):
-            return types.MethodType(_jit_trace()(fn.__func__), fn.__self__)
+            return types.MethodType(_jit_trace(jit_config)(fn.__func__), fn.__self__)
 
         if not isinstance(fn, types.FunctionType):
             logger.warning(f"The fn should be function, method or cell instance/class, but got {fn}")
@@ -150,6 +151,10 @@ def _jit_trace():
             if jit_context():
                 return fn(*args, **kwargs)
             # Start trace process.
+            if jit_config:
+                jit_config_dict = jit_config.jit_config_dict
+            else:
+                jit_config_dict = JitConfig().jit_config_dict
             if kwargs:
                 bound_arguments = inspect.signature(fn).bind(*args, **kwargs)
                 bound_arguments.apply_defaults()
@@ -170,14 +175,16 @@ def _jit_trace():
             line_str = fn.__code__.co_filename + ":" + str(fn.__code__.co_firstlineno)
             generate_name = generate_name + '#[' + line_str + ']'
 
-            new_compile = _jit_trace_begin(generate_name, *jit_args)
+            new_compile = _jit_trace_begin(
+                generate_name, *jit_args, jit_config=jit_config_dict)
             if new_compile:
                 fn_res = fn(*args, **kwargs)
                 logger.debug(f'fn: {fn}, fn_res: {fn_res}, line: {line_str}')
                 # Use fn's output to build func graph's output.
-                output = _jit_trace_end(fn_res)
+                output = _jit_trace_end(fn_res, jit_config=jit_config_dict)
             else:
-                output = _jit_trace_end(None)  # Run with compilation.
+                # Run with compilation.
+                output = _jit_trace_end(None, jit_config=jit_config_dict)
             logger.debug(f'output: {output}')
             return output
 
@@ -224,7 +231,7 @@ def _get_args_for_run(args):
     return tuple(new_args)
 
 
-def _jit_trace_begin(fn_name, *args):
+def _jit_trace_begin(fn_name, *args, **kwargs):
     """
     Start to build a MindIR func graph for a code snippet by trace method.
 
@@ -257,6 +264,10 @@ def _jit_trace_begin(fn_name, *args):
         ...
         >>> out = tensor_add(x, y)
     """
+    if "jit_config" in kwargs:
+        jit_config = kwargs["jit_config"]
+    else:
+        jit_config = JitConfig().jit_config_dict
     global _using_trace
     if _using_trace:
         raise RuntimeError(
@@ -279,7 +290,7 @@ def _jit_trace_begin(fn_name, *args):
     if not _compile_only and phase in _trace_compile_cache:
         logger.debug('Had compiled, just run.')
         _trace_jit_context.compiled = True
-        output = tr.get_instance().run_graph(phase, args)
+        output = tr.get_instance().run_graph(phase, jit_config, args)
         from mindspore.common.api import _convert_python_data
         _trace_jit_context.result = _convert_python_data(output)
         logger.debug(f'jit trace result: {_trace_jit_context.result}')
@@ -295,7 +306,7 @@ def _jit_trace_begin(fn_name, *args):
     return True
 
 
-def _jit_trace_end(*output_args):
+def _jit_trace_end(*output_args, **kwargs):
     """
     Finish building a MindIR func graph for a code snippet by trace method.
 
@@ -330,19 +341,23 @@ def _jit_trace_end(*output_args):
         ...
         >>> out = tensor_add(x, y)
     """
+    if "jit_config" in kwargs:
+        jit_config = kwargs["jit_config"]
+    else:
+        jit_config = JitConfig().jit_config_dict
     if _trace_jit_context.compiled:
         output = _trace_jit_context.result
         logger.debug(f'jit trace result: {output}')
     else:
         logger.debug(f'output_args: {output_args}')
         file_names, linenos = _get_caller_lines()
-        tr.get_instance().end_graph(file_names, linenos, *output_args)
+        tr.get_instance().end_graph(file_names, linenos, jit_config, *output_args)
         if _compile_only:
             output = output_args[0] if len(output_args) == 1 else output_args
         else:
             args = _get_args_for_run(_trace_jit_context.args)
             output = tr.get_instance().run_graph(
-                _trace_jit_context.phase, args)
+                _trace_jit_context.phase, jit_config, args)
             from mindspore.common.api import _convert_python_data
             output = _convert_python_data(output)
             logger.debug(f'jit trace result: {output}')

mindspore/common/lazy_inline.py

@@ -32,9 +32,11 @@ def lazy_inline(fn=None, attrs=None, policy=None):
     static_graph_expert_programming.html#using-lazy-inline-decorator>`_ .
 
     .. warning::
-        This feature is only supported on Ascend and is not supported on other hardwares.
-        The construct parameters must be positional or key word arguments and have not default values.
-        The cell has not switch sub graph.
+        - This feature is only supported on Ascend and is not supported on other hardwares.
+        - The construct parameters must be positional or key word arguments and have not default values.
+        - The cell has not switch sub graph.
+        - In the gradient accumulation scenario, it is recommended to use the @lazy_inline decorator to
+          reduce compilation time, and this decorator is only allowed to configure on the outermost cell.
 
     Args:
         fn (function): `__init__` function of a cell.

mindspore/common/parameter.py

@@ -21,7 +21,6 @@ from copy import copy
 import time
 import os
 import sys
-import math
 import numbers
 import numpy as np
 
@@ -29,8 +28,6 @@ from mindspore import log as logger
 from mindspore.log import _LogActionOnce
 from mindspore._c_expression import ParamInfo
 from mindspore.common import dtype as mstype
-from mindspore import context
-from mindspore.common._utils import get_slice_num, get_slice_shape
 from mindspore.common.initializer import initializer
 from mindspore.common.tensor import Tensor, _TensorMeta
 from mindspore.common.hook_handle import _update_hook_version
@@ -39,10 +36,6 @@ from mindspore._check_jit_forbidden_api import jit_forbidden_register
 from mindspore._c_expression import TensorPy as Tensor_
 from mindspore.parallel._tensor import _get_slice_index
 from mindspore.parallel._auto_parallel_context import auto_parallel_context
-from mindspore.parallel._ps_context import _is_role_worker, _is_role_pserver, _is_role_sched, _clone_hash_table, \
-                                           _is_ps_mode
-from mindspore.parallel._ps_context import _reinsert_hash_table_size, _insert_accumu_init_info, _cache_enable
-from mindspore.common._decorator import deprecated
 from mindspore.communication._comm_helper import _is_initialized
 from mindspore.communication import get_group_size, get_rank
 import mindspore.common._monad as monad
@@ -138,11 +131,7 @@ def _offload_if_config(data):
     Args:
         data: The parameter data to offload.
     """
-    if not context.get_context("memory_offload") or data is None:
-        return
-
-    offload_context = context.get_offload_context()
-    if offload_context.get("offload_param", None) != "disk":
+    if data is None:
         return
 
     data_size_threshold = 512
@@ -219,7 +208,10 @@ class Parameter(Tensor_):
                 self.param_a = Parameter(Tensor([1], ms.float32), name="name_a")
                 self.param_tuple = (self.param_a, self.param_a)
 
-        requires_grad (bool): True if the parameter requires gradient. Default: ``True`` .
+        requires_grad (bool): It is Used to filter parameters in :func:`mindspore.nn.Cell.trainable_params()`.
+            If it is ``False``, the filter parameters will not be returned in
+            :func:`mindspore.nn.Cell.trainable_params()`.
+            Default: ``True`` .
         layerwise_parallel (bool): When `layerwise_parallel` is true in data/hybrid parallel mode,
             broadcast and gradients communication would not be applied to the `Parameter`. Default: ``False`` .
         parallel_optimizer (bool): It is used to filter the weight shard operation in parallel mode. It works only when
@@ -230,10 +222,8 @@ class Parameter(Tensor_):
         device(str): Only Ascend device target is supported. It is used to specify the device which the parameter is
             stored. By default, the parameter will be stored on NPU while computing. When the device is specified as
             ``"CPU"``, the parameter will be loaded into the device when it needs to be used, and unloaded to the CPU
-            after use. It takes effext only when `memory_offload` is ``"ON"``, `jit_level` is not ``"O2"`` and
-            `memory_optimize_level` is ``O0`` in :func:`mindspore.set_context`.
-            Less device memory is needed when device is
-            specified as ``"CPU"``.
+            after use. It takes effext only when `jit_level` is not ``"O2"`` and `memory_optimize_level` is ``O0``
+            in :func:`mindspore.set_context`. Less device memory is needed when device is specified as ``"CPU"``.
 
     Examples:
         >>> import numpy as np
@@ -272,8 +262,6 @@ class Parameter(Tensor_):
         obj.is_default_input_init = init_data_flag
         if obj.has_init:
             obj.init_mode = default_input
-        else:
-            _offload_if_config(obj)
         return obj
 
     def __reduce_ex__(self, _):
@@ -289,7 +277,6 @@ class Parameter(Tensor_):
     def __init__(self, default_input, name=None, requires_grad=True, layerwise_parallel=False, parallel_optimizer=True,
                  storage_format="", device=None):
         self.param_info = ParamInfo()
-        self.init_in_server = False
         self.name = name
         self.requires_grad = requires_grad
         self.layerwise_parallel = layerwise_parallel
@@ -300,32 +287,15 @@ class Parameter(Tensor_):
         self.is_init = False
         self._inited_param = None
         self._sliced = False
-        self.is_param_ps = False
-        self.push_weight_to_server = False
-        self.pull_weight_from_server = False
         self.requires_aggr = True
         self._cast_type = None
         self._unique = False
         self.is_in_parallel = _is_in_auto_parallel_mode()
         self._pipeline_stage_list = []
-        self.slice_num = 1
         if -1 in self.shape:
             raise ValueError(f"All shape elements of the Parameter must be positive. But got None.")
         if isinstance(default_input, (Tensor_, Tensor)):
-            # At embedding cache scenes, we need limit the size of memory for parameter.
-            # And save out range data to persistent storage to support TB-Level size parameter.
-            slice_num_of_persistent_data = get_slice_num(default_input.dtype, default_input.shape)
-            if slice_num_of_persistent_data > 1:
-                data_shape = list(default_input.shape)
-                slice_first_dim = math.ceil(data_shape[0] / slice_num_of_persistent_data)
-                data_shape[0] = slice_first_dim
-                self.param_info.use_persistent_storage = True
-                self.param_info.origin_shape = default_input.shape
-                self.slice_num = slice_num_of_persistent_data
-                Tensor_.__init__(self, dtype=default_input.dtype, shape=tuple(data_shape))
-            else:
-                Tensor_.__init__(self, dtype=default_input.dtype, shape=default_input.shape)
-
+            Tensor_.__init__(self, dtype=default_input.dtype, shape=default_input.shape)
         elif isinstance(default_input, int):
             Tensor_.__init__(self, dtype=mstype.int64, shape=())
         elif isinstance(default_input, float):
@@ -387,11 +357,10 @@ class Parameter(Tensor_):
                     return (Tensor, data.asnumpy(), mstype.qint4x2)
                 return (Tensor, data.asnumpy())
 
-            not_init_data = not init_param or _is_role_sched() or (_is_role_pserver() and _cache_enable()) \
-                            or _is_in_auto_parallel_mode() or _is_parallel_mode()
+            not_init_data = not init_param or _is_in_auto_parallel_mode() or _is_parallel_mode()
             if not_init_data:
                 # do not init data while in auto parallel.
-                return (Tensor, None, data.dtype, get_slice_shape(data.dtype, data.shape), data.init)
+                return (Tensor, None, data.dtype, data.shape, data.init)
             return (Tensor, data.init_data())
         if isinstance(data, int):
             return (Tensor, data, mstype.int32)
@@ -399,29 +368,6 @@ class Parameter(Tensor_):
             return (Tensor, data, mstype.float32)
         return (Tensor, data)
 
-    def set_param_ps(self, init_in_server=False):
-        """
-        Set whether the trainable parameter is updated by parameter server and whether the
-        trainable parameter is initialized on server.
-
-        Note:
-            It only works when a running task is in the parameter server mode.
-            It is supported only in graph mode.
-
-        Args:
-            init_in_server (bool): Whether trainable parameter updated by parameter server is
-                initialized on server. Default: ``False``.
-
-        """
-        if not _is_ps_mode() or not (_is_role_worker() or _is_role_pserver() or _is_role_sched()):
-            raise RuntimeError("Must complete following two steps before calling set_param_ps: \n"
-                               "1. context.set_ps_context(enable_ps=True) \n"
-                               "2. export MS_ROLE environment variable \n"
-                               "Please refer to the official website for detailed usage.")
-        self.is_param_ps = True
-        self.init_in_server = init_in_server
-        self.param_info.init_in_server = init_in_server
-
     def copy(self):
         """
         Copy the parameter.
@@ -437,16 +383,6 @@ class Parameter(Tensor_):
         """
         return self.clone(init='same')
 
-    @deprecated("1.8", "set_param_fl")
-    def set_param_fl(self, push_to_server=False, pull_from_server=False, requires_aggr=True):
-        if push_to_server:
-            self.push_weight_to_server = True
-        if pull_from_server:
-            self.pull_weight_from_server = True
-        if not requires_aggr:
-            self.requires_aggr = False
-            self.param_info.requires_aggr = False
-
     @property
     def inited_param(self):
         """
@@ -512,8 +448,6 @@ class Parameter(Tensor_):
             raise ValueError("The type of the Parameter's name should be 'string' or 'None', "
                              "but got {}.".format(type(name_)))
 
-        if _is_role_worker() and self.cache_enable:
-            _reinsert_hash_table_size(name_, self.param_info.name)
         self.param_info.name = name_
 
     @property
@@ -642,8 +576,6 @@ class Parameter(Tensor_):
         x.param_info = param_info_clone
         x.is_init = False
         x.init = self.init
-        x.is_param_ps = self.is_param_ps
-        x.init_in_server = self.init_in_server
         x.cache_enable = self.cache_enable
         if x.cache_enable:
             x.key = _get_unique_parameter_key()
@@ -651,7 +583,7 @@ class Parameter(Tensor_):
         if self.cache_shape:
             x.cache_shape = self.cache_shape
         if init != 'same':
-            shape = self.shape if self.slice_num == 1 else self.param_info.origin_shape
+            shape = self.shape
             dtype = self.dtype
             tensor = initializer(init, shape=shape, dtype=dtype)
             x.set_data(tensor)
@@ -796,6 +728,7 @@ class Parameter(Tensor_):
             raise TypeError("The argument `requires_grad` must be bool type")
         Tensor_.wait_pipeline(self)
         self.param_info.requires_grad = value
+        self._requires_grad = value
 
     @property
     def data(self):
@@ -862,20 +795,6 @@ class Parameter(Tensor_):
             raise TypeError("The original tensor data is initialized, but the argument 'data' is not initialized."
                             "Please initialize 'data' before call this method.")
 
-    @staticmethod
-    def _from_tensor(tensor, *args, **kwargs):
-        """Create a `Parameter` that data is shared from a `Tensor`."""
-        if not isinstance(tensor, Tensor_):
-            raise TypeError(f"The type of input must be Tensor, but got {type(tensor)}.")
-        param = Tensor_.__new__(Parameter)
-        Tensor_.__init__(param, tensor)
-        param.init = None
-        param.init_mode = None
-        param.has_init = False
-        param.is_default_input_init = False
-        Parameter.__init__(param, tensor, *args, **kwargs)
-        return param
-
     @jit_forbidden_register
     def set_data(self, data, slice_shape=False):
         """
@@ -981,16 +900,7 @@ class Parameter(Tensor_):
 
         init_data_args = self._get_init_data_args(layout)
 
-        if _is_role_sched():
-            return self
-        if self.init_in_server and self.is_param_ps and isinstance(self.init_mode, Tensor) and \
-                self.init_mode.init is not None and _is_role_worker():
-            if self.cache_enable:
-                data = self.init_mode.init_data(*init_data_args)
-            else:
-                data = self.init_mode.init_data(0, [1])
-        else:
-            data = self.init_mode.init_data(*init_data_args)
+        data = self.init_mode.init_data(*init_data_args)
         origin_dtype = self.dtype
         obj = self._update_tensor_data(data)
         if self.dtype != origin_dtype:
@@ -999,7 +909,6 @@ class Parameter(Tensor_):
             self._inited_param = obj
         obj.init_mode = None
         obj.sliced = set_sliced
-        _offload_if_config(obj)
         return obj
 
     def register_hook(self, hook_fn):
@@ -1154,9 +1063,6 @@ class ParameterTuple(tuple):
             if not x1.cache_enable:
                 continue
 
-            if _is_role_worker():
-                _clone_hash_table(x.name, x.key, x1.name, x1.key)
-                _insert_accumu_init_info(x1.name, init_to_value(init))
         return ParameterTuple(new)
 
     def __parameter_tuple__(self):

mindspore/common/recompute.py

@@ -22,11 +22,10 @@ from mindspore.common.tensor import Tensor
 from mindspore import ops
 from mindspore.ops.composite import GradOperation
 from mindspore.common._register_for_recompute import recompute_registry
-from mindspore.common.api import _pynative_executor, _no_grad
+from mindspore.common.api import _pynative_executor, _no_grad, _run_in_jit
 from mindspore.common.generator import get_rng_state, set_rng_state
 from mindspore.train.amp import AmpDecorator
 from mindspore._c_expression.amp import get_curr_amp_strategy
-from mindspore._check_jit_forbidden_api import jit_forbidden_register
 
 
 class _WrapCell(Cell):
@@ -211,22 +210,15 @@ def _detach_input(input_arg):
 def _check_validation(block):
     if not isinstance(block, Cell):
         raise TypeError("Recompute function now only support block which inherited from Cell!")
-    if block.construct.__code__.co_name == "staging_specialize":
-        logger.warning('Block\'s construct method decorated by @jit that recompute '
-                       'function will not come into effect.')
 
 
-@jit_forbidden_register
 def recompute(block, *args, **kwargs):
     r"""
     This function is used to reduce memory, when run block, rather than
     storing the intermediate activation computed in forward pass, we will recompute it in backward pass.
 
     Note:
-        - Recompute function only support block which inherited from Cell object.
-        - This function interface now only support pynative mode. you can use Cell.recompute interface
-          in graph mode.
-        - When use recompute function, block object should not decorated by @jit.
+        Recompute function only support block which inherited from Cell object.
 
     Args:
         block (Cell): Block to be recompute.
@@ -238,7 +230,6 @@ def recompute(block, *args, **kwargs):
 
     Raises:
         TypeError: If `block` is not Cell object.
-        AssertionError: If execute mode is not PYNATIVE_MODE.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -272,6 +263,8 @@ def recompute(block, *args, **kwargs):
     """
 
     _check_validation(block)
+    if _run_in_jit():  # @jit.cond: True
+        return ops.recompute_block(block)(*args, **kwargs)
     return _RecomputeCell(block)(*args, **kwargs)