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

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)
@@ -115,11 +121,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 +198,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 +352,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)
 
 
@@ -407,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))
@@ -417,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))
@@ -431,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(
@@ -450,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)
@@ -546,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
@@ -559,3 +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):