mindspore 2.4.10__cp311-cp311-win_amd64.whl → 2.6.0rc1__cp311-cp311-win_amd64.whl

mindspore/parallel/_auto_parallel_context.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2023 Huawei Technologies Co., Ltd
+# Copyright 2020-2025 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.
@@ -21,6 +21,7 @@ from mindspore import context
 import mindspore.log as logger
 from mindspore.parallel._dp_allreduce_fusion import _set_fusion_strategy_by_idx, _set_fusion_strategy_by_size
 from mindspore.parallel._ps_context import _is_role_pserver
+from mindspore.parallel.shard import Layout
 from mindspore._c_expression import AutoParallelContext
 from mindspore._checkparam import args_type_check
 from mindspore import _checkparam as Validator
@@ -63,6 +64,7 @@ class _ParallelOptimizerConfig:
     GRADIENT_ACCUMULATION_SHARD = "gradient_accumulation_shard"
     PARALLEL_OPTIMIZER_THRESHOLD = "parallel_optimizer_threshold"
     OPTIMIZER_WEIGHT_SHARD_SIZE = "optimizer_weight_shard_size"
+    OPTIMIZER_LEVEL = "optimizer_level"
 
 
 class _PipelineConfig:
@@ -77,6 +79,8 @@ class _PipelineScheduler:
     PIPELINE_1F1B = "1f1b"
     PIPELINE_GPIPE = "gpipe"
     PIPELINE_SEQPIPE = "seqpipe"
+    PIPELINE_SEQVPP = "seqvpp"
+    PIPELINE_SEQSMARTVPP = "seqsmartvpp"
 
 
 class _AutoParallelContext:
@@ -100,6 +104,7 @@ class _AutoParallelContext:
     def __init__(self):
         self._context_handle = AutoParallelContext.get_instance()
         self._dataset_strategy_using_str = True
+        self._dataset_layout = None
 
     def check_context_handle(self):
         """
@@ -187,6 +192,25 @@ 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.
@@ -195,7 +219,7 @@ class _AutoParallelContext:
             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 in not a bool value.
+            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)
@@ -422,6 +446,9 @@ class _AutoParallelContext:
             raise ValueError("The context configuration parameter 'parallel_mode' only support 'stand_alone', "
                              "'data_parallel', 'hybrid_parallel', 'semi_auto_parallel' and 'auto_parallel', "
                              "but got the value : {}.".format(parallel_mode))
+        if run_mode == context.ParallelMode.DATA_PARALLEL and self.get_enable_parallel_optimizer():
+            logger.warning("'enable_parallel_optimizer' is not suggested in 'data_parallel' mode, "
+                           "consider using 'semi_auto_parallel' or 'auto_parallel' mode.")
 
     def get_parallel_mode(self):
         """Get parallel mode."""
@@ -566,6 +593,9 @@ class _AutoParallelContext:
         if not isinstance(dataset_strategy, tuple):
             raise TypeError("For 'set_auto_parallel_context', the argument 'dataset_strategy' "
                             "must be str or tuple type, but got the type : {}.".format(type(dataset_strategy)))
+        if dataset_strategy and isinstance(dataset_strategy[0], Layout):
+            self._set_dataset_strategy_layout(dataset_strategy)
+            return
         for ele in dataset_strategy:
             if not isinstance(ele, tuple):
                 raise TypeError("For 'set_auto_parallel_context', the element of argument "
@@ -580,8 +610,36 @@ class _AutoParallelContext:
         self._dataset_strategy_using_str = False
         self._context_handle.set_dataset_strategy(dataset_strategy)
 
+    def _set_dataset_strategy_layout(self, dataset_strategy):
+        """set dataset layout to c++ by using pybind."""
+        dataset_devmat = []
+        dataset_tensormap = []
+        dataset_alias_name = []
+        self._dataset_layout = dataset_strategy
+        for ele in dataset_strategy:
+            if not isinstance(ele, Layout):
+                raise TypeError(f"All the dataset_strategy elements should be Layout, but got {type(ele)}")
+            layout_to_dict = ele.to_dict()
+            dataset_devmat.append(layout_to_dict["device_matrix"])
+            dataset_alias_name.append(layout_to_dict["alias_name"])
+            if layout_to_dict["interleaved_parallel"]:
+                raise ValueError("For dataset_strategy, layout does not support interleaved_parallel")
+            tensor_map = []
+            for value in layout_to_dict["tensor_map"]:
+                if isinstance(value, tuple):
+                    tensor_map.append(value)
+                elif isinstance(value, int):
+                    tensor_map.append((value,))
+                else:
+                    raise TypeError(f"value in tensor map must be tuple or int, but got {type(value)}")
+            dataset_tensormap.append(tuple(tensor_map))
+        self._context_handle.set_dataset_layout(dataset_devmat, dataset_tensormap, dataset_alias_name)
+
+
     def get_dataset_strategy(self):
         """Get dataset sharding strategy."""
+        if self._dataset_layout is not None:
+            return self._dataset_layout
         self.check_context_handle()
         if self._dataset_strategy_using_str:
             if self._context_handle.get_full_batch():
@@ -869,6 +927,9 @@ class _AutoParallelContext:
                             "the argument 'enable_parallel_optimizer' must be bool, but got the type : {}."
                             .format(type(enable_parallel_optimizer)))
         self._context_handle.set_enable_parallel_optimizer(enable_parallel_optimizer)
+        if enable_parallel_optimizer and self.get_parallel_mode() == context.ParallelMode.DATA_PARALLEL:
+            logger.warning("'enable_parallel_optimizer' is not suggested in 'data_parallel' mode, "
+                           "consider using 'semi_auto_parallel' or 'auto_parallel' mode.")
 
     def set_force_fp32_communication(self, force_fp32_communication):
         """
@@ -899,7 +960,7 @@ class _AutoParallelContext:
 
             - 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"
+            - pipeline_scheduler(str): 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.
@@ -935,7 +996,9 @@ class _AutoParallelContext:
 
         Validator.check_string(pipeline_config[pp_scheduler], [_PipelineScheduler.PIPELINE_1F1B,
                                                                _PipelineScheduler.PIPELINE_GPIPE,
-                                                               _PipelineScheduler.PIPELINE_SEQPIPE])
+                                                               _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")
 
@@ -984,10 +1047,12 @@ class _AutoParallelContext:
         grad_shard_name = _ParallelOptimizerConfig.GRADIENT_ACCUMULATION_SHARD
         threshold_name = _ParallelOptimizerConfig.PARALLEL_OPTIMIZER_THRESHOLD
         optimizer_weight_shard_size_name = _ParallelOptimizerConfig.OPTIMIZER_WEIGHT_SHARD_SIZE
+        optimizer_level_name = _ParallelOptimizerConfig.OPTIMIZER_LEVEL
 
         for config_name in parallel_optimizer_config:
             unknown_config = []
-            if config_name not in [grad_shard_name, threshold_name, optimizer_weight_shard_size_name]:
+            if config_name not in [grad_shard_name, threshold_name, optimizer_weight_shard_size_name,
+                                   optimizer_level_name]:
                 unknown_config.append(config_name)
 
             if unknown_config:
@@ -998,6 +1063,10 @@ class _AutoParallelContext:
                 parallel_optimizer_config[grad_shard_name], grad_shard_name, grad_shard_name)
             self._context_handle.set_grad_accumulation_shard(
                 parallel_optimizer_config[grad_shard_name])
+            if optimizer_level_name in parallel_optimizer_config \
+                    and parallel_optimizer_config[optimizer_level_name] != "level2":
+                raise ValueError(f"The optimizer_level is set as {parallel_optimizer_config[optimizer_level_name]}, "
+                                 "thus cannot set grad_accumulation_shard as True.")
 
         if threshold_name in parallel_optimizer_config:
             Validator.check_non_negative_int(
@@ -1010,6 +1079,20 @@ class _AutoParallelContext:
             Validator.check_positive_int(value)
             self.set_optimizer_weight_shard_size(value)
 
+        if optimizer_level_name in parallel_optimizer_config:
+            optimizer_level = parallel_optimizer_config[optimizer_level_name]
+            if optimizer_level not in ["level1", "level2", "level3"]:
+                raise ValueError("Optimizer level should in ['level1', 'level2', 'level3'], but got {}"
+                                 .format(optimizer_level))
+
+            if self._context_handle.get_grad_accumulation_shard() and optimizer_level != "level2":
+                raise ValueError("The grad_accumulation shard is set, thus cannot set optimizer_level != 'level2'")
+            if optimizer_level == "level2":
+                self._context_handle.set_grad_accumulation_shard(True)
+            if optimizer_level == "level3":
+                self._context_handle.set_zero3(True)
+                self._context_handle.set_grad_accumulation_shard(False)
+
     def get_grad_accumulation_shard(self):
         """Get grad accumulation shard."""
         self.check_context_handle()
@@ -1089,13 +1172,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.
@@ -1124,6 +1200,7 @@ class _AutoParallelContext:
         self.check_context_handle()
         self._context_handle.reset()
         _ParallelFusionConfig.reset()
+        self._dataset_layout = None
 
     def _check_and_default_group(self, group):
         """Validate the given group, if group is empty, returns a default fusion group"""
@@ -1233,30 +1310,35 @@ class _AutoParallelContext:
             self.set_enable_all_gather_fusion(openstate)
             self.set_enable_reduce_scatter_fusion(openstate)
 
+    def set_auto_parallel_new_interface(self, auto_parallel_new_interface):
+        """
+        Set AutoParallel(cell) new interface flag.
 
-def _set_ops_strategy_json_config(type="SAVE", path="", mode="all"):
-    """
-    Set strategy json configuration.
+        Args:
+            auto_parallel_new_interface (bool): Mark whether to use the new interface.
+        """
+        self.check_context_handle()
+        self._context_handle.set_auto_parallel_new_interface(auto_parallel_new_interface)
 
-    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.
+    def get_auto_parallel_new_interface(self):
+        """Get auto_parallel_new_interface."""
+        self.check_context_handle()
+        return self._context_handle.get_auto_parallel_new_interface()
 
-    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'")
+    def set_init_param_in_compile(self, init_param_in_compile):
+        """
+        Set flag marking whether to init parameters in compiling process.
+
+        Args:
+            init_param_in_compile (bool): Mark whether to init parameters in compiling process.
+        """
+        self.check_context_handle()
+        self._context_handle.set_init_param_in_compile(init_param_in_compile)
 
+    def get_init_param_in_compile(self):
+        """Get init_param_in_compile."""
+        self.check_context_handle()
+        return self._context_handle.get_init_param_in_compile()
 
 _AUTO_PARALLEL_CONTEXT = None
 
@@ -1307,7 +1389,11 @@ _set_auto_parallel_context_func_map = {
     "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_device_local_norm": auto_parallel_context().set_dump_device_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,
+    "auto_parallel_new_interface": auto_parallel_context().set_auto_parallel_new_interface,
+    "init_param_in_compile": auto_parallel_context().set_init_param_in_compile}
+
 
 _get_auto_parallel_context_func_map = {
     "device_num": auto_parallel_context().get_device_num,
@@ -1341,7 +1427,10 @@ _get_auto_parallel_context_func_map = {
     "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_device_local_norm": auto_parallel_context().get_dump_device_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,
+    "auto_parallel_new_interface": auto_parallel_context().get_auto_parallel_new_interface,
+    "init_param_in_compile": auto_parallel_context().get_init_param_in_compile}
 
 
 @args_type_check(device_num=int, global_rank=int, gradients_mean=bool, gradient_fp32_sync=bool,
@@ -1452,6 +1541,8 @@ def _set_auto_parallel_context(**kwargs):
                     - reducescatter: If communication fusion type is `reducescatter`. The `mode` contains: `auto`
                         and `size`. Config is same as `allgather`.
 
+
+
     Raises:
         ValueError: If input key is not attribute in auto parallel context.
     """
@@ -1481,6 +1572,33 @@ def _get_auto_parallel_context(attr_key):
     return get_func()
 
 
+def _get_all_auto_parallel_context():
+    """get auto parallel context before reset"""
+    _auto_paralell_context_value_map = {}
+    _pipeline_config = {}
+    for key, value in _get_auto_parallel_context_func_map.items():
+        if key == "pipeline_interleave":
+            _pipeline_config[key] = value()
+        elif key == "pipeline_scheduler":
+            _pipeline_config[key] = value()
+        else:
+            _auto_paralell_context_value_map[key] = value()
+    return _auto_paralell_context_value_map, _pipeline_config
+
+
+def _recover_auto_parallel_context(context_value_map, pp_config):
+    """set auto parallel context after transformation"""
+    # set the same auto parallel context after transform
+    from mindspore.context import reset_auto_parallel_context
+    reset_auto_parallel_context()
+    for key, value in context_value_map.items():
+        # list is empty or full_batch_is_set is not needed to set
+        if (isinstance(value, list) and not value) or (key == "full_batch_is_set"):
+            continue
+        _set_auto_parallel_context_func_map[key](value)
+    _set_auto_parallel_context_func_map["pipeline_config"](pp_config)
+
+
 def _reset_auto_parallel_context():
     """
     Reset auto parallel context attributes to the default values:

mindspore/parallel/_cell_wrapper.py

@@ -18,17 +18,20 @@ from __future__ import division
 
 import numpy as np
 
+import mindspore.log as logger
 from mindspore import context
 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.communication import GlobalComm, get_rank
 from mindspore.common import jit
-from mindspore.communication import create_group, destroy_group
+from mindspore.communication import create_group, destroy_group, get_group_size
 from mindspore.communication._comm_helper import _get_group_map
 from mindspore.train._utils import get_parameter_redundancy, remove_param_redundancy
+from mindspore.parallel.shard import Layout
 
 _ALLGATHER_CELL = None
+ALLREDUCE_GROUP_LIST = []
 
 
 class AllGatherCell(Cell):
@@ -134,7 +137,7 @@ def _restore_parallel_context(origin_parallel_mode, origin_dataset_strategy):
 
 def _get_group_name(group_map, group):
     """get group name"""
-    group_name = str(group)
+    group_name = "remove_redundancy" + str(group)
     is_manual_communication_group = True
     if group_map:
         for name, rank_list in group_map.items():
@@ -142,20 +145,37 @@ def _get_group_name(group_map, group):
                 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):
+def _get_param_redundancy_reversed(param_redundancy, cur_rank):
+    """Generate the reverse mapping of parameter redundancy based on the current rank."""
+    param_redundancy_reversed = {}
+    for key, redundancy in param_redundancy.items():
+        for item in redundancy:
+            if len(item) == 1:
+                continue
+            if cur_rank in item:
+                param_redundancy_reversed.setdefault(item, []).append(key)
+    return param_redundancy_reversed
+
+
+def _remove_param_not_load(param_name, param_not_load):
+    """Remove param_name from param_not_load."""
+    if param_not_load is not None and param_name in param_not_load:
+        param_not_load.remove(param_name)
+
+
+def _single_parameter_broadcast(net, layout, param_not_load=None):
     """
     Broadcast single parameter to other rank in data parallel dimension.
     """
     from mindspore import Tensor
     origin_parallel_mode = context.get_auto_parallel_context("parallel_mode")
     origin_dataset_strategy = context.get_auto_parallel_context("dataset_strategy")
+    cur_rank = get_rank()
     if layout:
-        param_redundancy = get_parameter_redundancy(layout, initial_rank)
+        param_redundancy = get_parameter_redundancy(layout)
     else:
         param_redundancy = get_parameter_redundancy(net)
     if not param_redundancy:
@@ -163,33 +183,128 @@ def _single_parameter_broadcast(net, layout, cur_rank=0, initial_rank=0):
     single_params = remove_param_redundancy(param_redundancy)
     if not single_params:
         return
-    param_redundancy_reversed = {}
-    for key, redundancy in param_redundancy.items():
-        for item in redundancy:
-            if len(item) == 1:
-                continue
-            if cur_rank in item:
-                param_redundancy_reversed.setdefault(item, []).append(key)
+    param_redundancy_reversed = _get_param_redundancy_reversed(param_redundancy, cur_rank)
     if not param_redundancy_reversed or cur_rank not in single_params:
         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())}
     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(real_param), real_param.sliced)
+            real_param.set_data(communicator(Tensor(real_param)), real_param.sliced)
         if is_manual_communication_group:
             destroy_group(group_name)
     _restore_parallel_context(origin_parallel_mode, origin_dataset_strategy)
+
+
+def _insert_virtual_pp_dim(layout):
+    """insert virtual pp dim in device matrix and create new layout"""
+    if len(layout.to_dict()["rank_list"]) == get_group_size():
+        return layout
+    remain_pp = get_group_size() // len(layout.to_dict()["rank_list"])
+    layout_info = layout.to_dict()
+    device_matrix = layout_info["device_matrix"]
+    tensor_map = layout_info["tensor_map"]
+    alias_name = layout_info["alias_name"]
+    new_devmat = Layout((remain_pp,) + device_matrix, ("remain_pp",) + alias_name)
+    tensor_map_alias_name = []
+    for val in tensor_map:
+        sub_alias_name = []
+        if isinstance(val, tuple):
+            for sub_val in val:
+                if sub_val == -1:
+                    sub_alias_name.append("None")
+                else:
+                    sub_alias_name.append(alias_name[len(device_matrix) - sub_val - 1])
+            tensor_map_alias_name.append(tuple(sub_alias_name))
+        else:
+            if val == -1:
+                tensor_map_alias_name.append("None")
+            else:
+                tensor_map_alias_name.append(alias_name[len(device_matrix) - val - 1])
+    new_layout = new_devmat(*tensor_map_alias_name)
+    return new_layout
+
+
+class CommTensorDataForPP(Cell):
+    """Communicate tensor data for pipeline parallel scenario."""
+
+    def __init__(self, src_dtensor_info, dst_dtensor_info):
+        super().__init__()
+        self.zeros = P.Zeros()
+
+        self._current_rank_id = get_rank()
+        self._from_dev_num_in_stage = len(src_dtensor_info.layout.to_dict()["rank_list"])
+        self._from_rank_id = src_dtensor_info.layout.to_dict()["rank_list"]
+        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())
+
+    def comm_data(self, comm_data):
+        """communicate data"""
+        out_tensor = self.all_reduce(comm_data)
+        return out_tensor
+
+    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)
+        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])
+            rank_pos_in_stage += self._from_dev_num_in_stage
+        all_reduce_rank_list.sort()
+        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
+        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
+
+
+class RedistributionCell(Cell):
+    """Redistribute src_layout to dst_layout"""
+
+    def __init__(self, src_layout, dst_layout):
+        super().__init__()
+        if src_layout is None or dst_layout is None:
+            raise ValueError("src_layout and dst_layout should not be None.")
+        self._total_dev_num = get_group_size()
+        src_layout = _insert_virtual_pp_dim(src_layout)
+        dst_layout = _insert_virtual_pp_dim(dst_layout)
+        self.src_identity = P.Identity().shard(in_strategy=(src_layout,), out_strategy=(src_layout,))
+        self.src_identity.add_prim_attr("self_define_shard", True)
+        self.dst_identity = P.Identity().shard(in_strategy=(dst_layout,), out_strategy=(dst_layout,))
+        self.dst_identity.add_prim_attr("self_define_shard", True)
+
+    def construct(self, input_tensor):
+        """run redistribution"""
+        src_tensor = self.src_identity(input_tensor)
+        dst_tensor = self.dst_identity(src_tensor)
+        return dst_tensor