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

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,130 @@ 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._group, self._root_idx = self._create_all_reduce_group()
+
+    def comm_data(self, comm_data):
+        """communicate data"""
+        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)
+        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])
+            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, 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, root_idx
+
+
+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

mindspore/parallel/_parallel_serialization.py

@@ -19,6 +19,7 @@ import os
 import json
 import numpy as np
 import mindspore as ms
+from mindspore import _checkparam as Validator
 from mindspore.parallel._tensor import _get_tensor_strategy, _construct_from_to_tensor_layout, \
     _get_needed_rank_list_by_layouts, _get_needed_rank_transform_operator_map_by_layouts, \
     _generate_transform_operator_stack, _apply_tensor_transform_operators, _construct_tensor_layout_for_opt_shard, \
@@ -34,7 +35,12 @@ def _convert_to_list(strategy, rank_id=None):
         try:
             layout = strategy.get(param_name)
             dev_mat = list(layout.dev_matrix[0].dim)
-            tensor_map = list(layout.tensor_map[0].dim)
+            # for layout one axis two slices, layout(("dp", "mp"), "None")
+            if len(layout.tensor_map) > 1:
+                tensor_map = [list(tensor_map.dim) for tensor_map in layout.tensor_map
+                              if list(tensor_map.dim)]
+            else:
+                tensor_map = list(layout.tensor_map[0].dim)
             param_split_shape = list(layout.param_split_shape[0].dim)
             field_size = int(layout.field)
             shard_stride = int(layout.opt_weight_shard_step)
@@ -417,7 +423,7 @@ def _transform_parallel_checkpoint(rank_id, param_total_dict, param_attr_dict, s
         from_opt_shard_size = 0
         if src_strategy_list is not None:
             if param_name not in src_strategy_list:
-                ms.log.warning("The parameter {} is not in src_strategy.".format(param_name))
+                ms.log.info("The parameter {} is not in src_strategy.".format(param_name))
                 continue
             from_dev_matrix, from_tensor_map, from_opt_shard_step, from_opt_shard_size = _extract_layout_item(
                 src_strategy_list.get(param_name))
@@ -427,7 +433,7 @@ def _transform_parallel_checkpoint(rank_id, param_total_dict, param_attr_dict, s
         to_opt_shard_size = 0
         if dst_strategy_list is not None:
             if param_name not in dst_strategy_list:
-                ms.log.warning("The parameter {} is not in dst_strategy.".format(param_name))
+                ms.log.info("The parameter {} is not in dst_strategy.".format(param_name))
                 continue
             to_dev_matrix_origin, to_tensor_map_origin, to_opt_shard_step, to_opt_shard_size = _extract_layout_item(
                 dst_strategy_list.get(param_name))
@@ -441,6 +447,9 @@ def _transform_parallel_checkpoint(rank_id, param_total_dict, param_attr_dict, s
                 continue
             origin_tensor_shape += (item * param_strategy[i],)
 
+        has_layout_from = any(isinstance(i, (list, tuple)) for i in from_tensor_map)
+        has_layout_to = any(isinstance(i, (list, tuple)) for i in to_tensor_map_origin)
+
         from_dev_matrix, from_tensor_map, from_full_tensor_shape = _construct_tensor_layout_for_opt_shard(
             from_dev_matrix, from_tensor_map, from_opt_shard_step, from_opt_shard_size, origin_tensor_shape)
         to_dev_matrix, to_tensor_map, to_full_tensor_shape = _construct_tensor_layout_for_opt_shard(
@@ -460,6 +469,7 @@ def _transform_parallel_checkpoint(rank_id, param_total_dict, param_attr_dict, s
         from_info_tuple = (from_opt_shard_size, from_dev_matrix, from_tensor_map, from_full_tensor_shape)
         to_info_tuple = (to_opt_shard_size, to_dev_matrix_origin, to_tensor_map_origin, origin_tensor_shape)
         _insert_opt_shard_reshape(param_rank_map, from_info_tuple, to_info_tuple)
+        _insert_expand_layout_reshape(param_rank_map, from_info_tuple, to_info_tuple, has_layout_from, has_layout_to)
         transform_operator_stack = _generate_transform_operator_stack(param_rank_map, rank_id)
         param_total_dict_copy = param_total_dict[param_name].copy()
         _apply_tensor_transform_operators(transform_operator_stack, param_total_dict_copy, device_num)
@@ -556,6 +566,32 @@ def _insert_opt_shard_reshape(param_rank_map, from_info_tuple, to_info_tuple):
             param_rank_map.get(param_rank).append(('Reshape', list(to_slice_tensor_shape)))
 
 
+def _insert_expand_layout_reshape(param_rank_map, from_info_tuple, to_info_tuple,
+                                  insert_from_reshape, insert_to_reshape):
+    """ insert layout expand op reshape """
+    from_opt_shard_size = from_info_tuple[0]
+    from_dev_matrix = from_info_tuple[1]
+    from_tensor_map = from_info_tuple[2]
+    from_full_tensor_shape = from_info_tuple[3]
+    to_opt_shard_size = to_info_tuple[0]
+    to_dev_matrix_origin = to_info_tuple[1]
+    to_tensor_map_origin = to_info_tuple[2]
+    origin_tensor_shape = to_info_tuple[3]
+    for param_rank, _ in param_rank_map.items():
+        if from_opt_shard_size == 0 and insert_from_reshape:
+            from_slice_tensor_shape = ()
+            from_tensor_strategy = _get_tensor_strategy(from_dev_matrix, from_tensor_map)
+            for i, item in enumerate(from_full_tensor_shape):
+                from_slice_tensor_shape += (item // from_tensor_strategy[i],)
+            param_rank_map.get(param_rank).insert(0, ('Reshape', list(from_slice_tensor_shape)))
+        if to_opt_shard_size == 0 and insert_to_reshape:
+            to_tensor_strategy = _get_tensor_strategy(to_dev_matrix_origin, to_tensor_map_origin)
+            to_slice_tensor_shape = ()
+            for i, item in enumerate(origin_tensor_shape):
+                to_slice_tensor_shape += (item // to_tensor_strategy[i],)
+            param_rank_map.get(param_rank).append(('Reshape', list(to_slice_tensor_shape)))
+
+
 def _get_param_list_when_first_dim_sharded(device_arrangement, first_dim_sharded_device_index, rank):
     """Calculate rank list for optimizer parallel when first dim of parameter is sharded by other parallel method"""
     total_device_num = 1
@@ -569,4 +605,59 @@ 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
-            
+
+
+def _gather_tasks_load_dis(unified_safetensors_dir, predict_strategy, network, dst_safetensors_dir, dst_device_num,
+                           output_format, name_map, return_param_dict):
+    """gather transform tasks"""
+    tasks = []
+    for rank in range(0, dst_device_num):
+        tasks.append(
+            (unified_safetensors_dir, predict_strategy, network, dst_safetensors_dir, rank, output_format, name_map,
+             return_param_dict))
+    return tasks
+
+
+def _check_checkpoint_file(checkpoint_filenames):
+    """Check checkpoint file name."""
+    for index, filename in enumerate(checkpoint_filenames):
+        if not isinstance(filename, str) or not os.path.exists(filename) \
+                or filename[-5:] != ".ckpt" or os.path.getsize(filename) == 0:
+            raise ValueError(f"For 'load_distributed_checkpoint', please check 'checkpoint_filenames', and "
+                             f"make sure the {filename} at index {index} is a valid checkpoint file, it must "
+                             f"be a string ending with '.ckpt', and the checkpoint file it represents must "
+                             f"be exist and not empty.")
+
+
+def _check_predict_strategy(predict_strategy):
+    """Check predict strategy."""
+
+    def _check_int_list(arg):
+        if not isinstance(arg, list):
+            return False
+        for item in arg:
+            if not isinstance(item, int):
+                return False
+        return True
+
+    if predict_strategy is None:
+        return
+
+    flag = True
+    predict_strategy = Validator.check_isinstance("predict_strategy", predict_strategy, dict)
+    for key in predict_strategy.keys():
+        if not isinstance(key, str) or not isinstance(predict_strategy[key], (list, tuple)) \
+                or len(predict_strategy[key]) < 4:
+            flag = False
+        dev_matrix, tensor_map, param_split_shape, field_size = predict_strategy[key][:4]
+        if not _check_int_list(dev_matrix) or not _check_int_list(tensor_map) or \
+                not (_check_int_list(param_split_shape) or not param_split_shape) or \
+                not (isinstance(field_size, int) and field_size == 0):
+            flag = False
+
+    if not flag:
+        raise ValueError(f"For 'load_distributed_checkpoint', the argument 'predict_strategy' is dict, "
+                         f"the key of it must be string, and the value of it must be list or tuple that "
+                         f"the first four elements must be dev_matrix (list[int]), tensor_map (list[int]), "
+                         f"param_split_shape (list[int]) and field_size (int, which value is 0)."
+                         f"Please check whether 'predict_strategy' is correct.")

mindspore/parallel/_ps_context.py

@@ -115,7 +115,7 @@ def _set_ps_context(**kwargs):
         enable_ps (bool): Whether to enable parameter server training mode.
                           Only after enable_ps is set True, the environment variables will be effective.
                           Default: ``False``.
-        config_file_path (string): Configuration file path used by recovery. Default: ''.
+        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``.
         client_password (str): Password to decrypt the secret key stored in the client certificate. Default: ''.

mindspore/parallel/_recovery_context.py

@@ -33,18 +33,23 @@ def recovery_context():
         RECOVERY_CONTEXT = RecoveryContext.get_instance()
     return RECOVERY_CONTEXT
 
+
 _set_recovery_context_func_map = {
     "ckpt_path": recovery_context().set_ckpt_path,
-    "need_reset": recovery_context().set_need_reset
+    "need_reset": recovery_context().set_need_reset,
+    "is_reboot_node": recovery_context().set_is_reboot_node,
+    "is_arf": recovery_context().set_is_arf
 }
 
 _get_recovery_context_func_map = {
     "enable_recovery": recovery_context().enable_recovery,
+    "enable_repeat_register": recovery_context().enable_repeat_register,
     "latest_ckpt_file": recovery_context().latest_ckpt_file,
     "latest_ckpt_epoch": recovery_context().latest_ckpt_epoch,
     "latest_ckpt_step": recovery_context().latest_ckpt_step,
     "need_reset": recovery_context().need_reset,
     "recovery_path": recovery_context().recovery_path,
+    "is_arf": recovery_context().is_arf,
     "ckpt_path": recovery_context().ckpt_path
 }
 
@@ -64,7 +69,7 @@ def _set_recovery_context(**kwargs):
             MS_RECOVERY_INTERVAL  # The persistent interval for recovery
 
     Args:
-        ckpt_path (string): Set the recovery path used to save checkpoint. Default: ''.
+        ckpt_path (str): Set the recovery path used to save checkpoint. Default: ''.
         need_reset (bool): Set whether should call reset minddata and load ckpt for disaster recovery.
             Default: ``False``.
 

mindspore/parallel/_tensor.py

@@ -38,10 +38,17 @@ def _get_tensor_strategy(dev_mat, tensor_map):
     """
     tensor_strategy = []
     for dim in tensor_map:
-        if dim == -1:
-            tensor_strategy.append(1)
+        if isinstance(dim, (tuple, list)):
+            acc_stra = 1
+            for i in dim:
+                if i != -1:
+                    acc_stra *= dev_mat[len(dev_mat) - i - 1]
+            tensor_strategy.append(acc_stra)
         else:
-            tensor_strategy.append(dev_mat[-dim - 1])
+            if dim == -1:
+                tensor_strategy.append(1)
+            else:
+                tensor_strategy.append(dev_mat[-dim - 1])
     return tensor_strategy
 
 
@@ -182,7 +189,7 @@ def _get_slice_index(dev_mat, tensor_map, opt_shard_group):
     Args:
         dev_mat (list): The device matrix of devices.
         tensor_map (list): The split strategy of tensor.
-        opt_shard_group(string): The group of optimizer shard
+        opt_shard_group(str): The group of optimizer shard
 
     Returns:
         Integer, the slice index for slice on this device.
@@ -388,6 +395,124 @@ def _construct_from_to_tensor_layout(from_full_tensor_shape, from_dev_matrix,
     return from_tensor_layout, to_tensor_layout
 
 
+def _expand_layout(dev_matrix, tensor_map, tensor_shape):
+    """
+    expand nested tensor_map and reshape tensor shape according to tensor_map
+    dev_matrix = [4, 2, 2]
+    tensor_map = [[2, 1], 0]
+    tensor_shape = [8, 8]
+    =>
+    expanded_tensor_map = [2, 1, 0]
+    expanded_tensor_map = [4, 8/4, 8]
+    """
+    new_tensor_map = []
+    new_tensor_shape = []
+    for index, dim in enumerate(tensor_map):
+        if isinstance(dim, (tuple, list)):
+            accu_shape = 1
+            for i in range(len(dim) - 1):
+                new_tensor_map.append(dim[i])
+                new_tensor_shape.append(dev_matrix[len(dev_matrix) - 1 - dim[i]])
+                accu_shape *= dev_matrix[len(dev_matrix) - 1 - dim[i]]
+            new_tensor_map.append(dim[-1])
+            new_tensor_shape.append(tensor_shape[index] // accu_shape)
+        else:
+            new_tensor_map.append(dim)
+            new_tensor_shape.append(tensor_shape[index])
+    return dev_matrix, new_tensor_map, new_tensor_shape
+
+
+def _construct_tensor_layout_for_opt_shard_by_layout(dev_matrix, tensor_map, opt_shard_step, opt_shard_size,
+                                                     origin_full_tensor_shape):
+    """
+    Construct tensor layout for optimizer parallel when using layout.
+    For example, For Tensor with shape (4,2)
+    dev_matrix = [2, 2, 2, 2]
+    tensor_map = [[1, 0], -1]
+    opt_shard_size = 2
+    ==>
+    dev_matrix = [2, 2, 2, 2]
+    tensor_map = [[1, 0], 2, -1]
+    the new strategy is [4, 2, 1]
+    the tensor_shape should reshape to (model_parallel_size, -1, xx, xx)
+    first 4 means the model parallel sharding of data_dim
+    second 2 means the opt sharding of data_dim.
+    """
+    if opt_shard_step == 0 or opt_shard_size == 0:
+        return dev_matrix, tensor_map, list(origin_full_tensor_shape)
+    tensor_strategy = _get_tensor_strategy(dev_matrix, tensor_map)
+    repeated_dim = []
+    dev_sharded_index = []
+    dev_matrix, expanded_tensor_map, _ = _expand_layout(dev_matrix, tensor_map, origin_full_tensor_shape)
+    for dim in expanded_tensor_map:
+        if dim != -1:
+            dev_sharded_index.append(len(dev_matrix) - dim - 1)
+    for index, value in enumerate(dev_matrix):
+        if index not in dev_sharded_index and value > 1:
+            repeated_dim.append(index)
+    if not repeated_dim:
+        raise ValueError("The device_matrix {} and tensor_map {} cannot sharding opt_shard".
+                         format(dev_matrix, tensor_map))
+    return _construct_tensor_layout_helper(dev_matrix, tensor_map, opt_shard_size, origin_full_tensor_shape,
+                                           tensor_strategy, repeated_dim)
+
+
+def _construct_tensor_layout_helper(dev_matrix, tensor_map, opt_shard_size, origin_full_tensor_shape,
+                                    tensor_strategy, repeated_dim):
+    """
+    helper function to assign repeated device_matrix dim for opt shard.
+    """
+    new_dev_matrix = list(copy.deepcopy(dev_matrix))
+    new_dev_matrix_map = list(range(len(dev_matrix)))
+    opt_shard_dim = []
+    remained_opt_shard_size = opt_shard_size if opt_shard_size != -1 else \
+        int(np.prod([dev_matrix[i] for i in repeated_dim]))
+    for dim in repeated_dim[::-1]:
+        opt_sharding_size = dev_matrix[dim]
+        if remained_opt_shard_size // opt_sharding_size == 0:
+            if opt_sharding_size % remained_opt_shard_size != 0:
+                raise ValueError("dev_matrix value {} at dim {} cannot be divided by needed opt sharding "
+                                 "size {}".format(dev_matrix[dim], len(dev_matrix) - dim - 1,
+                                                  remained_opt_shard_size))
+            opt_sharding_size = remained_opt_shard_size
+            # update dev_matrix
+            new_dev_matrix[dim] = dev_matrix[dim] // opt_sharding_size
+            new_dev_matrix.insert(dim + 1, opt_sharding_size)
+            for i in range(len(dev_matrix) - dim - 1, len(dev_matrix)):
+                new_dev_matrix_map[i] += 1
+        if remained_opt_shard_size % opt_sharding_size != 0:
+            raise ValueError("Remained opt_shard_size {} cannot be divided by current sharding size {}, "
+                             "the repeat dim is {} with dev_matrix value {}".
+                             format(remained_opt_shard_size, opt_sharding_size,
+                                    len(dev_matrix) - dim - 1, dev_matrix[dim]))
+        remained_opt_shard_size //= opt_sharding_size
+        opt_shard_dim.insert(0, dim)
+        if remained_opt_shard_size == 1:
+            break
+    tensor_map_new = list(copy.deepcopy(tensor_map))
+    if len(new_dev_matrix) != len(dev_matrix):
+        opt_shard_dim = list(map(lambda x: x + 1, opt_shard_dim))
+        for index, item in enumerate(tensor_map_new):
+            if isinstance(item, (tuple, list)):
+                item = list(map(lambda x: new_dev_matrix_map[x] if x >= 0 else x, item))
+                tensor_map_new[index] = item
+            else:
+                if item >= 0:
+                    tensor_map_new[index] = new_dev_matrix_map[item]
+    tensor_shape_new = list(copy.deepcopy(origin_full_tensor_shape))
+    tensor_shape_new[0] = tensor_strategy[0]
+    first_dim_no_sharding_size = origin_full_tensor_shape[0] // tensor_strategy[0]
+    accu_shape = 1
+    for i in range(len(opt_shard_dim) - 1):
+        opt_sharding_size = new_dev_matrix[opt_shard_dim[i]]
+        tensor_shape_new.insert(i + 1, opt_sharding_size)
+        accu_shape = accu_shape * opt_sharding_size
+    tensor_shape_new.insert(len(opt_shard_dim), first_dim_no_sharding_size // accu_shape)
+    for index, r_dim in enumerate(opt_shard_dim):
+        tensor_map_new.insert(index + 1, len(new_dev_matrix) - r_dim - 1)
+    return list(new_dev_matrix), tensor_map_new, tensor_shape_new
+
+
 def _construct_tensor_layout_for_opt_shard(dev_matrix, tensor_map, opt_shard_step, opt_shard_size,
                                            origin_full_tensor_shape):
     """
@@ -404,6 +529,11 @@ def _construct_tensor_layout_for_opt_shard(dev_matrix, tensor_map, opt_shard_ste
     And the model parallel sharding dim is the right of opt sharding dim, so it would be 0-1-2-3 model parallel sharding
     then 0-4 optimizer sharding.
     """
+    has_layout = any(isinstance(i, (list, tuple)) for i in tensor_map)
+    if has_layout:
+        output = _construct_tensor_layout_for_opt_shard_by_layout(dev_matrix, tensor_map, opt_shard_step,
+                                                                  opt_shard_size, origin_full_tensor_shape)
+        return _expand_layout(*output)
 
     if opt_shard_step == 0 or opt_shard_size == 0:
         return dev_matrix, tensor_map, list(origin_full_tensor_shape)
@@ -424,18 +554,8 @@ def _construct_tensor_layout_for_opt_shard(dev_matrix, tensor_map, opt_shard_ste
                          format(opt_shard_step, np.prod(dev_matrix[repeated_dim[0] + 1:])))
     first_dim_no_sharding_size = origin_full_tensor_shape[0] // tensor_strategy[0]
     if (len(repeated_dim) < len(dev_matrix) and len(repeated_dim) > 1) or repeated_dim[0] > 0:
-        tensor_shape_new = list(origin_full_tensor_shape)
-        tensor_shape_new[0] = tensor_strategy[0]
-        accu_shp = 1
-        for i in range(len(repeated_dim) - 1):
-            opt_sharding_size = dev_matrix[repeated_dim[i]]
-            tensor_shape_new.insert(i + 1, opt_sharding_size)
-            accu_shp = accu_shp * opt_sharding_size
-        tensor_shape_new.insert(len(repeated_dim), first_dim_no_sharding_size // accu_shp)
-        tensor_map_new = list(copy.deepcopy(tensor_map))
-        for index, r_dim in enumerate(repeated_dim):
-            tensor_map_new.insert(index + 1, len(dev_matrix) - r_dim - 1)
-        return list(dev_matrix), tensor_map_new, tensor_shape_new
+        return _construct_tensor_layout_helper(dev_matrix, tensor_map, opt_shard_size, origin_full_tensor_shape,
+                                               tensor_strategy, repeated_dim)
 
     full_tensor_shape = list(origin_full_tensor_shape)
     full_tensor_shape[0] = tensor_strategy[0]
@@ -610,9 +730,13 @@ def _apply_operator(operator_name):
         """
         if not isinstance(numpy_data_list, list):
             raise TypeError("The data_list should be a list.")
+        new_numpy_data_list = []
         for numpy_data in numpy_data_list:
-            if not isinstance(numpy_data, np.ndarray):
-                raise TypeError("The data should be a numpy.ndarray.")
+            if str(type(numpy_data)) == "<class 'builtins.PySafeSlice'>":
+                new_numpy_data_list.append(numpy_data[:])
+            else:
+                new_numpy_data_list.append(numpy_data)
+        numpy_data_list = new_numpy_data_list
         _check_operator(allgather_op)
         concat_group = allgather_op[1][:-1]
         if len(concat_group) != len(numpy_data_list):