mindspore 2.2.14__cp39-cp39-win_amd64.whl → 2.4.0__cp39-cp39-win_amd64.whl

mindspore/common/parameter.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2022 Huawei Technologies Co., Ltd
+# Copyright 2020-2024 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.
@@ -41,6 +41,8 @@ from mindspore.parallel._ps_context import _is_role_worker, _is_role_pserver, _i
                                            _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
 import mindspore.common._monad as monad
 
 __all__ = ['Parameter', 'ParameterTuple']
@@ -52,11 +54,22 @@ PARAMETER_NAME_PREFIX_MAX_LEN = 1024
 _GLOBAL_PARAMETER_KEY = -1
 
 
-def _is_in_parallel_mode():
+def _is_in_auto_parallel_mode():
     """Get parallel mode."""
     return auto_parallel_context().get_parallel_mode() in ["semi_auto_parallel", "auto_parallel"]
 
 
+def _is_parallel_mode():
+    """ Whether is parallel mode """
+    if not _is_initialized() or context.get_context('mode') == context.PYNATIVE_MODE:
+        return False
+    if os.getenv("RUN_MODE") != "predict":
+        return False
+    if get_group_size() > 1 and _get_parallel_mode() == "stand_alone":
+        return True
+    return False
+
+
 def init_to_value(init):
     """
     Get value of initializer.
@@ -91,6 +104,15 @@ def _get_unique_parameter_key():
     return _GLOBAL_PARAMETER_KEY
 
 
+def _gen_offload_file_path(offload_dir):
+    offload_dir = os.path.relpath(offload_dir)
+    if not os.path.exists(offload_dir):
+        os.makedirs(offload_dir, mode=0o700, exist_ok=True)
+    offload_file_path = offload_dir + "/" + str(_get_global_rank()) + "_" + str(
+        _get_unique_parameter_key()) + "_" + str(time.time()) + ".data"
+    return offload_file_path
+
+
 def _offload_if_config(data):
     """
     Offload parameter(data size > 512) to file when enable memory offload and offload parameter to disk.
@@ -111,11 +133,7 @@ def _offload_if_config(data):
     offload_file_path = data.offload_file_path()
     if offload_file_path is None or offload_file_path == "":
         offload_dir = offload_context.get("offload_path", "./offload")
-        offload_dir = os.path.relpath(offload_dir)
-        if not os.path.exists(offload_dir):
-            os.makedirs(offload_dir)
-        offload_file_path = offload_dir + "/" + str(_get_global_rank()) + "_" + str(
-            _get_unique_parameter_key()) + "_" + str(time.time()) + ".data"
+        offload_file_path = _gen_offload_file_path(offload_dir)
     data.offload(offload_file_path)
 
 
@@ -145,7 +163,9 @@ class Parameter(Tensor_):
     Args:
         default_input (Union[Tensor, int, float, numpy.ndarray, list]): Parameter data,
             to initialize the parameter data.
-        name (str): Name of the parameter. Default: ``None`` .
+        name (str): Name of the parameter. Default: ``None`` . If two or more `Parameter`
+            objects with the same name exist in a network,
+            you will be prompted to set a unique name when defining them.
 
             1) If the parameter is not given a name, the default name is its variable name. For example, the name of
             param_a below is name_a, and the name of param_b is the variable name param_b.
@@ -186,6 +206,15 @@ class Parameter(Tensor_):
         parallel_optimizer (bool): It is used to filter the weight shard operation in `SEMI_AUTO_PARALLEL` or
             `AUTO_PARALLEL` mode. It works only when enable parallel optimizer in
             `mindspore.set_auto_parallel_context()`. Default: ``True`` .
+        storage_format (str): Only Ascend device target is supported. It is used to specify the format of the weight
+            loaded to the device. By default, the format is not changed. The optional values are ``"FRACTAL_NZ"`` ,
+            ``"NC1HWC0"`` , ``"FRACTAL_Z"`` , etc. Default: ``""`` .
+        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 `mindspore.set_context()`. Less device memory is needed when device is
+            specified as ``"CPU"``.
 
     Examples:
         >>> import numpy as np
@@ -221,6 +250,7 @@ class Parameter(Tensor_):
         # it's better to make the Initializer a kind of tensor.
         obj.init_mode = None
         obj.is_default_input_init = init_data_flag
+        obj.from_ckpt = False
         if obj.has_init:
             obj.init_mode = default_input
         else:
@@ -237,7 +267,8 @@ class Parameter(Tensor_):
         return (
             Parameter, (data, self.name, self.requires_grad, self.layerwise_parallel))
 
-    def __init__(self, default_input, name=None, requires_grad=True, layerwise_parallel=False, parallel_optimizer=True):
+    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
@@ -256,10 +287,11 @@ class Parameter(Tensor_):
         self.requires_aggr = True
         self._cast_type = None
         self._unique = False
-        self.is_in_parallel = _is_in_parallel_mode()
+        self.is_in_parallel = _is_in_auto_parallel_mode()
         self.is_in_shard = False
         self._pipeline_stage_list = []
         self.slice_num = 1
+        self.from_ckpt = False
         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)):
@@ -287,6 +319,11 @@ class Parameter(Tensor_):
             raise TypeError(f"The type of the argument 'default_input' must be in ['Tensor', 'int', 'float',"
                             f" 'numpy.ndarray', 'list']. But got type {type(default_input)}.")
         self.param_info.parameter_shape = self.shape
+        self.param_info.storage_format = storage_format
+        if device is not None:
+            if device != "CPU":
+                raise ValueError(f"Only 'CPU' is supported for device, but got ${device}.")
+            self._set_user_data("parameter_device", device)
 
         import mindspore.ops.operations.other_ops as other_ops
         self.load = other_ops.Load()
@@ -329,11 +366,12 @@ class Parameter(Tensor_):
                     # in other place, so we can make a Tensor without copy data.
                     return (Tensor, data)
                 # make a copy of Tensor to init the parameter.
-                if data.dtype == mstype.bfloat16:
-                    return (Tensor, data.float().asnumpy(), mstype.bfloat16)
+                if data.dtype == mstype.qint4x2:
+                    return (Tensor, data.asnumpy(), mstype.qint4x2)
                 return (Tensor, data.asnumpy())
 
-            not_init_data = _is_role_sched() or (_is_role_pserver() and _cache_enable()) or _is_in_parallel_mode()
+            not_init_data = _is_role_sched() or (_is_role_pserver() and _cache_enable()
+                                                 ) 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)
@@ -359,7 +397,7 @@ class Parameter(Tensor_):
 
         Tutorial Examples:
             - `Parameter Server Mode
-              <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/parameter_server_training.html>`_
+              <https://www.mindspore.cn/docs/en/master/model_train/parallel/parameter_server_training.html>`_
         """
         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"
@@ -607,6 +645,9 @@ class Parameter(Tensor_):
             shape = self.shape if self.slice_num == 1 else self.param_info.origin_shape
             dtype = self.dtype
             x.set_data(initializer(init, shape=shape, dtype=dtype))
+        device = self._get_user_data("parameter_device")
+        if device is not None:
+            x._set_user_data("parameter_device", device)
         return x
 
     @property
@@ -741,6 +782,7 @@ class Parameter(Tensor_):
     def requires_grad(self, value=True):
         if not isinstance(value, bool):
             raise TypeError("The argument `requires_grad` must be bool type")
+        Tensor_.wait_pipeline(self)
         self.param_info.requires_grad = value
 
     @property
@@ -783,6 +825,15 @@ class Parameter(Tensor_):
 
     @_LogActionOnce(logger=logger, key='add_pipeline_stage')
     def add_pipeline_stage(self, stage):
+        """
+        Add a pipeline stage to the parameter.
+
+        Args:
+            stage(int): The pipeline stage to be added.
+
+        Raise:
+            TypeError: If `stage` is not a positive number or not int type.
+        """
         logger.warning(f"This interface may be deleted in the future.")
         if not isinstance(stage, int) or stage < 0:
             raise TypeError("`stage` must be a positive number of int type")
@@ -794,9 +845,9 @@ class Parameter(Tensor_):
                         f"Use .set_dtype(xxx) to change the dtype.")
 
     @staticmethod
-    def _set_data_check_input_valid(current_shape, data_shape, current_tensor_is_init,
-                                    incoming_tensor_is_init, slice_shape=False, slice_num=1):
-        if incoming_tensor_is_init and not current_tensor_is_init:
+    def _set_data_check_input_valid(current_shape, data_shape, current_tensor_is_init, incoming_tensor_is_init,
+                                    from_ckpt, slice_shape=False, slice_num=1):
+        if not from_ckpt and incoming_tensor_is_init and not current_tensor_is_init:
             raise TypeError("The original tensor data is initialized, but the argument 'data' is not initialized."
                             "Please initialize 'data' before call this method.")
         if tuple(current_shape) != tuple(data_shape):
@@ -851,7 +902,7 @@ class Parameter(Tensor_):
         incoming_tensor_is_init = isinstance(data, Tensor) and not data.has_init
         current_tensor_is_init = isinstance(self, Tensor) and not self.has_init
         Parameter._set_data_check_input_valid(self.shape, data.shape, current_tensor_is_init, incoming_tensor_is_init,
-                                              slice_shape, self.slice_num)
+                                              self.from_ckpt, slice_shape, self.slice_num)
         if self.dtype != data.dtype:
             if mstype.implicit_conversion_seq.get(self.dtype) < mstype.implicit_conversion_seq.get(data.dtype):
                 self._raise_type_error(data.dtype)
@@ -923,7 +974,7 @@ class Parameter(Tensor_):
             >>> x = Parameter(Tensor(np.array([[1, 2], [3, 4]], dtype=np.float32)), name="param")
             >>> x.init_data()
         """
-        if self.is_default_input_init and self.is_in_parallel != _is_in_parallel_mode():
+        if self.is_default_input_init and self.is_in_parallel != _is_in_auto_parallel_mode():
             raise RuntimeError("Must set or change parallel mode before any initializer Tensor created.")
         if self.init_mode is None:
             return self
@@ -1007,8 +1058,9 @@ class ParameterTuple(tuple):
             Tuple, the new Parameter tuple.
 
         Tutorial Examples:
-            - `Cell and Parameter - Parameter Tuple
-              <https://mindspore.cn/tutorials/en/r2.2/advanced/modules/layer.html#parameter-tuple>`_
+            - `Tensor and Parameter - Parameter Tuple
+              <https://mindspore.cn/docs/en/master/model_train/model_building/tensor_and_parameter.html
+              #parameter-tuple>`_
         """
         Validator.check_str_by_regular(prefix)
         new = []

mindspore/common/recompute.py

@@ -0,0 +1,292 @@
+# Copyright 2023 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.
+# ============================================================================
+"""Defines other operators with functional form."""
+
+from collections import OrderedDict
+from types import MethodType
+from mindspore import log as logger
+from mindspore.nn.cell import Cell
+from mindspore import context
+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.generator import get_rng_state, set_rng_state
+from mindspore.train.amp import amp_decorator
+from mindspore._c_expression.amp import get_curr_amp_strategy
+
+
+class _WrapCell(Cell):
+    """
+    The warp cell is used by recompute cell,
+    which can set mixed precision to warp cell
+    """
+
+    def __init__(self, function):
+        super(_WrapCell, self).__init__(auto_prefix=False)
+        self.function = function
+
+    def construct(self, *args, **kwargs):
+        return self.function(*args, **kwargs)
+
+
+class _RecomputeCell(Cell):
+    """
+    Recompute cell, given the sub block, this cell will recompute the block, rather than
+    storing the intermediate activation computed in forward pass, we will recompute it in backward pass.
+    Note:
+     - RecomputeCell now only support pynative mode.
+     - When use recompute function, block object should not decorated by @jit.
+    """
+
+    def __init__(self, block):
+        """Initialize Recompute cell."""
+        super(_RecomputeCell, self).__init__(auto_prefix=False)
+        self.args = []
+        self.kwargs = []
+        self.wrap_cell = _WrapCell(block)
+        self.wrap_cell.set_inputs()
+
+        self.net = block
+        self.internal_params = []
+        self.save_rng_state = False
+        self.cpu_rng_state = None
+        self._add_attr("is_cell_recompute", "True")
+        self.grad = GradOperation(get_all=True, get_by_list=True, sens_param=True)
+        self.init_mixed_precision_type(block)
+        self.amp_strategy = None
+
+    def construct(self, *args, **kwargs):
+        _check_input_args_validate(self.net, args, kwargs)
+        self.args.append(args)
+        self.kwargs.append(kwargs)
+        self.save_rng_state = kwargs.pop("save_rng_state", True)
+        if self.save_rng_state:
+            self.cpu_rng_state = get_rng_state()
+        self.amp_strategy = get_curr_amp_strategy()
+        with _no_grad():
+            return self.net(*args, **kwargs)
+
+    def bprop(self, *args):
+        """
+        Custom grad method for recompute
+        :param args:
+        :return: input grad and weight grads
+        """
+        grad_input = args[-1]
+        input_args = self.args[-1]
+        kwargs = self.kwargs[-1]
+        self.args.pop()
+        self.kwargs.pop()
+        if kwargs:
+            input_args_for_check = list(input_args) + list(kwargs.values())
+        else:
+            input_args_for_check = list(input_args)
+        # To detach inputs to avoid erasing auto grad meta info of origin inputs.
+        input_args = _detach_input(input_args)
+        kwargs = _detach_input(kwargs)
+        kwargs['sens'] = grad_input
+        try:
+            pre_rng_state = get_rng_state()
+            set_rng_state(self.cpu_rng_state)
+            _pynative_executor.set_is_run_recompute(True)
+            if self.amp_strategy:
+                with amp_decorator(self.amp_strategy.get_amp_level(), self.amp_strategy.get_amp_dtype(),
+                                   self.amp_strategy.get_white_list(), self.amp_strategy.get_black_list()):
+                    grads = self.grad(self.net, self.internal_params)(*input_args, **kwargs)
+            else:
+                grads = self.grad(self.net, self.internal_params)(*input_args, **kwargs)
+            _pynative_executor.set_is_run_recompute(False)
+            set_rng_state(pre_rng_state)
+        except Exception as err:
+            _pynative_executor.clear_res()
+            raise err
+        weights = OrderedDict()
+        input_grads = list(grads[0])
+        _padding_input_grads(input_args_for_check, input_grads)
+        for i, param in enumerate(self.internal_params):
+            weights[param] = grads[1][i]
+        return tuple(input_grads), weights
+
+    def init_mixed_precision_type(self, block):
+        """
+        init mix precision
+        :param block:
+        :return:
+        """
+        if isinstance(block, Cell):
+            # To avoid sub cell same name
+            block.check_names_and_refresh_name()
+            self.internal_params = block.trainable_params()
+            return
+        if isinstance(block, MethodType) and isinstance(block.__self__, Cell):
+            # To avoid sub cell same name
+            block.__self__.check_names_and_refresh_name()
+            self.internal_params = block.__self__.trainable_params()
+            self.wrap_cell.mixed_precision_type = block.__self__.get_mixed_precision_type()
+            self.wrap_cell.set_mixed_precision_type(block.__self__.get_mixed_precision_type())
+            self.net = self.wrap_cell
+        else:
+            raise TypeError("For Recompute cell, it not support FunctionType function, "
+                            "only support Cell object or MethodType function!")
+
+
+def _check_input_args_validate(block, args, kwargs):
+    """
+    Check recompute input args validate
+    :param args:
+    :return:
+    """
+    if not (any([isinstance(arg, Tensor) for arg in args]) or \
+        any([isinstance(arg, Tensor) for arg in kwargs.values()])):
+        logger.warning("None of the inputs of function are tensors, which not need use recompute!")
+    for arg in args:
+        if isinstance(arg, (tuple, list)):
+            for data in arg:
+                if isinstance(data, Tensor):
+                    logger.info("For recompute block {}, tensor input in Tuple or list "
+                                "will not calculate grads!".format(block))
+                    break
+
+
+def _padding_input_grads(args, input_grads):
+    """
+    Padding input grads to same as input args
+    :param args:
+    :param input_grads:
+    :return:
+    """
+    for i, arg in enumerate(args):
+        if isinstance(arg, (list, tuple)):
+            if all([not isinstance(data, Tensor) for data in arg]):
+                input_grads.insert(i, None)
+            else:
+                # None is placeholder
+                grads = [None for data in arg]
+                input_grads.insert(i, grads)
+        elif not isinstance(arg, Tensor):
+            input_grads.insert(i, None)
+    if len(args) != len(input_grads):
+        raise ValueError("For recompute cell, the input grads size should be same as input args size: {}, "
+                         "but got {}".format(len(args), len(input_grads)))
+
+
+def _detach_input(input_arg):
+    """
+    Detach input
+    :param input_arg:
+    :return: detach output
+    """
+    if isinstance(input_arg, Tensor):
+        return ops.stop_gradient(input_arg)
+    if isinstance(input_arg, (list, tuple)):
+        detach_inputs = []
+        for arg in input_arg:
+            detach_inputs.append(_detach_input(arg))
+        return detach_inputs if isinstance(input_arg, list) else tuple(detach_inputs)
+    if isinstance(input_arg, dict):
+        detach_inputs = {}
+        for key, val in input_arg.items():
+            if isinstance(val, Tensor):
+                detach_inputs[key] = ops.stop_gradient(val)
+            else:
+                detach_inputs[key] = val
+        return detach_inputs
+    return input_arg
+
+
+def _check_validation(block):
+    if not isinstance(block, Cell):
+        raise TypeError("Recompute function now only support block which inherited from Cell!")
+    if context.get_context("mode") != context.PYNATIVE_MODE:
+        raise AssertionError("Recompute function now only support pynative mode, you can use "
+                             "Cell.recompute() in graph mode.")
+    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.')
+
+
+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.
+
+    Args:
+        block (Cell): Block to be recompute.
+        args(tuple): Inputs for block object to run forward pass.
+        kwargs(dict): Optional input for recompute function.
+
+    Returns:
+        Same as return type of block.
+
+    Raises:
+        TypeError: If `block` is not Cell object.
+        AssertionError: If execute mode is not PYNATIVE_MODE.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import numpy as np
+        >>> import mindspore.nn as nn
+        >>> from mindspore import ops
+        >>> from mindspore import Tensor, recompute
+        >>> class MyCell(nn.Cell):
+        ...     def __init__(self):
+        ...         super(MyCell, self).__init__(auto_prefix=False)
+        ...         self.conv = nn.Conv2d(2, 2, 2, has_bias=False, weight_init='ones')
+        ...         self.relu = ops.ReLU()
+        ...
+        ...     def construct(self, x):
+        ...         y = recompute(self.conv, x)
+        ...         return self.relu(y)
+        >>> inputs = Tensor(np.ones([2, 2, 2, 2]).astype(np.float32) * 2)
+        >>> my_net = MyCell()
+        >>> grad = ops.grad(my_net)(inputs)
+        >>> print(grad)
+        [[[[2. 4.]
+           [4. 8.]]
+          [[2. 4.]
+           [4. 8.]]]
+         [[[2. 4.]
+           [4. 8.]]
+          [[2. 4.]
+           [4. 8.]]]]
+    """
+
+    _check_validation(block)
+    return _RecomputeCell(block)(*args, **kwargs)
+
+
+def recompute_generator(block):
+    """
+    generator of recompute object.
+    :param block:
+    :return:
+    """
+    return _RecomputeCell(block)
+
+
+recompute_registry.register(recompute_generator)
+
+__all__ = ['recompute']
+__all__.sort()

mindspore/common/seed.py

@@ -41,18 +41,18 @@ def set_seed(seed):
     Set global seed.
 
     Note:
-        The global seed is used by numpy.random, mindspore.common.Initializer and
-        mindspore.nn.probability.distribution.
+        - The global seed is used by numpy.random, mindspore.common.Initializer and
+          mindspore.nn.probability.distribution.
 
-        If global seed is not set, these packages will use their own default seed independently, numpy.random and
-        mindspore.common.Initializer will choose a random seed, mindspore.nn.probability.distribution will use zero.
+        - If global seed is not set, these packages will use their own default seed independently, numpy.random and
+          mindspore.common.Initializer will choose a random seed, mindspore.nn.probability.distribution will use zero.
 
-        Seed set by numpy.random.seed() only used by numpy.random, while seed set by this API will also used by
-        numpy.random, so just set all seed by this API is recommended.
+        - Seed set by numpy.random.seed() only used by numpy.random, while seed set by this API will also used by
+          numpy.random, so just set all seed by this API is recommended.
 
-        In semi_auto_parallel/auto_parallel mode, when using set_seed, weights with same shape and same sharding
-        strategy in the same device would be initialized to the same result, otherwise, they would be initialized to
-        the different result.
+        - In semi_auto_parallel/auto_parallel mode, when using set_seed, weights with same shape and same sharding
+          strategy in the same device would be initialized to the same result, otherwise, they would be initialized to
+          the different result.
 
     Args:
         seed (int): The seed to be set.