mindspore 2.2.14__cp38-cp38-manylinux1_x86_64.whl → 2.3.0rc2__cp38-cp38-manylinux1_x86_64.whl

mindspore/ops_generate/template.py

@@ -0,0 +1,238 @@
+# 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.
+# ============================================================================
+"""Template."""
+import re
+import os
+
+
+class CppTemplate:
+    """
+    template for generate c++ code
+    """
+    regular_str = r"(^[^\n\S]*)?\$([^\d\W]\w*|\{,?[^\d\W]\w*\,?})"
+    regular_match = re.compile(regular_str, re.MULTILINE)
+
+    def __init__(self, code_pattern):
+        self.code_pattern = code_pattern
+
+    @staticmethod
+    def load_from_file(file_path):
+        with open(file_path, "r") as f:
+            return CppTemplate(f.read())
+
+    def replace(self, **kwargs):
+        """
+        replace param.
+        :param kwargs:
+        :return:
+        """
+
+        def find(key: str):
+            if key in kwargs:
+                return kwargs[key]
+            raise TypeError(f"{key} should be in kwargs!")
+
+        def add_indent(indent, var):
+            return "".join([indent + line + "\n" for data in var for line in str(data).splitlines()]).rstrip()
+
+        def extract_variable(key):
+            start = ""
+            end = ""
+            if key[0] == "{":
+                key = key[1:-1]
+                if key[0] == ",":
+                    start = ","
+                    key = key[1:]
+                if key[-1] == ",":
+                    end = ", "
+                    key = key[:-1]
+            return find(key), start, end
+
+        def match_rule(match):
+            indent = match.group(1)
+            key = match.group(2)
+            var, start, end = extract_variable(key)
+            if indent is not None:
+                if not isinstance(var, list):
+                    return add_indent(indent, [var])
+                return add_indent(indent, var)
+            if isinstance(var, list):
+                code = ", ".join(str(x) for x in var)
+                if not var:
+                    return code
+                return start + code + end
+            return str(var)
+
+        return self.regular_match.sub(match_rule, self.code_pattern)
+
+
+NEW_LINE = "\n"
+WORK_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../../../")
+
+PYTHON_PRIM_TEMPLATE = CppTemplate("""
+
+class _Pyboost${class_name}Prim(${class_name}Prim_):
+    def __call__(self, ${input_args}):
+        ${process_func}
+        return _convert_stub(super().__call__(${input_args}))
+
+
+${func_impl_name}_impl = _Pyboost${class_name}Prim()
+""")
+
+IMPORT_PYBOOST_PRIM_HEADER = f"""
+from mindspore.common._stub_tensor import _convert_stub
+from mindspore.ops.auto_generate.gen_arg_handler import *
+"""
+
+IMPORT_PYBOOST_FUNC_HEADER = f"""
+from mindspore.common import dtype as mstype
+from mindspore.ops.auto_generate.pyboost_inner_prim import *
+
+"""
+
+REGISTER_DEFINE_TEMPLATE = CppTemplate(
+    """
+    (void)py::class_<${class_name}PrimAdapter, PrimitiveFunctionAdapter, std::shared_ptr<${class_name}PrimAdapter>>(
+      *m, "${class_name}Prim_")
+      .def(py::init<>())
+      .def("__call__", &${class_name}PrimAdapter::Call, "Call ${class_name} op.");
+    m->def(\"${pyboost_op_name}\", &mindspore::pynative::${pyboost_cfunc_name}, \"Encrypt the data.\");""")
+REGISTER_TEMPLATE = CppTemplate("void RegisterPyBoostFunction(py::module *m) {${register_func}\n}")
+
+REGISTER_PYBOOST_GRAD_DEFINE_TEMPLATE = CppTemplate(
+    "MS_REG_PYBOOST_GRAD_OP(${pyboost_op_name}, mindspore::runtime::${pyboost_cfunc_name});\n")
+REGISTER_PYBOOST_GRAD_TEMPLATE = CppTemplate("${register_func}")
+
+PYBOOST_FUNCTION_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH, './mindspore/ccsrc/pipeline/pynative/op_function/template/pyboost_function.tpl'))
+
+PYBOOST_HEADER_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH, './mindspore/ccsrc/pipeline/pynative/op_function/template/pyboost_function_header.tpl'))
+
+PYBOOST_GRAD_FUNCTION_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH, './mindspore/ccsrc/runtime/pynative/op_function/template/pyboost_grad_function.tpl'))
+
+PYBOOST_GRAD_HEADER_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH,
+                 './mindspore/ccsrc/runtime/pynative/op_function/template/pyboost_grad_function_header.tpl'))
+
+GEN_OPS_DEF_HEADER_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH, './mindspore/python/mindspore/ops_generate/gen_ops_def_header.tpl'))
+
+PYBOOST_BASE_OP_DEFINE_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH, './mindspore/ccsrc/kernel/pyboost/template/pyboost_op_header.tpl'))
+
+PYBOOST_OP_REGISTER_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH, './mindspore/ccsrc/kernel/pyboost/template/pyboost_op_register.tpl'))
+
+# Ascend op generate
+PYBOOST_ASCEND_OP_HEADER_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH,
+                 './mindspore/ccsrc/plugin/device/ascend/kernel/pyboost/template/pyboost_aclnn_header_template.tpl'))
+
+PYBOOST_ASCEND_OP_SOURCE_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH,
+                 './mindspore/ccsrc/plugin/device/ascend/kernel/pyboost/template/pyboost_aclnn_source_template.tpl'))
+
+PYBOOST_ASCEND_CALL_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH,
+                 './mindspore/ccsrc/plugin/device/ascend/kernel/pyboost/template/pyboost_ascend_call_template.tpl'))
+
+PYBOOST_ASCEND_VIEW_CALL_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH,
+                 './mindspore/ccsrc/kernel/pyboost/template/'
+                 'pyboost_view_template.tpl'))
+
+PYBOOST_ASCEND_CUSTOMIZE_CALL_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH,
+                 './mindspore/ccsrc/plugin/device/ascend/kernel/pyboost/template'
+                 '/pyboost_ascend_customize_call_template.tpl'))
+
+# GPU op generate
+PYBOOST_GPU_OP_HEADER_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH,
+                 './mindspore/ccsrc/plugin/device/gpu/kernel/pyboost/template/pyboost_gpu_header_template.tpl'))
+
+PYBOOST_GPU_OP_SOURCE_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH,
+                 './mindspore/ccsrc/plugin/device/gpu/kernel/pyboost/template/pyboost_gpu_source_template.tpl'))
+
+PYBOOST_GPU_CALL_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH,
+                 './mindspore/ccsrc/plugin/device/gpu/kernel/pyboost/template/pyboost_gpu_call_template.tpl'))
+
+PYBOOST_GPU_VIEW_CALL_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH,
+                 './mindspore/ccsrc/kernel/pyboost/template/pyboost_view_template.tpl'))
+
+PYBOOST_GPU_CUSTOMIZE_CALL_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH,
+                 './mindspore/ccsrc/plugin/device/gpu/kernel/pyboost/template'
+                 '/pyboost_gpu_customize_call_template.tpl'))
+
+# CPU op generate
+PYBOOST_CPU_OP_HEADER_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH,
+                 './mindspore/ccsrc/plugin/device/cpu/kernel/pyboost/template/pyboost_cpu_header_template.tpl'))
+
+PYBOOST_CPU_OP_SOURCE_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH,
+                 './mindspore/ccsrc/plugin/device/cpu/kernel/pyboost/template/pyboost_cpu_source_template.tpl'))
+
+PYBOOST_CPU_CALL_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH,
+                 './mindspore/ccsrc/plugin/device/cpu/kernel/pyboost/template/pyboost_cpu_call_template.tpl'))
+
+PYBOOST_CPU_VIEW_CALL_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH,
+                 './mindspore/ccsrc/kernel/pyboost/template/pyboost_view_template.tpl'))
+
+PYBOOST_CPU_CUSTOMIZE_CALL_TEMPLATE = CppTemplate.load_from_file(
+    os.path.join(WORK_PATH,
+                 './mindspore/ccsrc/plugin/device/cpu/kernel/pyboost/template'
+                 '/pyboost_cpu_customize_call_template.tpl'))
+
+PYBOOST_PY_FUNC_IMPORT_HEADEAR = CppTemplate(
+    """from mindspore._c_expression import ${class_name}Prim_\n"""
+)
+
+PYBOOST_PY_FUNC_TEMPLATE = CppTemplate("""
+def ${func_name}(${func_args}):
+    r\"\"\"
+    ${description}
+    \"\"\"
+    return ${func_impl_name}_impl(${input_args})\n\n""")
+
+OP_PROTO_TEMPLATE = CppTemplate("""
+${class_name}FuncImpl g${class_name}FuncImpl;
+OpDef g${class_name} = {
+  /*.name_=*/"${class_name}",
+  /*.args_=*/ {
+    ${input_args}
+  },
+  /* .returns_ = */ {
+    ${return_args} 
+  },
+  /*.signatures_ =*/ {
+    ${signatures}
+  },
+  /*.indexes_ =*/ {
+    ${indexes}
+  },
+  /*.func_impl_=*/g${class_name}FuncImpl,
+  /*.enable_dispatch_ =*/${enable_dispatch},
+};
+""")

mindspore/parallel/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Huawei Technologies Co., Ltd
+# Copyright 2023-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.
@@ -18,9 +18,11 @@ from __future__ import absolute_import
 from mindspore.parallel.algo_parameter_config import get_algo_parameters, reset_algo_parameters, \
     set_algo_parameters
 from mindspore.parallel.checkpoint_transform import rank_list_for_transform, transform_checkpoint_by_rank, \
-    transform_checkpoints, merge_pipeline_strategys, load_segmented_checkpoints
-from mindspore.parallel.shard import shard
+    transform_checkpoints, merge_pipeline_strategys, sync_pipeline_shared_parameters, \
+    load_segmented_checkpoints
+from mindspore.parallel.parameter_broadcast import parameter_broadcast
+from mindspore.parallel.shard import shard, Layout
 
 __all__ = ["set_algo_parameters", "reset_algo_parameters", "get_algo_parameters", "rank_list_for_transform",
            "transform_checkpoint_by_rank", "transform_checkpoints", "merge_pipeline_strategys", "shard",
-           "load_segmented_checkpoints"]
+           "sync_pipeline_shared_parameters", "Layout", "parameter_broadcast", "load_segmented_checkpoints"]

mindspore/parallel/_auto_parallel_context.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2022 Huawei Technologies Co., Ltd
+# Copyright 2020-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.
@@ -242,6 +242,25 @@ class _AutoParallelContext:
         self.check_context_handle()
         return self._context_handle.get_pipeline_stage_split_num()
 
+    def set_pipeline_result_broadcast(self, pipeline_result_broadcast):
+        """
+        Set the value of enabling pipeline result broadcast. Default: ``False``.
+
+        Args:
+            pipeline_result_broadcast (bool): Enable/disable broadcast the last stage result to all other stages.
+        """
+        self.check_context_handle()
+        if not isinstance(pipeline_result_broadcast, bool):
+            raise TypeError("For 'set_auto_parallel_context().set_pipeline_result_broadcast', the argument "
+                            "'pipeline_result_broadcast' must be bool, but got the type : {}."
+                            .format(type(pipeline_result_broadcast)))
+        self._context_handle.set_pipeline_result_broadcast(pipeline_result_broadcast)
+
+    def get_pipeline_result_broadcast(self):
+        """Get the value of enabling pipeline result broadcast"""
+        self.check_context_handle()
+        return self._context_handle.get_pipeline_result_broadcast()
+
     def get_pipeline_interleave(self):
         """Get pipeline interleave flag"""
         self.check_context_handle()
@@ -800,6 +819,20 @@ class _AutoParallelContext:
                             .format(type(enable_parallel_optimizer)))
         self._context_handle.set_enable_parallel_optimizer(enable_parallel_optimizer)
 
+    def set_force_fp32_communication(self, force_fp32_communication):
+        """
+        Set enable/disable force fp32 communication.
+
+        Args:
+            set_force_fp32_communication (bool): Enable/disable force fp32 communication.
+        """
+        self.check_context_handle()
+        if not isinstance(force_fp32_communication, bool):
+            raise TypeError("For 'set_auto_parallel_context', "
+                            "the argument 'force_fp32_communication' must be bool, but got the type : {}."
+                            .format(type(force_fp32_communication)))
+        self._context_handle.set_force_fp32_communication(force_fp32_communication)
+
     def get_enable_fold_pipeline(self):
         """Get parallel optimizer flag."""
         self.check_context_handle()
@@ -861,6 +894,12 @@ class _AutoParallelContext:
         self.check_context_handle()
         return self._context_handle.get_enable_parallel_optimizer()
 
+    def get_force_fp32_communication(self):
+        """Get force fp32 communication flag."""
+        self.check_context_handle()
+        return self._context_handle.get_force_fp32_communication()
+
+
     def set_parallel_optimizer_config(self, parallel_optimizer_config):
         r"""
         Set the configure for parallel optimizer. The configure provides more detailed behavior control about parallel
@@ -1190,6 +1229,7 @@ _set_auto_parallel_context_func_map = {
     "gradient_fp32_sync": auto_parallel_context().set_gradient_fp32_sync,
     "loss_repeated_mean": auto_parallel_context().set_loss_repeated_mean,
     "pipeline_stages": auto_parallel_context().set_pipeline_stages,
+    "pipeline_result_broadcast": auto_parallel_context().set_pipeline_result_broadcast,
     "pipeline_segments": auto_parallel_context().set_pipeline_segments,
     "parallel_mode": auto_parallel_context().set_parallel_mode,
     "search_mode": auto_parallel_context().set_strategy_search_mode,
@@ -1201,6 +1241,7 @@ _set_auto_parallel_context_func_map = {
     "full_batch": auto_parallel_context().set_full_batch,
     "dataset_strategy": auto_parallel_context().set_dataset_strategy,
     "enable_parallel_optimizer": auto_parallel_context().set_enable_parallel_optimizer,
+    "force_fp32_communication": auto_parallel_context().set_force_fp32_communication,
     "parallel_optimizer_config": auto_parallel_context().set_parallel_optimizer_config,
     "pipeline_config": auto_parallel_context().set_pipeline_config,
     "grad_accumulation_step": auto_parallel_context().set_grad_accumulation_step,
@@ -1220,6 +1261,7 @@ _get_auto_parallel_context_func_map = {
     "gradient_fp32_sync": auto_parallel_context().get_gradient_fp32_sync,
     "loss_repeated_mean": auto_parallel_context().get_loss_repeated_mean,
     "pipeline_stages": auto_parallel_context().get_pipeline_stages,
+    "pipeline_result_broadcast": auto_parallel_context().get_pipeline_result_broadcast,
     "pipeline_interleave": auto_parallel_context().get_pipeline_interleave,
     "pipeline_scheduler": auto_parallel_context().get_pipeline_scheduler,
     "parallel_mode": auto_parallel_context().get_parallel_mode,
@@ -1231,6 +1273,7 @@ _get_auto_parallel_context_func_map = {
     "full_batch": auto_parallel_context().get_full_batch,
     "dataset_strategy": auto_parallel_context().get_dataset_strategy,
     "enable_parallel_optimizer": auto_parallel_context().get_enable_parallel_optimizer,
+    "force_fp32_communication": auto_parallel_context().get_force_fp32_communication,
     "grad_accumulation_step": auto_parallel_context().get_grad_accumulation_step,
     "all_reduce_fusion_config": auto_parallel_context().get_all_reduce_fusion_split_indices,
     "communi_parallel_mode": auto_parallel_context().get_communi_parallel_mode,
@@ -1250,7 +1293,7 @@ _get_auto_parallel_context_func_map = {
                  grad_accumulation_step=int, all_reduce_fusion_config=list, group_ckpt_save_file=str,
                  communi_parallel_mode=str, optimizer_weight_shard_size=int, sharding_propagation=bool,
                  optimizer_weight_shard_aggregated_save=bool, enable_alltoall=bool, comm_fusion=dict,
-                 strategy_ckpt_config=dict)
+                 strategy_ckpt_config=dict, force_fp32_communication=bool)
 def _set_auto_parallel_context(**kwargs):
     """
     Set auto parallel context.
@@ -1298,11 +1341,16 @@ def _set_auto_parallel_context(**kwargs):
         full_batch (bool): Whether to load the whole batch on each device. Default: ``False``.
         dataset_strategy Union[str, tuple]: Dataset sharding strategy. Default: "data_parallel".
         enable_parallel_optimizer (bool): Enable using optimizer segmentation or not. Default: ``False``.
+        force_fp32_communication (bool): A switch that determines whether reduce operators (AllReduce, ReduceScatter)
+                        are forced to use the fp32 data type for communication during communication. True is the enable
+                        switch. Default: ``False`` .
         all_reduce_fusion_config (list): Set allreduce fusion strategy by parameters indices.
         pipeline_stages (int): Set the stage information for pipeline parallel. This indicates how
                         the devices are distributed alone the pipeline. The total devices will be divided into
                         'pipeline_stags' stages. This currently could only be used when
                         parallel mode semi_auto_parallel is enabled. Default: 0
+        pipeline_result_broadcast (bool): A switch that broadcast the last stage result to all other stage in pipeline
+                        parallel inference. Default: ``False`` .
         communi_parallel_mode (str): There are tree kinds of communication parallel modes, "all_group_parallel",
                      "same_server_group_parallel" and "no_group_parallel". Default: "all_group_parallel".
 
@@ -1386,10 +1434,12 @@ def _reset_auto_parallel_context():
     - strategy_ckpt_load_file: ""
     - strategy_ckpt_save_file: ""
     - enable_parallel_optimizer: False
+    - force_fp32_communication: False
     - search_mode: 'recursive_programming
     - auto_parallel_search_mode: 'recursive_programming
     - sharding_propagation: False
     - pipeline_stages: 0
+    - pipeline_result_broadcast: False
     - gradient_accumulation_shard: True
     - fusion_threshold: 64
     """

mindspore/parallel/_cell_wrapper.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Huawei Technologies Co., Ltd
+# 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.
@@ -17,6 +17,7 @@ from __future__ import absolute_import
 from __future__ import division
 
 from mindspore.nn.cell import Cell
+from mindspore.ops import operations as P
 from mindspore.ops.operations.comm_ops import AllGather
 from mindspore.communication import GlobalComm
 from mindspore.common import jit
@@ -29,16 +30,18 @@ class AllGatherCell(Cell):
     Allgather cell, used in model parallel scenario.
     To allgather the selected parameter slice from each device.
     """
-    def __init__(self, group):
+    def __init__(self, group, do_reshape, after_reshape_slice_shape):
         super(AllGatherCell, self).__init__(auto_prefix=False)
-
         self.allgather = AllGather(group)
+        self.do_reshape = do_reshape
+        self.after_reshape_slice_shape = tuple(after_reshape_slice_shape)
         self.add_flags(skip_auto_parallel_compile=True)
 
     @jit()
     def construct(self, x):
+        if self.do_reshape:
+            x = P.Reshape()(x, self.after_reshape_slice_shape)
         x = self.allgather(x)
-
         return x
 
 
@@ -51,29 +54,33 @@ class SaveOptShardCkptCell(Cell):
     Note:
         This could be optimized later with less communication consumption.
     """
-    def __init__(self, group):
+    def __init__(self, group, do_reshape, after_reshape_slice_shape):
         super(SaveOptShardCkptCell, self).__init__(auto_prefix=False)
         self.allgather1 = AllGather(group)
         self.allgather2 = AllGather()
+        self.do_reshape = do_reshape
+        self.after_reshape_slice_shape = tuple(after_reshape_slice_shape)
         self.add_flags(skip_auto_parallel_compile=True)
 
     def construct(self, x):
         x = self.allgather1(x)
+        if self.do_reshape:
+            x = P.Reshape()(x, self.after_reshape_slice_shape)
         x = self.allgather2(x)
 
         return x
 
 
-def get_allgather_cell(group, need_merge_twice=False):
+def get_allgather_cell(group, need_merge_twice=False, do_reshape=False, after_reshape_slice_shape=()):
     """Get AllGatherCell object."""
     global _ALLGATHER_CELL
     if need_merge_twice:
-        _ALLGATHER_CELL = SaveOptShardCkptCell(group)
+        _ALLGATHER_CELL = SaveOptShardCkptCell(group, do_reshape, after_reshape_slice_shape)
     else:
         if group:
-            _ALLGATHER_CELL = AllGatherCell(group)
+            _ALLGATHER_CELL = AllGatherCell(group, do_reshape, after_reshape_slice_shape)
         else:
-            _ALLGATHER_CELL = AllGatherCell(GlobalComm.WORLD_COMM_GROUP)
+            _ALLGATHER_CELL = AllGatherCell(GlobalComm.WORLD_COMM_GROUP, do_reshape, after_reshape_slice_shape)
     return _ALLGATHER_CELL
 
 

mindspore/parallel/_cost_model_context.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Huawei Technologies Co., Ltd
+# 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.