mindspore 2.2.11__cp39-cp39-win_amd64.whl → 2.3.0__cp39-cp39-win_amd64.whl

mindspore/ops_generate/template.py

@@ -0,0 +1,239 @@
+# 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},
+  /*.is_view_ =*/${is_view},
+};
+""")

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,8 +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
-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"]
+           "transform_checkpoint_by_rank", "transform_checkpoints", "merge_pipeline_strategys", "shard",
+           "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.
@@ -65,6 +65,19 @@ class _ParallelOptimizerConfig:
     OPTIMIZER_WEIGHT_SHARD_SIZE = "optimizer_weight_shard_size"
 
 
+class _PipelineConfig:
+    """
+    The key of the Pipeline parallelism.
+    """
+    PIPELINE_INTERLEAVE = "pipeline_interleave"
+    PIPELINE_SCHEDULER = "pipeline_scheduler"
+
+
+class _PipelineScheduler:
+    PIPELINE_1F1B = "1f1b"
+    PIPELINE_GPIPE = "gpipe"
+
+
 class _AutoParallelContext:
     """
     _AutoParallelContext is the environment in which operations are executed
@@ -105,11 +118,11 @@ class _AutoParallelContext:
             device_num (int): The device number.
 
         Raises:
-            ValueError: If the device num is not in [1, 4096].
+            ValueError: If the device num is not a positive integer.
         """
         self.check_context_handle()
-        if device_num < 1 or device_num > 4096:
-            raise ValueError("The context configuration parameter 'device_num' must be in [1, 4096], "
+        if device_num < 1:
+            raise ValueError("The context configuration parameter 'device_num' must be a positive integer, "
                              "but got the value of device_num : {}.".format(device_num))
         from mindspore.communication._comm_helper import _HCCL_TEST_AVAILABLE
         self._context_handle.set_hccl_test_avaible(_HCCL_TEST_AVAILABLE)
@@ -229,6 +242,48 @@ class _AutoParallelContext:
         self.check_context_handle()
         return self._context_handle.get_pipeline_stage_split_num()
 
+    def set_auto_pipeline(self, auto_pipeline):
+        """Set the pipeline stage number to automatic"""
+        if not isinstance(auto_pipeline, bool):
+            raise TypeError("For 'set_auto_parallel_context', the argument 'auto_pipeline' "
+                            "must be bool, but got the type : {}.".format(type(auto_pipeline)))
+        self.check_context_handle()
+        self._context_handle.set_auto_pipeline(auto_pipeline)
+
+    def get_auto_pipeline(self):
+        """Get whether the pipeline stage number is automatic"""
+        self.check_context_handle()
+        return self._context_handle.get_auto_pipeline()
+
+    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()
+        return self._context_handle.get_pipeline_interleave()
+
+    def get_pipeline_scheduler(self):
+        """Get pipeline scheduler"""
+        self.check_context_handle()
+        return self._context_handle.get_pipeline_scheduler()
+
     def set_pipeline_segments(self, segments):
         """Set the segments of the pipeline"""
         if isinstance(segments, bool) or not isinstance(segments, int):
@@ -529,7 +584,7 @@ class _AutoParallelContext:
         self.check_context_handle()
         dir_path = os.path.dirname(strategy_ckpt_save_file)
         if dir_path and not os.path.exists(dir_path):
-            os.makedirs(dir_path)
+            os.makedirs(dir_path, exist_ok=True)
         self._context_handle.set_strategy_ckpt_save_file(strategy_ckpt_save_file)
 
     def get_strategy_ckpt_save_file(self):
@@ -777,16 +832,87 @@ 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()
         return self._context_handle.get_enable_fold_pipeline()
 
+    def set_pipeline_config(self, pipeline_config):
+        r"""
+        Set the configuration for pipeline parallelism. The configuration provides more detailed behavior control about
+        parallel training when pipeline parallelism is enabled.
+
+        Args:
+            pipeline_config (dict): The configuration for pipeline parallelism. It supports following keys:
+
+            - pipeline_interleave(bool): Setting true enable interleave scheduler for pipeline parallelism. This
+                                         scheduler requires more memory but less bubble.
+            - pipeline_scheduler(string): There are two choices, "1f1b" and "gpipe". default is "1f1b"
+
+              - 1f1b: It requires less memory and bubble ratio, for it run backward pass when corresponding forward pass
+                      finished.
+              - gpipe: It requires more memory and bubble ratio, for it run backward pass after all forward pass
+                       finished.
+
+        Raises:
+            TypeError: If the type of `pipeline_config` is not `dict`.
+            ValueError: If the key in `pipeline_config` not  in ["pipeline_interleave", "pipeline_scheduler"].
+            ValueError: If pipeline interleave is False, pipeline scheduler is not `1f1b`.
+        """
+        self.check_context_handle()
+
+        if not isinstance(pipeline_config, dict):
+            raise TypeError("For 'set_pipeline_config', the argument 'pipeine_config' "
+                            "must be dict, but got the type : {}.".format(type(pipeline_config)))
+
+        pp_interleave = _PipelineConfig.PIPELINE_INTERLEAVE
+        pp_scheduler = _PipelineConfig.PIPELINE_SCHEDULER
+
+        for config_name in pipeline_config:
+            unknown_config = []
+            if config_name not in [pp_interleave, pp_scheduler]:
+                unknown_config.append(config_name)
+
+            if unknown_config:
+                raise ValueError("Unknown config: {}".format(unknown_config))
+
+        Validator.check_bool(
+            pipeline_config[pp_interleave], pp_interleave, pp_interleave)
+        self._context_handle.set_pipeline_interleave(
+            pipeline_config[pp_interleave])
+
+        Validator.check_string(pipeline_config[pp_scheduler], [_PipelineScheduler.PIPELINE_1F1B,
+                                                               _PipelineScheduler.PIPELINE_GPIPE])
+        if not pipeline_config[pp_interleave] and pipeline_config[pp_scheduler] != _PipelineScheduler.PIPELINE_1F1B:
+            raise ValueError(f"When pipeline_interleave is False, {pp_scheduler} is not supported")
+
+        self._context_handle.set_pipeline_scheduler(pipeline_config[pp_scheduler])
+
     def get_enable_parallel_optimizer(self):
         """Get parallel optimizer flag."""
         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
@@ -1068,6 +1194,7 @@ class _AutoParallelContext:
             self.set_enable_all_gather_fusion(openstate)
             self.set_enable_reduce_scatter_fusion(openstate)
 
+
 def _set_ops_strategy_json_config(type="SAVE", path="", mode="all"):
     """
     Set strategy json configuration.
@@ -1091,6 +1218,7 @@ def _set_ops_strategy_json_config(type="SAVE", path="", mode="all"):
     else:
         raise KeyError("Type must be 'SAVE' or 'LOAD' and mode must be 'all' or 'principal'")
 
+
 _AUTO_PARALLEL_CONTEXT = None
 
 
@@ -1114,6 +1242,8 @@ _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,
+    "auto_pipeline": auto_parallel_context().set_auto_pipeline,
+    "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,
@@ -1125,7 +1255,9 @@ _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,
     "all_reduce_fusion_config": auto_parallel_context().set_all_reduce_fusion_split_indices,
     "communi_parallel_mode": auto_parallel_context().set_communi_parallel_mode,
@@ -1143,6 +1275,10 @@ _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,
+    "auto_pipeline": auto_parallel_context().get_auto_pipeline,
+    "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,
     "search_mode": auto_parallel_context().get_strategy_search_mode,
     "auto_parallel_search_mode": auto_parallel_context().get_auto_parallel_search_mode,
@@ -1152,6 +1288,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,
@@ -1171,7 +1308,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.
@@ -1219,11 +1356,19 @@ 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
+        auto_pipeline (bool): Set the pipeline stage number to automatic. Its value will be selected between 1 and the
+                        parameter `pipeline_stages`. This option requires the `parallel_mode` to be ``auto_parallel``
+                        and the `search_mode` to be ``recursive_programming``. Default: ``False`` .
+        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".
 
@@ -1265,6 +1410,7 @@ def _set_auto_parallel_context(**kwargs):
                         and `size`. Config is same as `allgather`.
 
 
+
     Raises:
         ValueError: If input key is not attribute in auto parallel context.
     """
@@ -1307,10 +1453,13 @@ 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
+    - auto_pipeline: False
+    - 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.