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

mindspore/common/dynamic_shape/enable_dynamic.py

@@ -0,0 +1,197 @@
+# This is the Python adaptation and derivative work of Myia (https://github.com/mila-iqia/myia/).
+#
+# Copyright 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.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Define enable_dynamic decorator."""
+import types
+import inspect
+from mindspore import log as logger
+from mindspore.common.tensor import Tensor
+from mindspore.common._utils import get_func, is_dim_unknown
+from mindspore.common.dynamic_shape.auto_dynamic_shape import SHAPE_DIM_ANY
+
+
+ENABLE_DYNAMIC = "__enable_dynamic__"
+
+
+def _check_element_valid(item, shape, name):
+    """Check elements in shape."""
+    if item is not SHAPE_DIM_ANY and (isinstance(item, int) and item <= 0):
+        raise TypeError(f"The argument '{name}' has invalid shape '{shape}', only supports None " \
+                        f"or a tuple/list of positive integers and None.")
+    return True
+
+
+def _check_arg_shape_valid(arg, name):
+    """Check if the shape of arg is valid"""
+    #if the shape of arg is None
+    if isinstance(arg, Tensor) and is_dim_unknown(arg.shape):
+        return True
+    if isinstance(arg, Tensor) and \
+        SHAPE_DIM_ANY in arg.shape and \
+        all(_check_element_valid(item, arg.shape, name) for item in arg.shape):
+        return True
+    if isinstance(arg, (tuple, list)) and any(_check_arg_shape_valid(item, name) for item in arg):
+        return True
+    return False
+
+
+def _check_arg_type_valid(arg, name):
+    """Check if the type of arg is valid."""
+    if isinstance(arg, Tensor):
+        return
+    if isinstance(arg, (tuple, list)):
+        for item in arg:
+            _check_arg_type_valid(item, name)
+    else:
+        raise TypeError(f"The decorator enable_dynamic only supports Tensor " \
+                        f"or a tuple/list of Tensor, but the argument : {name} is type of:{type(arg)}.")
+
+
+def _check_input_valid(arg):
+    """Check if real argument is valid."""
+    if isinstance(arg, Tensor):
+        if not all(isinstance(item, int) and item > 0 for item in arg.shape):
+            raise ValueError(f"When using decorator enable_dynamic, the corresponding shape of inputs should be " \
+                             f"a tuple/list of positive integers")
+    elif isinstance(arg, (tuple, list)):
+        for item in arg:
+            _check_input_valid(item)
+    else:
+        raise TypeError(f"When using decorator enable_dynamic, the corresponding inputs only supports Tensor " \
+                    f"or a tuple/list of Tensor.")
+
+
+def _check_arg_type_shape(arg, dyn_arg, name):
+    """Check the type, shape and dtype of real argument."""
+    if isinstance(arg, Tensor) and isinstance(dyn_arg, Tensor):
+        if arg.dtype != dyn_arg.dtype:
+            raise TypeError(f"When using decorator enable_dynamic, input tensor dtype = {arg.dtype}, " \
+                            f"dynamic tensor dtype = {dyn_arg.dtype}, tensor dtypes are not the same.")
+        if is_dim_unknown(dyn_arg.shape):
+            return
+        if len(arg.shape) != len(dyn_arg.shape) or \
+            any(y is not SHAPE_DIM_ANY and x != y for x, y in zip(arg.shape, dyn_arg.shape)):
+            raise ValueError(f"When using decorator enable_dynamic, input tensor shape = {arg.shape}, " \
+                             f"dynamic tensor shape = {dyn_arg.shape}, tensor shapes are not the same.")
+    elif isinstance(arg, (tuple, list)) and isinstance(dyn_arg, (tuple, list)):
+        if len(arg) != len(dyn_arg):
+            raise ValueError("Input sequences must have the same structure and length.")
+        for x, y in zip(arg, dyn_arg):
+            _check_arg_type_shape(x, y, name)
+    else:
+        raise TypeError(f"When using decorator enable_dynamic, the type between argument '{name}' " \
+                        f"and corresponding input are not the same.")
+
+
+def generate_dynamic_sequence_args(args_list, dyn_args_list):
+    """Generate dynamic shapes for input sequence"""
+    if isinstance(args_list, Tensor):
+        return dyn_args_list if args_list.shape != dyn_args_list.shape else args_list
+    result = []
+    for x, y in zip(args_list, dyn_args_list):
+        result.append(generate_dynamic_sequence_args(x, y))
+    return type(args_list)(result)
+
+
+def generate_dynamic_tensor_args(args_list, dynamic_shapes):
+    """Generate compile args with dynamic_shapes"""
+    new_compile_args = list(args_list)
+    for index, arg in enumerate(args_list):
+        if isinstance(arg, (tuple, list)) and not hasattr(arg, "__ms_mutable__"):
+            raise ValueError(f"When using decorator enable_dynamic, the corresponding attribute of input should be " \
+                             f"mutable(tuple/list)")
+        if index not in dynamic_shapes:
+            continue
+        _check_input_valid(arg)
+        name, dyn_arg = dynamic_shapes[index]
+        _check_arg_type_shape(arg, dyn_arg, name)
+        new_compile_args[index] = generate_dynamic_sequence_args(arg, dyn_arg)
+    logger.debug(f"args_list: {args_list}, dynamic_shapes: {dynamic_shapes}, " \
+                 f"new_compile_args: {new_compile_args}")
+    return new_compile_args
+
+
+def enable_dynamic(**kwargs):
+    """
+    Use to specify whether the shape of the parameter is dynamic shape or dynamic rank.
+
+    Note:
+        - It needs to be used in conjunction with the JIT interface. Without using the JIT decorator,
+          the dynamic shape and dynamic rank functions will not be enabled.
+        - In the scenario where both set_context(mode=GRAPH_MODE) and nn.Cell are set simultaneously,
+          use enabled_dynamic to report an error.
+
+    Args:
+        kwargs (dict): The input types are Tensor, tuple[Tensor] and list[Tensor]. If one or
+            more dimensions in the shape of the parameter need to be specified as dynamic shapes,
+            the corresponding dimensions in the shape can be set to None. If the shape that needs
+            to generate specified parameters is dynamic rank, the shape can be set to None.
+
+    Returns:
+        Function, return a function that specifies the dynamic shape information of the parameter.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import numpy as np
+        >>> import mindspore as ms
+        >>> from mindspore import Tensor
+        >>> from mindspore import enable_dynamic
+        >>> from mindspore import jit
+        ...
+        >>> x = Tensor(np.random.randn(2, 3), ms.float32)
+        >>> y = Tensor(np.random.randn(2, 3), ms.float32)
+        ...
+        >>> # Specify parameter y as dynamic shape
+        >>> @enable_dynamic(y=Tensor(shape=None, dtype=ms.float32))
+        >>> @jit
+        >>> def func(x, y):
+        ...     return x + 1, y + 1
+        ...
+        >>> out = func(x, y)
+    """
+    # Check inputs at first.
+    if not kwargs:
+        raise ValueError(f"When using decorator enable_dynamic, the input cannot be empty!")
+    for name, arg in kwargs.items():
+        _check_arg_type_valid(arg, name)
+        if not _check_arg_shape_valid(arg, name):
+            raise TypeError(f"When using decorator enable_dynamic, the shape of argument '{name}' " \
+                            f"at least have one None.")
+
+    def decorator(func):
+        if not isinstance(func, (types.FunctionType, types.MethodType)):
+            raise ValueError(f"Decorator enable_dynamic can only be used for function or method " \
+                             f"decrocated by ms.jit, but got {func}.")
+        signature = inspect.signature(func)
+        sigs_name = [sig_name for sig_name in signature.parameters if sig_name != "self"]
+        if len(kwargs) > len(sigs_name):
+            raise ValueError(f"When using decorator enable_dynamic, the number of arguments {len(kwargs)} " \
+                             f"exceeds the number of function arguments {len(sigs_name)}.")
+        # Generate dynamic args.
+        dynamic_args = dict()
+        for key, value in kwargs.items():
+            index = sigs_name.index(key)
+            if index in dynamic_args:
+                raise ValueError(f"keyword argument repeated: {key}")
+            dynamic_args[index] = (key, value)
+        # Set dynamic_tensor_shape to func.
+        inner_func = inspect.unwrap(func, stop=lambda f: not hasattr(f, '__wrapped__'))
+        setattr(get_func(inner_func), ENABLE_DYNAMIC, dynamic_args)
+        logger.info(f"Set enable dynamic: {dynamic_args} to {inner_func}")
+        return func
+    return decorator

mindspore/common/file_system.py

@@ -14,10 +14,14 @@
 # ============================================================================
 """File system registration management"""
 from mindspore import log as logger
+from mindspore import _checkparam as Validator
+
+mindio_server_info = {"memfs.data_block_pool_capacity_in_gb": "100"}
 
 
 class FileSystem:
     """File operation interface manager"""
+
     def __init__(self):
         self.create = open
         self.create_args = ("ab",)
@@ -35,20 +39,33 @@ def _register_basic_file_system(fs: FileSystem):
     return True
 
 
-def _register_mindio_file_system(fs: FileSystem):
-    """register mindio file system"""
+def _init_mindio():
+    """Initialize MindIO and return the module if successful"""
     try:
-        import mindio
+        import mindio_acp as mindio
+        ret = mindio.initialize(server_info=mindio_server_info)
+        if ret == 0:
+            return mindio
+        logger.warning(f"Failed to initialize mindio_acp: ret = {ret}")
     except ImportError:
-        return False
+        pass
     try:
+        import mindio
         ret = mindio.initialize()
-    except AttributeError as e:
-        logger.warning(f"Failed to initialize MindIO: {e}")
-        return False
-    if ret != 0:
-        logger.warning(f"Failed to initialize MindIO: ret = {ret}")
+        if ret == 0:
+            return mindio
+        logger.warning(f"Failed to initialize mindio: ret = {ret}")
+    except ImportError:
+        pass
+    return None
+
+
+def _register_mindio_file_system(fs: FileSystem):
+    """register mindio file system"""
+    mindio = _init_mindio()
+    if mindio is None:
         return False
+
     fs.create = mindio.create_file
     fs.create_args = ()
     fs.open = mindio.open_file
@@ -56,3 +73,36 @@ def _register_mindio_file_system(fs: FileSystem):
     fs.backend = "mindio"
     logger.info("The weights are stored using MindIO as the backend.")
     return True
+
+
+def set_mindio_server_info(data_block_pool_capacity_in_gb=100):
+    """
+    Configure MindIO server settings.
+
+    Args:
+        data_block_pool_capacity_in_gb (int): Memory pool capacity for data blocks in gigabytes.
+    """
+    global mindio_server_info
+    Validator.check_positive_int(data_block_pool_capacity_in_gb, "data_block_pool_capacity_in_gb")
+    mindio_server_info["memfs.data_block_pool_capacity_in_gb"] = str(data_block_pool_capacity_in_gb)
+
+
+def mindio_preload(ckpt_file_name):
+    """
+    Preload data into memory using MindIO for faster access.
+
+    Args:
+        ckpt_file_name (str): Checkpoint file name.
+
+    Returns:
+        bool: True if preloading is successful, False otherwise.
+    """
+    Validator.check_value_type('ckpt_file_name', ckpt_file_name, str, "mindio_preload")
+    mindio = _init_mindio()
+    if mindio is None:
+        return False
+    if not hasattr(mindio, 'preload'):
+        logger.warning("MindIO module does not have preload method")
+        return False
+    mindio.preload(ckpt_file_name)
+    return True

mindspore/common/hook_handle.py

@@ -15,7 +15,23 @@
 """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
+
+
+# Global variable to mark the `Parameter` hook and `Cell` hook version
+_HOOK_VERSION = 0
+
+
+def _update_hook_version():
+    global _HOOK_VERSION
+    _HOOK_VERSION += 1
+
+
+def _hook_version():
+    global _HOOK_VERSION
+    return _HOOK_VERSION
 
 
 class _TensorHookHandle:
@@ -31,8 +47,9 @@ class _TensorHookHandle:
 
     def __init__(self, tensor):
         self.id = None
-        self.tensor_ref = weakref.ref(tensor)
+        self.tensor_weakref = weakref.ref(tensor)
 
+    @jit_forbidden_register
     def remove(self):
         """
         Remove the tensor hook function, which corresponds to this '_TensorHookHandle' object.
@@ -67,9 +84,9 @@ class _TensorHookHandle:
         """
         if self.id is not None:
             Tensor_.remove_hook(self.id)
-            tensor = self.tensor_ref()
+            tensor = self.tensor_weakref()
             if tensor is not None:
-                tensor._remove_hook() # pylint:disable=protected-access
+                tensor._remove_hook()  # pylint:disable=protected-access
 
 
 class HookHandle:
@@ -99,6 +116,7 @@ class HookHandle:
             if extra_dict is not None:
                 self.extra_dict_ref = weakref.ref(extra_dict)
 
+    @jit_forbidden_register
     def remove(self):
         """
         Remove the cell hook function, which corresponds to this 'HookHandle' object.
@@ -145,6 +163,8 @@ class HookHandle:
             (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], dtype=Float32,
             value= [ 2.00000000e+00]))
         """
+        _update_hook_version()  # pylint:disable=protected-access
+
         if self.hook_dict_ref is not None:
             hook_dict = self.hook_dict_ref()
             if hook_dict is not None and self.handle_id in hook_dict:
@@ -154,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/np_dtype.py

@@ -16,10 +16,10 @@
 # ============================================================================
 """Numpy data type for MindSpore."""
 
-from mindspore._c_expression.np_dtypes import np_version_valid
-if np_version_valid(True):
+from mindspore._c_expression.np_dtypes import np_dtype_valid
+if np_dtype_valid(True):
     from mindspore._c_expression.np_dtypes import bfloat16 # pylint: disable=unused-import
 
 __all__ = []
-if np_version_valid(False):
+if np_dtype_valid(False):
     __all__.extend(["bfloat16"])