mindspore 2.6.0rc1__cp39-cp39-win_amd64.whl → 2.7.0rc1__cp39-cp39-win_amd64.whl

mindspore/ops_generate/pyboost/pyboost_overload_functions_cpp_generator.py

@@ -71,7 +71,7 @@ class PyboostOverloadFunctionsGenerator(BaseGenerator):
             '${arg_handler_processor}\n'
             'MS_LOG(INFO) << "Call Tensor${class_name}";\n'
             'auto res = ${pyboost_base_func_name}_OP(${prim_name}, parse_args.src_types_, ${convert_args});\n'
-            'trace::Capture(parse_args.arg_list_, "${class_name}", &res);\n'
+            'trace::Capture(parse_args.arg_list_, mindspore::prim::kPrim${class_name}, &res);\n'
             'return res;\n'
         )
         self.callback_python_template = Template(

mindspore/ops_generate/pyboost/pyboost_utils.py

@@ -68,7 +68,7 @@ def get_index(index: int):
     return "kIndex" + str(index)
 
 
-def get_convert_type_str(dtype: str, optional):
+def get_convert_type_str(dtype: str, optional, use_basic_type=False):
     """
     Convert type
     """
@@ -106,16 +106,30 @@ def get_convert_type_str(dtype: str, optional):
         'list[bool]': 'ToBoolListOptional<py::list>',
         'list[tensor]': 'ToTensorListOptional<py::list>',
     }
+    basic_optional_type_convert = {
+        'tuple[int]': "ToBasicIntVectorOptional",
+        'list[int]': "ToBasicIntVectorOptional",
+        'int': "ToBasicIntOptional",
+    }
+    basic_type_convert = {
+        'tuple[int]': "ToBasicIntVector",
+        'list[int]': "ToBasicIntVector",
+        'int': "ToBasicInt",
+    }
     if optional:
+        if use_basic_type and dtype in basic_optional_type_convert:
+            return basic_optional_type_convert[dtype]
         if dtype in optional_type_convert:
             return optional_type_convert[dtype]
         raise TypeError(f"""Unsupported convert optional type {dtype} for args.""")
+    if use_basic_type and dtype in basic_type_convert:
+        return basic_type_convert[dtype]
     if dtype in native_type_convert:
         return native_type_convert[dtype]
     raise TypeError(f"""Unsupported convert type {dtype} for args.""")
 
 
-def get_input_args_type_str(dtype: str, optional):
+def get_input_args_type_str(dtype: str, optional, use_basic_type=False):
     """
     Convert type
     """
@@ -153,17 +167,52 @@ def get_input_args_type_str(dtype: str, optional):
         'list[bool]': 'std::optional<ValueTuplePtr>',
         'list[tensor]': 'std::optional<ValueTuplePtr>',
     }
+    basic_optional_type_convert = {
+        'tuple[int]': "std::optional<std::vector<int64_t>>",
+        'list[int]': "std::optional<std::vector<int64_t>>",
+        'int': "std::optional<int64_t>",
+    }
+    basic_type_convert = {
+        'tuple[int]': "std::vector<int64_t>",
+        'list[int]': "std::vector<int64_t>",
+        'int': "int64_t",
+    }
     if optional:
+        if use_basic_type and dtype in basic_optional_type_convert:
+            return basic_optional_type_convert[dtype]
         if dtype in optional_type:
             return optional_type[dtype]
         raise TypeError(f"""Unknown optional type {dtype} for args.""")
+    if use_basic_type and dtype in basic_type_convert:
+        return basic_type_convert[dtype]
     if dtype in native_type:
         return native_type[dtype]
     raise TypeError(f"""Unknown type {dtype} for args.""")
 
 
+def basic_type_convert_str(dtype: str, optional):
+    """
+    Convert type
+    """
+    optional_type = {
+        'tuple[int]': "ToBasicIntVectorOptional",
+        'list[int]': "ToBasicIntVectorOptional",
+        'int': "ToBasicIntOptional",
+    }
+    native_type = {
+        'tuple[int]': "ToBasicIntVector",
+        'list[int]': "ToBasicIntVector",
+        'int': "ToBasicInt",
+    }
+    if optional:
+        if dtype in optional_type:
+            return optional_type[dtype]
+    if dtype in native_type:
+        return native_type[dtype]
+    return ""
+
 
-def get_value_convert_type_str(dtype: str, optional):
+def get_value_convert_type_str(dtype: str, optional, use_basic_type=False):
     """
     Convert type
     """
@@ -193,10 +242,24 @@ def get_value_convert_type_str(dtype: str, optional):
         'tuple[bool]': 'ToValueTupleOptional',
         'tuple[tensor]': 'ToValueTupleOptional',
     }
+    basic_optional_type_convert = {
+        'tuple[int]': "ToBasicIntVectorOptional",
+        'list[int]': "ToBasicIntVectorOptional",
+        'int': "ToBasicIntOptional",
+    }
+    basic_type_convert = {
+        'tuple[int]': "ToBasicIntVector",
+        'list[int]': "ToBasicIntVector",
+        'int': "ToBasicInt",
+    }
     if optional:
+        if use_basic_type and dtype in basic_optional_type_convert:
+            return basic_optional_type_convert[dtype]
         if dtype in optional_type_convert:
             return optional_type_convert[dtype]
         raise TypeError(f"""Unsupported convert optional type {dtype} for args.""")
+    if use_basic_type and dtype in basic_type_convert:
+        return basic_type_convert[dtype]
     if dtype in native_type_convert:
         return native_type_convert[dtype]
     raise TypeError(f"""Unsupported convert type {dtype} for args.""")
@@ -212,7 +275,7 @@ def tuple_input_to_cpp_type(dtype: str):
         'tuple[int]': 'int64_t',
         'tuple[float]': 'float',
         'tuple[bool]': 'bool',
-        'tuple[str]': 'string',
+        'tuple[str]': 'std::string',
         'tuple[tensor]': 'mindspore::tensor::TensorPtr',
         'list[int]': 'int64_t',
         'list[float]': 'float',
@@ -227,12 +290,44 @@ def number_input_to_cpp_type(dtype: str):
         'int': 'int64_t',
         'float': 'float',
         'bool': 'bool',
-        'str': 'string'
+        'str': 'std::string'
     }
     return types_map.get(dtype)
 
 
-def get_input_dtype(dtype: str, optional):
+def input_dtype_to_cpp_type(dtype: str, optional):
+    """
+    Map input dtype to cpp dtype
+    """
+    type_convert = {
+        'int': 'int64_t',
+        'float': 'float',
+        'bool': 'bool',
+        'number': 'mindspore::ScalarPtr',
+        'str': 'std::string',
+        'tensor': 'mindspore::tensor::TensorPtr',
+        'tuple[int]': 'std::vector<int64_t>',
+        'tuple[float]': 'std::vector<float>',
+        'tuple[bool]': 'std::vector<bool>',
+        'tuple[tensor]': 'std::vector<mindspore::tensor::TensorPtr>',
+        'list[int]': 'std::vector<int64_t>',
+        'list[float]': 'std::vector<float>',
+        'list[bool]': 'std::vector<bool>',
+        'list[tensor]': 'std::vector<mindspore::tensor::TensorPtr>',
+    }
+    optional_tensor_type_convert = {
+        'tensor': 'std::optional<mindspore::tensor::TensorPtr>',
+        'tuple[tensor]': 'std::vector<mindspore::tensor::TensorPtr>',
+        'list[tensor]': 'std::vector<mindspore::tensor::TensorPtr>'
+    }
+    if optional and dtype in optional_tensor_type_convert:
+        return optional_tensor_type_convert[dtype]
+    if dtype in type_convert:
+        return type_convert[dtype]
+    raise TypeError(f"""Unsupported dtype {dtype} for args.""")
+
+
+def get_input_dtype(dtype: str, optional, use_basic_type=False):
     """
     Convert type
     """
@@ -244,7 +339,7 @@ def get_input_dtype(dtype: str, optional):
         'bool': 'mindspore::BoolImmPtr',
         'number': 'mindspore::ScalarPtr',
         'str': 'mindspore::StringImmPtr',
-        'tensor': 'mindspore::tensor::BaseTensorPtr',
+        'tensor': 'mindspore::tensor::TensorPtr',
         'tuple[int]': value_tuple,
         'tuple[float]': value_tuple,
         'tuple[bool]': value_tuple,
@@ -261,16 +356,30 @@ def get_input_dtype(dtype: str, optional):
         'bool': 'std::optional<mindspore::BoolImmPtr>',
         'number': 'std::optional<mindspore::ScalarPtr>',
         'str': 'std::optional<mindspore::StringImmPtr>',
-        'tensor': 'std::optional<mindspore::tensor::BaseTensorPtr>',
+        'tensor': 'std::optional<mindspore::tensor::TensorPtr>',
         'tuple[int]': value_tuple_optional,
         'tuple[float]': value_tuple_optional,
         'tuple[bool]': value_tuple_optional,
         'tuple[tensor]': value_tuple_optional,
     }
+    basic_optional_type_convert = {
+        'tuple[int]': "std::optional<std::vector<int64_t>>",
+        'list[int]': "std::optional<std::vector<int64_t>>",
+        'int': "std::optional<int64_t>",
+    }
+    basic_type_convert = {
+        'tuple[int]': "std::vector<int64_t>",
+        'list[int]': "std::vector<int64_t>",
+        'int': "int64_t",
+    }
     if optional:
+        if use_basic_type and dtype in basic_optional_type_convert:
+            return basic_optional_type_convert[dtype]
         if dtype in optional_type_convert:
             return optional_type_convert[dtype]
         raise TypeError(f"""Unsupported convert optional type {dtype} for args.""")
+    if use_basic_type and dtype in basic_type_convert:
+        return basic_type_convert[dtype]
     if dtype in type_convert:
         return type_convert[dtype]
     raise TypeError(f"""Unsupported convert type {dtype} for args.""")
@@ -343,7 +452,7 @@ def get_tuple_input_convert(arg_name, arg_type):
     """
     cpp_type = tuple_input_to_cpp_type(arg_type)
     if cpp_type == "mindspore::tensor::TensorPtr":
-        cpp_type = "mindspore::tensor::BaseTensorPtr"
+        cpp_type = "mindspore::tensor::TensorPtr"
     return f"std::vector<{cpp_type}> {arg_name}_vector = ConvertValueTupleToVector<{cpp_type}>({arg_name});\n"
 
 

mindspore/ops_generate/tensor_py_cc_generator.py

@@ -37,15 +37,6 @@ class TensorPyCppGenerator(BaseGenerator):
             '{"${snake_api_name}"'
             ', (PyCFunction)TensorMethod${pascal_api_name}_CPyWrapper, METH_VARARGS | METH_KEYWORDS},'
         )
-        self.stubtensor_api_def_template = Template(
-            'py::cpp_function TensorMethod${snake_api_name}_wrapper(\n'
-            '  [](const py::object& self, const py::args& args, const py::kwargs& kwargs) {\n'
-            '      return  TensorMethod${pascal_api_name}(self, args, kwargs);\n'
-            '  },\n'
-            '  py::is_method(stubTensorClass)\n'
-            ');\n'
-            'stubTensorClass.attr("${snake_api_name}") = TensorMethod${snake_api_name}_wrapper;'
-        )
 
     def generate(self, work_path, tensor_method_protos, alias_func_mapping):
         """
@@ -61,7 +52,6 @@ class TensorPyCppGenerator(BaseGenerator):
         """
         wrapper_defs = []
         tensor_api_defs = []
-        stubtensor_api_defs = []
         for api_name, _ in tensor_method_protos.items():
             pascal_api_name = pyboost_utils.format_func_api_name(api_name)
             snake_api_name = api_name
@@ -72,12 +62,6 @@ class TensorPyCppGenerator(BaseGenerator):
                     pascal_api_name=pascal_api_name
                 )
             )
-            stubtensor_api_defs.append(
-                self.stubtensor_api_def_template.replace(
-                    snake_api_name=snake_api_name,
-                    pascal_api_name=pascal_api_name
-                )
-            )
             if api_name in alias_func_mapping:
                 alias_api_names = alias_func_mapping[api_name]
                 for alias_api_name in alias_api_names:
@@ -88,19 +72,12 @@ class TensorPyCppGenerator(BaseGenerator):
                             pascal_api_name=pascal_api_name
                         )
                     )
-                    stubtensor_api_defs.append(
-                        self.stubtensor_api_def_template.replace(
-                            snake_api_name=snake_api_name,
-                            pascal_api_name=pascal_api_name
-                        )
-                    )
 
         # delete the ' \' for the last wrapper macro definition
         wrapper_defs[-1] = wrapper_defs[-1][:-2]
 
         file_str = self.TENSOR_PY_CC_TEMPLATE.replace(
-            tensor_api_defs=tensor_api_defs,
-            stubtensor_api_defs=stubtensor_api_defs
+            tensor_api_defs=tensor_api_defs
         )
         save_file(
             os.path.join(work_path, K.TENSOR_PY_CC_PATH),

mindspore/parallel/_auto_parallel_context.py

@@ -999,8 +999,10 @@ class _AutoParallelContext:
                                                                _PipelineScheduler.PIPELINE_SEQPIPE,
                                                                _PipelineScheduler.PIPELINE_SEQVPP,
                                                                _PipelineScheduler.PIPELINE_SEQSMARTVPP])
-        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")
+        scheduler_val = pipeline_config[pp_scheduler]
+        if not pipeline_config[pp_interleave] and scheduler_val != _PipelineScheduler.PIPELINE_1F1B:
+            raise TypeError(f"When pipeline_interleave is False, {scheduler_val!r} is not supported, "
+                            "only '1f1b' is allowed.")
 
         self._context_handle.set_pipeline_scheduler(pipeline_config[pp_scheduler])
 
@@ -1038,10 +1040,10 @@ class _AutoParallelContext:
                                                  shape[n] \* size(dtype). Non-negative. Unit: KB. Default: 64.
             - optimizer_weight_shard_size(int): Set the optimizer weight shard group size if you want to specific the
                                                 maximum group size across devices when the parallel optimizer is
-                                                enabled. The numerical range can be (0, device_num]. Default value
-                                                is -1, which means the optimizer weight shard group size will
-                                                the data parallel group of each parameter. Default -1.
-
+                                                enabled. The numerical range can be (0, device_num] or -1. If pipeline
+                                                parallelism is enabled, the numerical range is (0, device_num/stage]
+                                                or -1. Default value is -1, which means the optimizer weight shard
+                                                group size will be equal to the data parallel group of each parameter.
         """
         self.check_context_handle()
         grad_shard_name = _ParallelOptimizerConfig.GRADIENT_ACCUMULATION_SHARD
@@ -1076,8 +1078,9 @@ class _AutoParallelContext:
 
         if optimizer_weight_shard_size_name in parallel_optimizer_config:
             value = parallel_optimizer_config[optimizer_weight_shard_size_name]
-            Validator.check_positive_int(value)
-            self.set_optimizer_weight_shard_size(value)
+            if value != -1:
+                Validator.check_positive_int(value, prim_name="optimizer_weight_shard_size")
+                self.set_optimizer_weight_shard_size(value)
 
         if optimizer_level_name in parallel_optimizer_config:
             optimizer_level = parallel_optimizer_config[optimizer_level_name]

mindspore/parallel/_cell_wrapper.py

@@ -26,7 +26,7 @@ from mindspore.ops.operations.comm_ops import AllGather
 from mindspore.communication import GlobalComm, get_rank
 from mindspore.common import jit
 from mindspore.communication import create_group, destroy_group, get_group_size
-from mindspore.communication._comm_helper import _get_group_map
+from mindspore.communication._comm_helper import _get_group_map, _remove_group_info
 from mindspore.train._utils import get_parameter_redundancy, remove_param_redundancy
 from mindspore.parallel.shard import Layout
 
@@ -119,20 +119,18 @@ def destroy_allgather_cell():
 
 def _chang_parallel_context(origin_dataset_strategy):
     """Change the original parallel state."""
-    if context.get_context("mode") == context.GRAPH_MODE:
-        context.set_auto_parallel_context(parallel_mode="hybrid_parallel")
-        if origin_dataset_strategy != "data_parallel":
-            context.set_auto_parallel_context(dataset_strategy="data_parallel")
+    context.set_auto_parallel_context(parallel_mode="hybrid_parallel")
+    if origin_dataset_strategy != "data_parallel":
+        context.set_auto_parallel_context(dataset_strategy="data_parallel")
 
 
 def _restore_parallel_context(origin_parallel_mode, origin_dataset_strategy):
     """Restore the original parallel state."""
-    if context.get_context("mode") == context.GRAPH_MODE:
-        context.set_auto_parallel_context(parallel_mode=origin_parallel_mode)
-        if origin_dataset_strategy != "data_parallel":
-            if origin_dataset_strategy is not None and isinstance(origin_dataset_strategy, list):
-                origin_dataset_strategy = tuple(tuple(ds_item) for ds_item in origin_dataset_strategy)
-            context.set_auto_parallel_context(dataset_strategy=origin_dataset_strategy)
+    context.set_auto_parallel_context(parallel_mode=origin_parallel_mode)
+    if origin_dataset_strategy != "data_parallel":
+        if origin_dataset_strategy is not None and isinstance(origin_dataset_strategy, list):
+            origin_dataset_strategy = tuple(tuple(ds_item) for ds_item in origin_dataset_strategy)
+        context.set_auto_parallel_context(dataset_strategy=origin_dataset_strategy)
 
 
 def _get_group_name(group_map, group):
@@ -166,11 +164,92 @@ def _remove_param_not_load(param_name, param_not_load):
         param_not_load.remove(param_name)
 
 
-def _single_parameter_broadcast(net, layout, param_not_load=None):
+def _get_param_index_in_group(total_param_loaded, group, param):
+    """Get param_index in group."""
+    param_rank_index = []
+    for rank_id in group:
+        if rank_id < len(total_param_loaded):
+            if param in total_param_loaded[rank_id]:
+                param_rank_index.append(rank_id)
+        else:
+            raise ValueError("rank_id should be smaller than total rank num")
+    return param_rank_index
+
+
+def _communicate_allreduce(allreduce_input, group_map, group):
+    """Communicate allreduce input."""
+    if not allreduce_input:
+        return
+    from mindspore import Tensor
+    group_name, is_manual_communication_group = _get_group_name(group_map, group)
+    if is_manual_communication_group:
+        create_group(group_name, list(group))
+    communicator = SingleCommunicator(group_name)
+    for real_param in allreduce_input:
+        real_param.set_data(communicator(Tensor(real_param)), real_param.sliced)
+    if is_manual_communication_group:
+        destroy_group(group_name)
+        _remove_group_info(group_name)
+
+
+def _create_allreduce_input(params, group, net_param_dict, total_param_loaded, param_not_load, cur_rank):
+    """Creates allreduce input."""
+    from mindspore import Tensor
+    allreduce_input = []
+    for param in params:
+        if param not in net_param_dict:
+            continue
+        if param.startswith("accu_grads") or param.endswith("expert_load"):
+            continue
+        param_rank_index = _get_param_index_in_group(total_param_loaded, group, param)
+        if not param_rank_index:
+            continue
+        elif len(param_rank_index) == 1:
+            real_param = net_param_dict[param]
+            _remove_param_not_load(real_param.name, param_not_load)
+            if cur_rank != param_rank_index[0]:
+                real_param.set_data(Tensor(np.zeros(real_param.shape), dtype=real_param.dtype), real_param.sliced)
+            allreduce_input.append(real_param)
+        elif len(param_rank_index) > 1:
+            raise ValueError(f"For param {param} in group {group} should be in one rank, but in {param_rank_index}.")
+    return allreduce_input
+
+
+def _get_sorted_group_map():
+    """Get the world group map."""
+    group_map = _get_group_map()
+    if group_map:
+        group_map = {key: group_map[key] for key in sorted(group_map.keys())}
+    return group_map
+
+
+def _check_total_param_loaded(total_param_loaded):
+    """Check total_param_loaded."""
+    flag = True
+    for rank_id, param_loaded in enumerate(total_param_loaded):
+        if rank_id not in param_loaded:
+            flag = False
+            logger.warning("The order of loaded parameters on each card obtained by all_gather_object is incorrect,"
+                           "and the parameter broadcast will reorder them.")
+            break
+    if not flag:
+        new_total_param_loaded = [None] * len(total_param_loaded)
+        for _, param_loaded in enumerate(total_param_loaded):
+            for param in param_loaded:
+                if isinstance(param, int):
+                    new_total_param_loaded[param] = param_loaded
+                    break
+        return new_total_param_loaded
+    return total_param_loaded
+
+
+def _single_parameter_broadcast(net, layout, param_not_load=None, param_loaded=None):
     """
     Broadcast single parameter to other rank in data parallel dimension.
     """
-    from mindspore import Tensor
+    logger.info("Start loading the parameter broadcast for removing redundant parameters.")
+    from mindspore.runtime import synchronize
+    from mindspore.mint.distributed import all_gather_object
     origin_parallel_mode = context.get_auto_parallel_context("parallel_mode")
     origin_dataset_strategy = context.get_auto_parallel_context("dataset_strategy")
     cur_rank = get_rank()
@@ -188,33 +267,20 @@ def _single_parameter_broadcast(net, layout, param_not_load=None):
         return
     net_param_dict = net.parameters_dict()
     _chang_parallel_context(origin_dataset_strategy)
-    group_map = _get_group_map()
-    if group_map:
-        group_map = {key: group_map[key] for key in sorted(group_map.keys())}
+    param_loaded.add(cur_rank)
+    total_num = get_group_size()
+    total_param_loaded = [None] * total_num
+    synchronize()
+    all_gather_object(total_param_loaded, param_loaded)
+    total_param_loaded = _check_total_param_loaded(total_param_loaded)
+    group_map = _get_sorted_group_map()
     for group, params in param_redundancy_reversed.items():
-        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:
-                continue
-            if param.startswith("accu_grads") or param.endswith("expert_load"):
-                continue
-            real_param = net_param_dict[param]
-            _remove_param_not_load(real_param.name, param_not_load)
-            if param not in single_params[cur_rank]:
-                real_param.set_data(Tensor(np.zeros(real_param.shape), dtype=real_param.dtype), real_param.sliced)
-            allreduce_input.append(real_param)
-        if not allreduce_input:
-            continue
-        if is_manual_communication_group:
-            create_group(group_name, list(group))
-        allreduce_input.sort(key=lambda param: (str(param.shape), str(param.dtype)))
-        communicator = SingleCommunicator(group_name)
-        for real_param in allreduce_input:
-            real_param.set_data(communicator(Tensor(real_param)), real_param.sliced)
-        if is_manual_communication_group:
-            destroy_group(group_name)
+        allreduce_input = _create_allreduce_input(params, group, net_param_dict, total_param_loaded, param_not_load,
+                                                  cur_rank)
+        _communicate_allreduce(allreduce_input, group_map, group)
     _restore_parallel_context(origin_parallel_mode, origin_dataset_strategy)
+    synchronize()
+    logger.info("End loading the parameter broadcast for removing redundant parameters.")
 
 
 def _insert_virtual_pp_dim(layout):
@@ -259,20 +325,22 @@ class CommTensorDataForPP(Cell):
         self._current_rank_has_data = self._current_rank_id in src_dtensor_info.layout.to_dict()["rank_list"]
         self._diff_rank_id = [
             rank_id for rank_id in dst_dtensor_info.layout.to_dict()["rank_list"] if rank_id not in self._from_rank_id]
-        self.all_reduce = P.AllReduce(group=self._create_all_reduce_group())
+        self._group, self._root_idx = self._create_all_reduce_group()
 
     def comm_data(self, comm_data):
         """communicate data"""
-        out_tensor = self.all_reduce(comm_data)
-        return out_tensor
+        from mindspore import mint
+        comm_handle = mint.distributed.broadcast(comm_data, self._root_idx, self._group, async_op=False)
+        return comm_handle
 
     def _create_all_reduce_group(self):
         """create all reduce group"""
         global ALLREDUCE_GROUP_LIST
         current_rank_stage_id = self._current_rank_id // self._from_dev_num_in_stage
         end_stage = self._from_dev_num_in_stage * (current_rank_stage_id + 1)
-        rank_pos_in_stage = [rank_id for rank_id in range(self._from_dev_num_in_stage * current_rank_stage_id,
-                                                          end_stage)].index(self._current_rank_id)
+        start_stage = self._from_dev_num_in_stage * current_rank_stage_id
+        rank_pos_in_stage = list(range(start_stage, end_stage)).index(self._current_rank_id)
+        root_idx = self._from_rank_id[rank_pos_in_stage]
         all_reduce_rank_list = [self._from_rank_id[rank_pos_in_stage]]
         while rank_pos_in_stage < len(self._diff_rank_id):
             all_reduce_rank_list.append(self._diff_rank_id[rank_pos_in_stage])
@@ -281,11 +349,11 @@ class CommTensorDataForPP(Cell):
         str_rank_list = '-'.join([str(rank) for rank in all_reduce_rank_list])
         all_reduce_group = f"pp_allreduce_group-{str_rank_list}"
         if all_reduce_group in ALLREDUCE_GROUP_LIST:
-            return all_reduce_group
+            return all_reduce_group, root_idx
         ALLREDUCE_GROUP_LIST.append(all_reduce_group)
         create_group(all_reduce_group, all_reduce_rank_list)
         logger.debug(f"Create group {all_reduce_group} for tensor data communication.")
-        return all_reduce_group
+        return all_reduce_group, root_idx
 
 
 class RedistributionCell(Cell):

mindspore/parallel/_parallel_serialization.py

@@ -278,7 +278,7 @@ def _extract_pipeline_stage_num(strategy_file):
             pipeline_stage_set.update(layout[6])
         pipeline_stage_num = len(pipeline_stage_set)
         if list(pipeline_stage_set) != list(range(pipeline_stage_num)):
-            raise ValueError("The strategy file for pipeline parallel dose not contains all stages.")
+            raise ValueError("The strategy file for pipeline parallel does not contains all stages.")
     return pipeline_stage_num
 
 

mindspore/parallel/_ps_context.py

@@ -65,7 +65,6 @@ _set_ps_context_func_map = {
     "enable_ssl": ps_context().set_enable_ssl,
     "client_password": ps_context().set_client_password,
     "server_password": ps_context().set_server_password,
-    "scheduler_manage_port": ps_context().set_scheduler_manage_port,
     "config_file_path": ps_context().set_config_file_path,
 }
 
@@ -78,12 +77,10 @@ _get_ps_context_func_map = {
     "scheduler_ip": ps_context().scheduler_ip,
     "scheduler_port": ps_context().scheduler_port,
     "enable_ssl": ps_context().enable_ssl,
-    "client_password": ps_context().client_password,
-    "server_password": ps_context().server_password,
-    "scheduler_manage_port": ps_context().scheduler_manage_port,
     "config_file_path": ps_context().config_file_path,
 }
 
+
 def _get_ps_mode_rank():
     ps_rank = ps_context().ps_rank_id()
     if ps_rank == -1:
@@ -116,8 +113,9 @@ def _set_ps_context(**kwargs):
                           Only after enable_ps is set True, the environment variables will be effective.
                           Default: ``False``.
         config_file_path (str): Configuration file path used by recovery. Default: ''.
-        scheduler_manage_port (int): scheduler manage port used to scale out/in. Default: 11202.
         enable_ssl (bool): Set PS SSL mode enabled or disabled. Default: ``False``.
+                           There might be risk when this is set to False.
+                           It is user's responsibility to ensure the network environment is safe.
         client_password (str): Password to decrypt the secret key stored in the client certificate. Default: ''.
         server_password (str): Password to decrypt the secret key stored in the server certificate. Default: ''.
 
@@ -126,7 +124,7 @@ def _set_ps_context(**kwargs):
 
     Examples:
         >>> import mindspore as ms
-        >>> ms.set_ps_context(enable_ps=True, enable_ssl=True, client_password='123456', server_password='123456')
+        >>> ms.set_ps_context(enable_ps=True, enable_ssl=True, client_password='', server_password='')
     """
     for key, value in kwargs.items():
         if key not in _set_ps_context_func_map: