mindspore 2.4.1__cp311-cp311-win_amd64.whl → 2.5.0__cp311-cp311-win_amd64.whl

mindspore/ops_generate/tensor_func_reg_cpp_generator.py

@@ -0,0 +1,422 @@
+# Copyright 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.
+# 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.
+# ============================================================================
+"""
+This module defines the PyboostInnerPrimGenerator class, which is responsible for generating Python primitive
+wrappers for Pyboost operations. The generator constructs Python function definitions based on operator prototypes,
+generates necessary import statements, and writes the generated content into Python source files.
+
+The primary functionality is to take operator prototypes, extract relevant fields, and create Python function wrappers
+that can be used to call the Pyboost primitive implementations.
+"""
+
+import os
+
+import template
+from template import Template
+import gen_constants as K
+import pyboost_utils
+import op_api_proto
+from gen_utils import save_file
+from base_generator import BaseGenerator
+from op_proto import OpProto
+from op_template_parser import OpTemplateParser
+
+
+class TensorFuncRegCppGenerator(BaseGenerator):
+    """
+    Generates C++ tensor function registration code for different backends (Ascend, CPU, GPU).
+
+    This class is responsible for generating header and implementation files required to register
+    tensor functions, including device-specific dispatchers and function definitions.
+    """
+
+    def __init__(self):
+        self.TENSOR_FUNC_CC_REG = template.TENSOR_FUNC_CC_REG
+        self.TENSOR_FUNC_CALL_BODY = template.TENSOR_FUNC_CALL_BODY
+        self.TENSOR_FUNC_OVERLOAD_CALL_BODY = template.TENSOR_FUNC_OVERLOAD_CALL_BODY
+        self.TENSOR_API_HEADER = template.TENSOR_API_HEADER
+        self.TENSOR_API_SOURCE = template.TENSOR_API_SOURCE
+        self.TENSOR_FUNC_UTILS = template.TENSOR_FUNC_UTILS
+        self.TENSOR_FUNC_UT_BODY = template.TENSOR_FUNC_UT_BODY
+        self.TENSOR_FUNC_UT_OVERLOAD_BODY = template.TENSOR_FUNC_UT_OVERLOAD_BODY
+        self.TENSOR_CPP_METHOD = template.TENSOR_CPP_METHOD
+
+        self.func_def_reg = Template("tensor_class->def(\"${func_name}\", TensorMethod${cpp_func_name});\n")
+        self.single_case_template = Template(
+            'case ${case_id}:\n'
+            '  ${device_dispatcher}\n'
+            '  break;\n'
+        )
+        self.single_case_in_ut_template = Template(
+            'case ${case_id}:\n'
+            '  ${device_dispatcher}\n'
+        )
+        self.device_dispatcher_template = Template(
+            'if (backend == kAscendDevice || backend == kDavinciDevice) {\n'
+            '  ${ascend_dispatcher}\n'
+            '} else if (backend == kCPUDevice) {\n'
+            '  ${cpu_dispatcher}\n'
+            '} else if (backend == kGPUDevice) {\n'
+            '  ${gpu_dispatcher}\n'
+            '} else {\n'
+            '  MS_LOG(ERROR) << "Device target is not supported!";\n'
+            '  return py::none();\n'
+            '}'
+        )
+        self.pyboost_return_template = Template(
+            '${arg_handler_processor}\n'
+            'MS_LOG(INFO) << "Call Tensor${class_name}";\n'
+            'return ToPython(TensorPyboostMethodRegister::'
+            'GetOp(tensor::TensorPyboostMethod::k${class_name}Reg)(arg_list));\n'
+        )
+        self.callback_python_template = Template(
+            'MS_LOG(INFO) << "Callback python method: ${py_method}";\n'
+            'py::function fn = python_adapter::GetPyFn(\"mindspore.ops.tensor_method\", \"${py_method}\");\n'
+            'py::object res = fn(self, *py_args, **py_kwargs);\n'
+            'return res;\n'
+        )
+        self.callback_python_in_ut_template = Template(
+            'MS_LOG(INFO) << "Callback python method in UT: ${py_method}";\n'
+            'fn = python_adapter::GetPyFn(\"mindspore.ops.tensor_method\", \"${py_method}\");\n'
+            'res = fn(self, *py_args, **py_kwargs);\n'
+            'break;\n'
+        )
+        self.header_func_header_template = Template(
+            "py::object TensorMethod${cpp_func_name}"
+            "(const py::object &self, const py::args &py_args, const py::kwargs &py_kwargs);\n"
+        )
+
+    def generate(self, work_path, op_protos, func_protos_data, alias_func_mapping):
+        """
+        Generates C++ header and source files for tensor function registrations.
+
+        Args:
+            work_path (str): The directory where the generated files will be saved.
+            op_protos (list): A list of tensor op prototypes.
+            func_protos_data (dict): Dictionary mapping function names to lists of TensorFuncProto objects.
+            alias_func_mapping (dict): A dictionary mapping function name to its alias function names.
+        """
+
+        all_op_func_data, single_op_func_data, overload_op_func_data = op_api_proto.categorize_func_data(
+            func_protos_data)
+
+        tensor_method_list = self._get_op_enum_name_list(op_protos)
+        func_call_body_list = []
+        self._create_single_op_source_files(single_op_func_data, func_call_body_list)
+        self._create_overload_op_source_files(overload_op_func_data, func_call_body_list)
+        merge_func_call_body = pyboost_utils.merge_strings_by_chunk_size(func_call_body_list)
+
+        for i, func_body_chunk_str in enumerate(merge_func_call_body):
+            tensor_api_source = self.TENSOR_API_SOURCE.replace(tenosr_func_call_body=func_body_chunk_str)
+            save_file(os.path.join(work_path, K.TENSOR_API_PATH), f"tensor_api_{i}.cc",
+                      tensor_api_source)
+
+        func_def_body_list, tensor_cpp_methods_list, tensor_api_declaration_list = self._get_sorted_func_def_body(
+            all_op_func_data, alias_func_mapping)
+        tensor_api_header = self.TENSOR_API_HEADER.replace(tensor_api_declaration_list=tensor_api_declaration_list)
+        save_file(os.path.join(work_path, K.TENSOR_API_PATH), f"tensor_api.h",
+                  tensor_api_header)
+        self._generate_func_name_for_stub_tensor(work_path, tensor_cpp_methods_list)
+        func_cc_reg = self.TENSOR_FUNC_CC_REG.replace(func_def_body=func_def_body_list)
+        tensor_methods = self.TENSOR_FUNC_UTILS.replace(tensor_methods=tensor_method_list)
+
+        save_file(os.path.join(work_path, K.TENSOR_FUNC_REGISTER_PATH), "tensor_func_utils.h", tensor_methods)
+        save_file(os.path.join(work_path, K.TENSOR_FUNC_REGISTER_PATH), "tensor_func_reg.cc", func_cc_reg)
+
+    def _get_op_enum_name_list(self, op_protos):
+        """
+        Extracts operation class names and returns them as a formatted list.
+
+        Args:
+            op_protos (list): A list of operation prototype objects, where each object has an `op_class`
+                              with a `name` attribute.
+
+        Returns:
+            str: A list of formatted strings, where each string is of the form 'k<name>,\n', where <name>
+                  is the class name from the `op_class` attribute.
+
+        """
+        tensor_method_list = ""
+        for op_proto in op_protos:
+            if op_proto.op_dispatch is None or not op_proto.op_dispatch.enable:
+                continue
+            class_name = op_proto.op_class.name
+            tensor_method_list += f"k{class_name}Reg,\n"
+        return tensor_method_list
+
+    def _generate_func_name_for_stub_tensor(self, work_path, tensor_cpp_methods_list):
+        """
+        Generates a Python file containing tensor C++ function methods list and saves it to the specified path.
+
+        This function takes a list of C++ tensor methods, formats them into a Python script as a string,
+        and writes this script to a file named `_tensor_cpp_method.py` under the provided working path.
+
+        Args:
+            work_path (str): The base directory where the generated file will be saved.
+            tensor_cpp_methods_list (list): A list of tensor C++ method definitions to be included in the Python file.
+        """
+        tensor_cpp_methods_str = self.TENSOR_CPP_METHOD.replace(
+            tensor_cpp_methods_list_str=str(tensor_cpp_methods_list))
+        save_file(os.path.join(work_path, K.ADD_TENSOR_DOCS_PY_PATH), "_tensor_cpp_method.py", tensor_cpp_methods_str)
+
+    def _get_sorted_func_def_body(self, all_op_func_data, alias_func_mapping):
+        """
+        Generate sorted function definitions and headers for operations.
+
+        This function processes a dictionary of operation function data and an alias mapping,
+        producing two lists: one containing function definition bodies and another containing
+        function header definitions.
+
+        Args:
+            all_op_func_data (dict): A dictionary where keys are function API names (str), and
+                values are lists of function prototypes.
+            alias_func_mapping (dict): A mapping of function names to a list of their alias names.
+
+        Returns:
+            tuple: A tuple containing two lists:
+                - func_def_body_list (list of str): A list of formatted function definition strings.
+                - tensor_cpp_methods_list (list of str): A list of formatted function header strings.
+        """
+        func_def_body_list = []
+        tensor_cpp_methods_list = []
+        tensor_api_declaration_list = ""
+        for func_api_name, func_protos in all_op_func_data.items():
+            cpp_func_name = pyboost_utils.format_func_api_name(func_api_name)
+            if len(func_protos) == 1:
+                func_proto = func_protos[0]
+                func_name = func_proto.func_name
+                func_def_body_list.append(self.func_def_reg.replace(func_name=func_name, cpp_func_name=cpp_func_name))
+                tensor_cpp_methods_list.append(func_name)
+                tensor_api_declaration_list += self.header_func_header_template.replace(cpp_func_name=cpp_func_name)
+                if func_name in alias_func_mapping:
+                    for alias_func_name in alias_func_mapping[func_name]:
+                        func_def_body_list.append(
+                            self.func_def_reg.replace(func_name=alias_func_name, cpp_func_name=cpp_func_name))
+                        tensor_cpp_methods_list.append(alias_func_name)
+            elif len(func_protos) > 1:
+                func_def_body_list.append(
+                    self.func_def_reg.replace(func_name=func_api_name, cpp_func_name=cpp_func_name))
+                tensor_cpp_methods_list.append(func_api_name)
+                tensor_api_declaration_list += self.header_func_header_template.replace(cpp_func_name=cpp_func_name)
+                if func_api_name in alias_func_mapping:
+                    for alias_func_name in alias_func_mapping[func_api_name]:
+                        func_def_body_list.append(self.func_def_reg.replace(func_name=alias_func_name,
+                                                                            cpp_func_name=cpp_func_name))
+                        tensor_cpp_methods_list.append(alias_func_name)
+        return func_def_body_list, tensor_cpp_methods_list, tensor_api_declaration_list
+
+    def _create_single_op_source_files(self, single_op_func_data, func_call_body_list):
+        """
+        Generates the list of call body strings for single operation functions.
+
+        Args:
+            single_op_func_data (dict): Dictionary of tensor function prototypes with only one definition.
+
+        Returns:
+            list: Updated str list for generating C++ function call bodies.
+        """
+        for func_api_name, func_proto in single_op_func_data.items():
+            func_name = func_proto.func_name
+            cpp_func_name = pyboost_utils.format_func_api_name(func_api_name)
+            device_dispatcher_str = self._get_device_dispatchers_str(func_proto)
+            signature_str = self._generate_single_signature_str(
+                func_proto.op_proto, func_proto.kw_only_args, func_proto.varargs
+            )
+            op_args = func_proto.op_proto.op_args
+            max_size = len(op_args)
+            self_index = self._get_input_tensor_index(func_proto)
+            ut_body = self.TENSOR_FUNC_UT_BODY.replace(py_method=func_proto.py_method)
+            tensor_func_single_call_body = self.TENSOR_FUNC_CALL_BODY.replace(cpp_func_name=cpp_func_name,
+                                                                              func_name=func_name,
+                                                                              device_dispatcher=device_dispatcher_str,
+                                                                              signatures=signature_str,
+                                                                              max_args=max_size,
+                                                                              self_index=self_index,
+                                                                              ut_body=ut_body)
+            func_call_body_list.append(tensor_func_single_call_body)
+
+    def _create_overload_op_source_files(self, overload_op_func_data, func_call_body_list):
+        """
+        Generates the list of call body strings for overloaded operation functions.
+
+        Args:
+            overload_op_func_data (dict): Dictionary of tensor function prototypes with overloaded definitions.
+
+        Returns:
+            list: Updated str list for generating C++ function call bodies.
+        """
+        for func_api_name, func_protos in overload_op_func_data.items():
+            tensor_func_overload_call_body = self._get_overload_func_call_str(func_api_name, func_protos)
+            func_call_body_list.append(tensor_func_overload_call_body)
+
+    def _get_overload_func_call_str(self, func_api_name, func_protos):
+        """
+        Generates C++ call body string for overloaded tensor functions.
+
+        Args:
+            func_api_name (str): Name of the function API.
+            func_protos (list): List of TensorFuncProto objects representing the function prototypes.
+
+        Returns:
+            str: Generated call body string for the overloaded functions.
+        """
+        signatures_str = self._generate_func_signatures_list_str(func_protos)
+        dispatch_cases = self._get_dispatch_cases(func_protos)
+        ut_dispatch_cases = self._get_ut_dispatch_cases(func_protos)
+        ut_overload_body = self.TENSOR_FUNC_UT_OVERLOAD_BODY.replace(ut_dispatch_cases=ut_dispatch_cases)
+
+        max_size = 0
+        self_index = 0
+        for tensor_proto in func_protos:
+            op_proto = tensor_proto.op_proto
+            op_args = op_proto.op_args
+            max_size = max(len(op_args), max_size)
+            self_index = self._get_input_tensor_index(tensor_proto)
+        cpp_func_name = pyboost_utils.format_func_api_name(func_api_name)
+        overload_func_call_str = self.TENSOR_FUNC_OVERLOAD_CALL_BODY.replace(cpp_func_name=cpp_func_name,
+                                                                             func_name=func_api_name,
+                                                                             signatures=signatures_str,
+                                                                             dispatch_cases=dispatch_cases,
+                                                                             max_args=max_size,
+                                                                             self_index=self_index,
+                                                                             ut_overload_body=ut_overload_body)
+        return overload_func_call_str
+
+    def _generate_func_signatures_list_str(self, func_protos) -> str:
+        """
+        Generates function signatures as a string from the given prototypes.
+
+        Args:
+            func_protos (list): List of TensorFuncProto objects representing the function prototypes.
+
+        Returns:
+            str: Generated function signatures string.
+        """
+        sig_str = ''
+        first_sig = True
+        for tensor_proto in func_protos:
+            op_proto = tensor_proto.op_proto
+            if not first_sig:
+                sig_str += ',\n'
+            first_sig = False
+            sig_str += self._generate_single_signature_str(op_proto, tensor_proto.kw_only_args, tensor_proto.varargs)
+        return sig_str
+
+    def _generate_single_signature_str(self, op_proto: OpProto, kw_only_args, varargs) -> str:
+        op_parser = OpTemplateParser(op_proto)
+        return op_parser.generate_signature_str(kw_only_args, varargs, is_tensor_api=True)
+
+    def _get_input_tensor_index(self, func_proto):
+        """
+        Get index of input.
+
+        Args:
+            func_proto (TensorFuncProto): Function prototype to generate dispatch strings for.
+
+        Returns:
+            int: Index of input.
+        """
+        op_name = func_proto.op_proto.op_class.name
+        op_args = func_proto.op_proto.op_args
+        if op_name in K.INPUT_NAME_MAP:
+            self_index = [i for i in range(len(op_args)) if op_args[i].arg_name == K.INPUT_NAME_MAP[op_name]]
+        else:
+            self_index = [i for i in range(len(op_args)) if op_args[i].arg_name in K.INPUT_ARGS_NAME]
+        if len(self_index) != 1:
+            raise ValueError(
+                f'There must be only one field named \'input\'. But got {len(self_index)} in {op_name}')
+        return self_index
+
+    def _get_dispatch_cases(self, func_protos):
+        """
+        Generates C++ switch-case statements for dispatching tensor function calls.
+
+        Args:
+            func_protos (list): List of TensorFuncProto objects representing the function prototypes.
+
+        Returns:
+            str: Generated switch-case dispatch statements.
+        """
+        dispatch_cases_str = ''
+        for idx, func_proto in enumerate(func_protos):
+            device_dispatcher_str = self._get_device_dispatchers_str(func_proto)
+            dispatch_cases_str += self.single_case_template.replace(case_id=idx,
+                                                                    device_dispatcher=device_dispatcher_str)
+        dispatch_cases_str += 'default:\n'
+        dispatch_cases_str += '  return py::none();'
+        return dispatch_cases_str
+
+    def _get_ut_dispatch_cases(self, func_protos):
+        """
+        Generates C++ switch-case statements for dispatching tensor function calls.
+
+        Args:
+            func_protos (list): List of TensorFuncProto objects representing the function prototypes.
+
+        Returns:
+            str: Generated switch-case dispatch statements.
+        """
+        dispatch_cases_str = ''
+        for idx, func_proto in enumerate(func_protos):
+            device_dispatcher_str = self.callback_python_in_ut_template.replace(py_method=func_proto.py_method)
+            dispatch_cases_str += self.single_case_in_ut_template.replace(case_id=idx,
+                                                                          device_dispatcher=device_dispatcher_str)
+        dispatch_cases_str += 'default:\n'
+        dispatch_cases_str += '  res = py::none();'
+        return dispatch_cases_str
+
+    def _get_device_dispatchers_str(self, func_proto):
+        """
+        Generates device-specific dispatch strings for the given function prototype.
+
+        Args:
+            func_proto (TensorFuncProto): Function prototype to generate dispatch strings for.
+
+        Returns:
+            str: Generated device-specific dispatch string.
+        """
+        ascend_dispatcher_str = self._get_single_device_dispatcher_str(func_proto, 'ascend')
+        cpu_dispatcher_str = self._get_single_device_dispatcher_str(func_proto, 'cpu')
+        gpu_dispatcher_str = self._get_single_device_dispatcher_str(func_proto, 'gpu')
+        device_dispatcher_str = self.device_dispatcher_template.replace(ascend_dispatcher=ascend_dispatcher_str,
+                                                                        cpu_dispatcher=cpu_dispatcher_str,
+                                                                        gpu_dispatcher=gpu_dispatcher_str)
+        return device_dispatcher_str
+
+    def _get_single_device_dispatcher_str(self, func_proto, device):
+        """
+        Generates the dispatch string for a specific device.
+
+        Args:
+            func_proto (TensorFuncProto): Function prototype to generate the dispatcher for.
+            device (str): Device type ('ascend', 'cpu', 'gpu').
+
+        Returns:
+            str: Generated device dispatcher string.
+        """
+        func_proto_device = getattr(func_proto, device)
+        if func_proto_device == 'pyboost':
+            arg_handler_processor_str = self._get_arg_handler_processor(func_proto.func_name, func_proto.op_proto)
+            return self.pyboost_return_template.replace(arg_handler_processor=arg_handler_processor_str,
+                                                        class_name=func_proto.op_proto.op_class.name)
+
+        if func_proto_device == 'py_method':
+            return self.callback_python_template.replace(py_method=func_proto.py_method)
+
+        raise TypeError("Only support pyboost or python_method.")
+
+    def _get_arg_handler_processor(self, func_name, op_proto):
+        op_parser = OpTemplateParser(op_proto)
+        return op_parser.get_arg_handler_processor(func_name, op_proto, is_tensor_api=True)

mindspore/parallel/__init__.py

@@ -19,12 +19,12 @@ from mindspore.parallel.algo_parameter_config import get_algo_parameters, reset_
     set_algo_parameters
 from mindspore.parallel.checkpoint_transform import rank_list_for_transform, transform_checkpoint_by_rank, \
     transform_checkpoints, merge_pipeline_strategys, sync_pipeline_shared_parameters, \
-    load_segmented_checkpoints
+    load_segmented_checkpoints, set_op_strategy_config
 from mindspore.parallel.parameter_broadcast import parameter_broadcast
 from mindspore.parallel.shard import shard, Layout
-from mindspore.parallel.transform_safetensors import safetensors_to_ckpt, ckpt_to_safetensors, unified_safetensors
+from mindspore.parallel.transform_safetensors import unified_safetensors
 
 __all__ = ["set_algo_parameters", "reset_algo_parameters", "get_algo_parameters", "rank_list_for_transform",
            "transform_checkpoint_by_rank", "transform_checkpoints", "merge_pipeline_strategys", "shard",
            "sync_pipeline_shared_parameters", "Layout", "parameter_broadcast", "load_segmented_checkpoints",
-           "safetensors_to_ckpt", "ckpt_to_safetensors", "unified_safetensors"]
+           "unified_safetensors", "set_op_strategy_config"]

mindspore/parallel/_auto_parallel_context.py

@@ -187,6 +187,44 @@ class _AutoParallelContext:
         self.check_context_handle()
         return self._context_handle.get_dump_local_norm()
 
+    def set_dump_local_norm_path(self, dump_local_norm_path):
+        """
+        Set dump local norm path for auto parallel.
+
+        Args:
+            dump_local_norm_path (str): User need to specify the path to save dump files
+                                        if he want to dump local norm.  Default: ''
+
+        Raises:
+            KeyError: When the value of dump_local_norm_path is not a str value.
+        """
+        self.check_context_handle()
+        self._context_handle.set_dump_local_norm_path(dump_local_norm_path)
+
+    def get_dump_local_norm_path(self):
+        """Get dump local norm path."""
+        self.check_context_handle()
+        return self._context_handle.get_dump_local_norm_path()
+
+    def set_dump_device_local_norm(self, dump_device_local_norm):
+        """
+        Set dump device local norm for auto parallel.
+
+        Args:
+            dump_device_local_norm (bool): User need to specify if he want to dump device local norm.  Default: False
+
+        Raises:
+            ValueError: If the dump_device_local_norm is not a bool value.
+        """
+        self.check_context_handle()
+        self._context_handle.set_dump_device_local_norm(dump_device_local_norm)
+
+    def get_dump_device_local_norm(self):
+        """Get dump device local norm."""
+        self.check_context_handle()
+        return self._context_handle.get_dump_device_local_norm()
+
+
     def set_fusion_threshold_mb(self, fusion_threshold=64, comm_type="allreduce"):
         """
         Set fusion threshold (MB) for auto parallel.
@@ -1070,13 +1108,6 @@ class _AutoParallelContext:
         self.check_context_handle()
         return self._context_handle.get_optimizer_weight_shard_size()
 
-    def set_ops_strategy_json_config(self, type, path, mode):
-        """
-         Set configuration of saving ops strategy in file .json.
-        """
-        self.check_context_handle()
-        self._context_handle.set_ops_strategy_json_config(type, path, mode)
-
     def set_optimizer_weight_shard_aggregated_save(self, optimizer_weight_shard_aggregated_save):
         """
         Set optimizer_weight_shard_aggregated_save.
@@ -1214,31 +1245,6 @@ 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.
-
-    Args:
-        type (str): The parameter for choosing save or load .json file.
-        path (str): Path to save or load parallel strategy json.
-        mode (str): The parameter for choosing save all or important operators.
-
-    Raises:
-        KeyError: When type is not 'SAVE' or 'LOAD'.
-        KeyError: When mode is not 'all' or 'principal'.
-    """
-    dir_path = os.path.dirname(path)
-    if dir_path and not os.path.exists(dir_path):
-        os.makedirs(dir_path, mode=0o700, exist_ok=True)
-    check_type = ["SAVE", "LOAD"]
-    check_mode = ["all", "principal"]
-    if type in check_type and mode in check_mode:
-        auto_parallel_context().set_ops_strategy_json_config(type, path, mode)
-    else:
-        raise KeyError("Type must be 'SAVE' or 'LOAD' and mode must be 'all' or 'principal'")
-
-
 _AUTO_PARALLEL_CONTEXT = None
 
 
@@ -1287,7 +1293,9 @@ _set_auto_parallel_context_func_map = {
     "enable_alltoall": auto_parallel_context().set_enable_alltoall,
     "strategy_ckpt_config": auto_parallel_context().set_strategy_ckpt_config,
     "comm_fusion": auto_parallel_context().set_comm_fusion,
-    "dump_local_norm": auto_parallel_context().set_dump_local_norm}
+    "dump_local_norm": auto_parallel_context().set_dump_local_norm,
+    "dump_local_norm_path": auto_parallel_context().set_dump_local_norm_path,
+    "dump_device_local_norm": auto_parallel_context().set_dump_device_local_norm}
 
 _get_auto_parallel_context_func_map = {
     "device_num": auto_parallel_context().get_device_num,
@@ -1320,7 +1328,9 @@ _get_auto_parallel_context_func_map = {
     "comm_fusion": auto_parallel_context().get_comm_fusion,
     "strategy_ckpt_config": auto_parallel_context().get_strategy_ckpt_config,
     "full_batch_is_set": auto_parallel_context().get_full_batch_is_set,
-    "dump_local_norm": auto_parallel_context().get_dump_local_norm}
+    "dump_local_norm": auto_parallel_context().get_dump_local_norm,
+    "dump_local_norm_path": auto_parallel_context().get_dump_local_norm_path,
+    "dump_device_local_norm": auto_parallel_context().get_dump_device_local_norm}
 
 
 @args_type_check(device_num=int, global_rank=int, gradients_mean=bool, gradient_fp32_sync=bool,

mindspore/parallel/_cell_wrapper.py

@@ -24,7 +24,8 @@ 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
-from mindspore.communication import create_group
+from mindspore.communication import create_group, destroy_group
+from mindspore.communication._comm_helper import _get_group_map
 from mindspore.train._utils import get_parameter_redundancy, remove_param_redundancy
 
 _ALLGATHER_CELL = None
@@ -131,6 +132,21 @@ def _restore_parallel_context(origin_parallel_mode, origin_dataset_strategy):
             context.set_auto_parallel_context(dataset_strategy=origin_dataset_strategy)
 
 
+def _get_group_name(group_map, group):
+    """get group name"""
+    group_name = str(group)
+    is_manual_communication_group = True
+    if group_map:
+        for name, rank_list in group_map.items():
+            if list(group) == rank_list:
+                group_name = name
+                is_manual_communication_group = False
+                break
+    if is_manual_communication_group:
+        create_group(str(group), list(group))
+    return group_name, is_manual_communication_group
+
+
 def _single_parameter_broadcast(net, layout, cur_rank=0, initial_rank=0):
     """
     Broadcast single parameter to other rank in data parallel dimension.
@@ -158,8 +174,9 @@ def _single_parameter_broadcast(net, layout, cur_rank=0, initial_rank=0):
         return
     net_param_dict = net.parameters_dict()
     _chang_parallel_context(origin_dataset_strategy)
+    group_map = _get_group_map()
     for group, params in param_redundancy_reversed.items():
-        create_group(str(group), list(group))
+        group_name, is_manual_communication_group = _get_group_name(group_map, group)
         allreduce_input = []
         for param in params:
             if param not in net_param_dict:
@@ -170,7 +187,9 @@ def _single_parameter_broadcast(net, layout, cur_rank=0, initial_rank=0):
             allreduce_input.append(real_param)
         if not allreduce_input:
             continue
-        communicator = SingleCommunicator(str(group))
+        communicator = SingleCommunicator(group_name)
         for real_param in allreduce_input:
             real_param.set_data(communicator(real_param), real_param.sliced)
+        if is_manual_communication_group:
+            destroy_group(group_name)
     _restore_parallel_context(origin_parallel_mode, origin_dataset_strategy)

mindspore/parallel/_parallel_serialization.py

@@ -115,11 +115,15 @@ def _check_strategy_file(strategy_filename):
                          f"be empty. Please check whether the 'strategy_filename' is correct.")
 
 
-def _load_protobuf_strategy(strategy_filename):
+def _load_protobuf_strategy(strategy_filename, strategy_set=None):
     """load strategy from protobuf file"""
     parallel_strategy_map = ms.train.node_strategy_pb2.ParallelStrategyMap()
     with open(strategy_filename, 'rb') as f:
         pb_content = f.read()
+        if strategy_set is not None:
+            if pb_content in strategy_set:
+                return {}
+            strategy_set.add(pb_content)
     try:
         parallel_strategy_map.ParseFromString(pb_content)
     except BaseException as e:
@@ -188,8 +192,11 @@ def _merge_protobuf_strategy(src_strategy_files, dst_strategy_file):
     """merge protobuf strategy"""
     dst_parallel_strategy_map = ms.train.node_strategy_pb2.ParallelStrategyMap()
     merged_stage = []
+    strategy_set = set()
     for src_strategy_file in src_strategy_files:
-        src_parallel_strategy_map = _load_protobuf_strategy(src_strategy_file)
+        src_parallel_strategy_map = _load_protobuf_strategy(src_strategy_file, strategy_set=strategy_set)
+        if not src_parallel_strategy_map:
+            continue
         strategy_items = src_parallel_strategy_map.parallel_strategy_item
         layout_items = src_parallel_strategy_map.parallel_layout_item
         if not strategy_items or not layout_items:
@@ -339,6 +346,9 @@ def _get_device_num_from_strategy(strategy_file=None):
         src_strategy = strategy_file
     strategy_list = _convert_to_list(src_strategy)
     device_mat = list(strategy_list.values())[0][0]
+    if not device_mat:
+        raise ValueError("The parallel strategy file only contains pipeline-parallelism, which is not supported for "
+                         "parallel strategy conversion now.")
     return np.prod(device_mat)
 
 
@@ -559,3 +569,4 @@ def _get_param_list_when_first_dim_sharded(device_arrangement, first_dim_sharded
     start = rank - offset
     param_total_list = list(range(start, start + range_size))
     return param_total_list
+            

mindspore/parallel/_utils.py

@@ -30,7 +30,7 @@ from mindspore.common.seed import get_seed
 from mindspore._c_expression import GraphExecutor_
 from mindspore.parallel._tensor import _load_tensor_by_layout, _load_tensor_shape_by_layout
 
-SUPPORTED_TUPLE_IN_TUPLE_STRATEGY = ["GroupedMatmul", "FusedInferAttentionScore", "Custom"]
+SUPPORTED_TUPLE_IN_TUPLE_STRATEGY = ["GroupedMatmul", "FusedInferAttentionScore", "Custom", "Index"]
 
 
 def _get_parallel_mode():
@@ -127,7 +127,6 @@ class ParallelParamInitProfCtx:
 
 def _slice_parameter(parameter, phase, layout):
     """Slice python parameter obj according to the layout."""
-    # graph_executor.updata_param_node_default_input(phase, {parameter.name: parameter})
     if getattr(parameter, "init_param", False):
         if layout is None:
             parameter.sliced = True
@@ -136,6 +135,9 @@ def _slice_parameter(parameter, phase, layout):
             rank = get_rank()
             new_tensor_shape = _load_tensor_shape_by_layout(parameter, layout, rank)
             parameter.shape = new_tensor_shape
+            if hasattr(parameter.init_mode, "shape") and parameter.init_mode.shape != parameter.shape:
+                parameter.init_mode.shape = new_tensor_shape
+            parameter.sliced = True
     else:
         graph_executor = GraphExecutor_.get_instance()
         new_param = parameter.init_data(layout, set_sliced=True)