mindspore 2.3.0__cp39-cp39-win_amd64.whl → 2.4.1__cp39-cp39-win_amd64.whl

mindspore/parallel/_utils.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 # ============================================================================
 """Utils of auto parallel"""
+import os
+from time import perf_counter
 from importlib import import_module
 import numpy as np
 import mindspore as ms
@@ -22,12 +24,13 @@ from mindspore.common.tensor import Tensor
 from mindspore.common.dtype import dtype_to_nptype
 from mindspore.common import dtype as mstype
 from mindspore.communication.management import get_group_size, get_rank
+from mindspore.communication._comm_helper import _is_initialized
 from mindspore.parallel._auto_parallel_context import auto_parallel_context
 from mindspore.common.seed import get_seed
 from mindspore._c_expression import GraphExecutor_
-from mindspore.parallel._tensor import _load_tensor_by_layout
+from mindspore.parallel._tensor import _load_tensor_by_layout, _load_tensor_shape_by_layout
 
-SUPPORTED_TUPLE_IN_TUPLE_STRATEGY = ["GroupedMatmul", "FusedInferAttentionScore"]
+SUPPORTED_TUPLE_IN_TUPLE_STRATEGY = ["GroupedMatmul", "FusedInferAttentionScore", "Custom"]
 
 
 def _get_parallel_mode():
@@ -45,6 +48,16 @@ def _is_in_auto_parallel_mode():
     return _get_parallel_mode() in [ms.ParallelMode.SEMI_AUTO_PARALLEL, ms.ParallelMode.AUTO_PARALLEL]
 
 
+def _is_parallel_mode():
+    if not _is_initialized() or context.get_context('mode') == context.PYNATIVE_MODE:
+        return False
+    if os.getenv("RUN_MODE") != "predict":
+        return False
+    if get_group_size() > 1 and _get_parallel_mode() == ms.ParallelMode.STAND_ALONE:
+        return True
+    return False
+
+
 def _is_in_data_parallel_mode():
     return _get_parallel_mode() == ms.ParallelMode.DATA_PARALLEL
 
@@ -92,31 +105,49 @@ def _need_to_full():
     return not _get_full_batch()
 
 
+class ParallelParamInitProfCtx:
+    """Collect parallel param initialization performance context mgr."""
+
+    def __init__(self, parameter, func_name):
+        self.parameter = parameter
+        self.func_name = func_name
+        self.start_timestamp = None
+
+    def __enter__(self):
+        self.start_timestamp = perf_counter()
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        end_timestamp = perf_counter()
+        duration = end_timestamp - self.start_timestamp
+        if os.getenv("MS_DEV_PARAM_INIT_PROF_COLLECT"):
+            logger.warning(f"{self.func_name}: {self.parameter.name}, shape: {self.parameter.shape}, "
+                           f"sliced: {self.parameter.sliced}, duration: {duration}")
+
+
 def _slice_parameter(parameter, phase, layout):
     """Slice python parameter obj according to the layout."""
-    is_train_phase = phase.startswith('train')
-    is_prefill_phase = phase.startswith('prefill')
-    if layout is not None and parameter.from_ckpt and not is_train_phase:
-        is_opt_shard_group = layout[5]
-        if not parameter.sliced and is_prefill_phase and is_opt_shard_group:
+    # graph_executor.updata_param_node_default_input(phase, {parameter.name: parameter})
+    if getattr(parameter, "init_param", False):
+        if layout is None:
+            parameter.sliced = True
+            return
+        if not parameter.sliced:
+            rank = get_rank()
+            new_tensor_shape = _load_tensor_shape_by_layout(parameter, layout, rank)
+            parameter.shape = new_tensor_shape
+    else:
+        graph_executor = GraphExecutor_.get_instance()
+        new_param = parameter.init_data(layout, set_sliced=True)
+        parameter = new_param
+        graph_executor.updata_param_node_default_input(phase, {parameter.name: parameter})
+        if layout is None:
+            parameter.sliced = True
+            return
+        if not parameter.sliced:
             rank = get_rank()
             new_tensor = _load_tensor_by_layout(parameter, layout, rank)
             parameter.set_data(new_tensor, True)
-            return
-        layout_shape = layout[2]
-        parameter.shape = tuple(layout_shape)
-        return
-    graph_executor = GraphExecutor_.get_instance()
-    new_param = parameter.init_data(layout, set_sliced=True)
-    parameter = new_param
-    graph_executor.updata_param_node_default_input(phase, {parameter.name: parameter})
-    if layout is None:
-        parameter.sliced = True
-        return
-    if not parameter.sliced:
-        rank = get_rank()
-        new_tensor = _load_tensor_by_layout(parameter, layout, rank)
-        parameter.set_data(new_tensor, True)
 
 
 def _slice_tensor(tensor, layout, rank_id):

mindspore/parallel/algo_parameter_config.py

@@ -234,7 +234,7 @@ def set_algo_parameters(**kwargs):
 
     Args:
         fully_use_devices (bool): Whether ONLY searching strategies that fully use all available devices.
-            Default: ``True`` . For example with 8 devices available, if set ``True`` , strategy (4, 1) will not be
+            Default: ``False`` . For example with 8 devices available, if set ``True`` , strategy (4, 1) will not be
             included in ReLU's candidate strategies, because strategy (4, 1) only utilizes 4 devices.
         elementwise_op_strategy_follow (bool): Whether the elementwise operator has the consistent strategies as its
             subsequent operators. Elementwise operators refer to operators that operate on input element by element,
@@ -264,14 +264,14 @@ def set_algo_parameters(**kwargs):
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
             Please see the `rank table startup
-            <https://www.mindspore.cn/tutorials/experts/en/master/parallel/rank_table.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/rank_table.html>`_
             for more details.
 
             For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun startup
-            <https://www.mindspore.cn/tutorials/experts/en/master/parallel/mpirun.html>`_ .
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/mpirun.html>`_ .
 
             For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
-            Startup <https://www.mindspore.cn/tutorials/experts/en/master/parallel/dynamic_cluster.html>`_ .
+            Startup <https://www.mindspore.cn/docs/en/master/model_train/parallel/dynamic_cluster.html>`_ .
 
         >>> import numpy as np
         >>> import mindspore as ms
@@ -386,7 +386,7 @@ def reset_algo_parameters():
 
     After reset, the values of the attributes are:
 
-    - fully_use_devices: True.
+    - fully_use_devices: False.
     - elementwise_op_strategy_follow: False.
     - enable_algo_approxi: False.
     - algo_approxi_epsilon: 0.1.

mindspore/parallel/checkpoint_transform.py

@@ -22,12 +22,12 @@ from collections import defaultdict
 import numpy as np
 import mindspore as ms
 from mindspore.common import dtype as mstype
-from mindspore.parallel._utils import _is_in_auto_parallel_mode
+from mindspore.parallel._utils import _is_in_auto_parallel_mode, _get_pipeline_stages
 from mindspore.parallel._parallel_serialization import _rank_list_for_transform_parallel_checkpoint, \
     _transform_parallel_checkpoint, _get_device_num_from_strategy, _make_dir, \
     _extract_layout_map, _extract_src_dst_layout_map, _parameter_not_in_local_stage, _extract_pipeline_stage_num, \
     _merge_protobuf_strategy, _merge_json_strategy, _extract_src_dst_layout_map_by_src
-
+from mindspore.parallel.transform_safetensors import _transform_safetensors, _collect_safetensor_files
 
 __all__ = ["merge_pipeline_strategys", "rank_list_for_transform", "transform_checkpoint_by_rank",
            "transform_checkpoints", "sync_pipeline_shared_parameters", "load_segmented_checkpoints"]
@@ -37,7 +37,7 @@ def merge_pipeline_strategys(src_strategy_dirs, dst_strategy_file):
     """
     Merge parallel strategy between all pipeline stages in pipeline parallel mode.
     For more details about converting distributed Checkpoint, please refer to
-    `Model Transformation <https://www.mindspore.cn/tutorials/experts/en/master/parallel/model_transformation.html>`_.
+    `Model Transformation <https://www.mindspore.cn/docs/en/master/model_train/parallel/model_transformation.html>`_.
 
     Note:
         Strategy file of each pipeline stage should be included in src_strategy_dirs.
@@ -72,12 +72,11 @@ def merge_pipeline_strategys(src_strategy_dirs, dst_strategy_file):
         _merge_json_strategy(src_strategy_files_json, dst_strategy_file)
 
 
-
 def rank_list_for_transform(rank_id, src_strategy_file=None, dst_strategy_file=None):
     """
     List of original distributed checkpoint rank index for obtaining the target checkpoint of a rank_id during the
     distributed checkpoint conversion. For more details about converting distributed Checkpoint, please refer to
-    `Model Transformation <https://www.mindspore.cn/tutorials/experts/en/master/parallel/model_transformation.html>`_.
+    `Model Transformation <https://www.mindspore.cn/docs/en/master/model_train/parallel/model_transformation.html>`_.
 
     Args:
         rank_id (int): The rank of which distributed checkpoint needs to be obtained after conversion.
@@ -132,7 +131,9 @@ def rank_list_for_transform(rank_id, src_strategy_file=None, dst_strategy_file=N
             src_rank_id_start = src_pipeline_stage_id * src_stage_device_num
             result_set.update([src_rank_id_start + rank for rank in needed_rank_list_in_local_stage])
             handled_pipeline_stage.append(src_pipeline_stage_id)
-    return list(result_set)
+    result_list = list(result_set)
+    result_list.sort(reverse=True)
+    return list(result_list)
 
 
 def transform_checkpoint_by_rank(rank_id, checkpoint_files_map, save_checkpoint_file_name,
@@ -140,7 +141,7 @@ def transform_checkpoint_by_rank(rank_id, checkpoint_files_map, save_checkpoint_
     """
     Transform distributed checkpoint from source sharding strategy to destination sharding strategy by rank
     for a network. For more details about converting distributed Checkpoint, please refer to
-    `Model Transformation <https://www.mindspore.cn/tutorials/experts/en/master/parallel/model_transformation.html>`_.
+    `Model Transformation <https://www.mindspore.cn/docs/en/master/model_train/parallel/model_transformation.html>`_.
 
     Args:
         rank_id (int): The rank of which distributed checkpoint needs to be obtained after conversion.
@@ -232,7 +233,7 @@ def _transform_checkpoint_by_stage(src_checkpoints_dir, dst_checkpoints_dir, ckp
     param_attr_dict = defaultdict(dict)
     param_type_dict = defaultdict(dict)
     src_strategy_list, dst_strategy_list, stage_id = _extract_src_dst_layout_map_by_src(src_strategy_file, \
-                                                                                             dst_strategy_file)
+                                                                                        dst_strategy_file)
     src_stage_device_num = np.prod(src_strategy_list.get(list(src_strategy_list.keys())[0])[0]) if src_strategy_list \
                                                                                                    is not None else 1
     dst_stage_device_num = np.prod(dst_strategy_list.get(list(dst_strategy_list.keys())[0])[0]) if dst_strategy_list \
@@ -357,29 +358,35 @@ def _transform_checkpoints(src_checkpoints_dir, dst_checkpoints_dir, ckpt_prefix
 
 
 def transform_checkpoints(src_checkpoints_dir, dst_checkpoints_dir, ckpt_prefix, src_strategy_file=None,
-                          dst_strategy_file=None):
+                          dst_strategy_file=None, process_num=1, output_format="ckpt"):
     """
     Transform distributed checkpoint from source sharding strategy to destination sharding strategy for a rank.
     For more details about converting distributed Checkpoint, please refer to
-    `Model Transformation <https://www.mindspore.cn/tutorials/experts/en/master/parallel/model_transformation.html>`_.
+    `Model Transformation <https://www.mindspore.cn/docs/en/master/model_train/parallel/model_transformation.html>`_.
 
     Note:
         The `src_checkpoints_dir` directory structure should be organized like "src_checkpoints_dir/rank_0/a.ckpt", the
         rank number should be set to a subdirectory and the checkpoint file is stored in this subdirectory. If multiple
         files exist in a rank directory, the last file in the lexicgraphic order would be selected.
 
+        The number of multiprocess settings is related to the size of the host, and it is not recommended to set it
+        too large, otherwise it may cause freezing.
+
     Args:
         src_checkpoints_dir (str): The source checkpoints directory.
         dst_checkpoints_dir (str): The destination checkpoints directory to save the converted checkpoints.
         ckpt_prefix (str): The destination checkpoint name prefix.
-        src_strategy_file (str): Name of source sharding strategy file which saved by
+        src_strategy_file (str, optional): Name of source sharding strategy file which saved by
                                  'mindspore.set_auto_parallel_context(strategy_ckpt_save_file)'.
                                  when the 'src_strategy_file' is None, it means that the source sharding strategy is
                                  without any sharing for each parameter. Default:None.
-        dst_strategy_file (str): Name of destination sharding strategy file which saved by
+        dst_strategy_file (str, optional): Name of destination sharding strategy file which saved by
                                  'mindspore.set_auto_parallel_context(strategy_ckpt_save_file)'.
                                  when the 'dst_strategy_file' is None, it means that the destination sharding strategy
                                  is without any sharing for each parameter. Default:None.
+        process_num (int, optional): Number of processes to use for parallel processing. Defaults: 1.
+        output_format (str, optional): Control the format of the output checkpoint after conversion.
+            It can be set to either "ckpt" or "safetensors". Default: "ckpt".
 
     Raises:
         ValueError: `src_strategy_file` or `dst_strategy_file` is incorrect.
@@ -393,6 +400,21 @@ def transform_checkpoints(src_checkpoints_dir, dst_checkpoints_dir, ckpt_prefix,
         ...                       "./src_strategy.ckpt", "./dst_strategy.ckpt")
 
     """
+    all_safetensor_files_map = _collect_safetensor_files(src_checkpoints_dir)
+    all_ckpt_files_map = _collect_safetensor_files(src_checkpoints_dir, format='ckpt')
+    if all_safetensor_files_map and all_ckpt_files_map:
+        raise ValueError("For 'transform_checkpoints', the 'src_checkpoints_dir' cannot contain "
+                         "both ckpt file and safetensors file simultaneously")
+    if all_safetensor_files_map and not all_ckpt_files_map:
+        _transform_safetensors(src_checkpoints_dir, dst_checkpoints_dir, ckpt_prefix, src_strategy_file,
+                               dst_strategy_file, process_num, output_format)
+        return
+    if not all_safetensor_files_map and not all_ckpt_files_map:
+        raise ValueError("For 'transform_checkpoints', the 'src_checkpoints_dir' can not be empty.")
+    if all_ckpt_files_map and not all_safetensor_files_map and output_format == 'safetensors':
+        raise ValueError("For 'transform_checkpoints', 'output_format' can not be 'safetensors' "
+                         "when 'src_checkpoints_dir' only contains ckpt file.")
+
     if not os.path.isdir(src_checkpoints_dir):
         raise NotADirectoryError("src_checkpoints_dir {} is not a directory.".format(src_checkpoints_dir))
     _make_dir(dst_checkpoints_dir, "path")
@@ -419,7 +441,7 @@ def transform_checkpoints(src_checkpoints_dir, dst_checkpoints_dir, ckpt_prefix,
     layout_is_passed = src_layout_map and dst_layout_map
 
     if layout_is_passed and pipeline_stage_num == 1 and dst_stage_num == 1 and \
-        src_param_keys.issubset(dst_param_keys) and len(src_param_keys) < len(dst_param_keys):
+            src_param_keys.issubset(dst_param_keys) and len(src_param_keys) < len(dst_param_keys):
         ms.log.info("Transform checkpoint by every pipeline stage.")
         _transform_checkpoint_by_stage(src_checkpoints_dir, dst_checkpoints_dir, ckpt_prefix,
                                        src_strategy_file, dst_strategy_file)
@@ -442,31 +464,13 @@ def _sync_params(name, param, layout):
     is_send = layout[9]
     peer_rank = layout[10]
     sr_tag = layout[11]
-
-    class SharedParameterSyncCell(ms.nn.Cell):
-        """synchronize cell"""
-        def __init__(self, param, is_send, peer_rank, sr_tag):
-            super().__init__()
-            self.param = param
-            self.is_send = is_send
-            self.ret = ms.Tensor([0])
-
-            from mindspore.ops import Send, Receive
-            if self.is_send:
-                self.send = Send(sr_tag=sr_tag, dest_rank=peer_rank)
-            else:
-                self.receive = Receive(sr_tag=sr_tag, src_rank=peer_rank, shape=param.shape, dtype=param.dtype)
-
-        def construct(self):
-            if self.is_send:
-                out = self.send(self.param)
-                return ms.ops.functional.depend(self.ret, out)
-
-            self.param = self.receive(self.ret)
-            return ms.ops.functional.depend(self.ret, self.param)
-
-    sync_net = SharedParameterSyncCell(param, is_send, peer_rank, sr_tag)
-    sync_net()
+    if is_send:
+        ms.ops.Send(sr_tag=sr_tag, dest_rank=peer_rank)(param)
+    else:
+        param.assign_value(ms.ops.Receive(sr_tag=sr_tag,
+                                          src_rank=peer_rank,
+                                          shape=param.shape,
+                                          dtype=param.dtype)(param))
 
 
 def sync_pipeline_shared_parameters(net):
@@ -489,7 +493,7 @@ def sync_pipeline_shared_parameters(net):
             Before running the following examples, you need to configure the communication environment variables.
 
             For the Ascend device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
-            Startup <https://www.mindspore.cn/tutorials/experts/en/master/parallel/dynamic_cluster.html>`_ .
+            Startup <https://www.mindspore.cn/docs/en/master/model_train/parallel/dynamic_cluster.html>`_ .
 
         >>> import numpy as np
         >>> import mindspore as ms
@@ -562,6 +566,9 @@ def sync_pipeline_shared_parameters(net):
                "but got {}.".format(type(net)))
         raise TypeError(msg)
 
+    if _get_pipeline_stages() < 2:
+        return
+
     layout_dict = net.parameter_layout_dict
     if _is_in_auto_parallel_mode() and not layout_dict:
         from mindspore.common.api import _get_parameter_layout

mindspore/parallel/cluster/process_entity/__init__.py

@@ -15,4 +15,4 @@
 """Interfaces for ms_run"""
 from ._api import _Node, _MetaServerNode, _ComputeGraphNode, _ProcessManager
 
-from ._utils import _generate_cmd, _generate_url, _is_local_ip, _send_scale_num, _get_status_and_params
+from ._utils import _generate_cmd, _generate_url, _is_local_ip, _send_scale_num

mindspore/parallel/cluster/process_entity/_api.py

@@ -19,7 +19,7 @@ import sys
 import subprocess
 import mindspore.log as logger
 from ._utils import _generate_cmd_args_list, _generate_cmd_args_list_with_core, _generate_url,\
-                    _is_local_ip, _send_scale_num, _get_status_and_params
+                    _is_local_ip, _send_scale_num
 
 class _Node:
     """
@@ -212,6 +212,7 @@ class _ProcessManager:
             raise ValueError(f"Simulation level is set, worker_num must be 1, but got {self.worker_num}.")
 
         for i in range(self.local_worker_num):
+            os.environ["DEVICE_ID"] = str(i)
             node_id, log_name = self._get_node_id_and_log_path(i)
             if node_id is None:
                 logger.warning(f"Rank ids will be assigned automatically, "
@@ -241,19 +242,6 @@ class _ProcessManager:
             process = cgn.run()
             self.cgn_processes.append(process)
 
-    def heartbeat_with_scheduler(self):
-        """
-        Sends a heartbeat to the scheduler and updates the worker_num and local_worker_num.
-
-        Returns:
-            bool: True if the network has changed, False otherwise.
-
-        """
-        network_changed, worker_num, local_worker_num = _get_status_and_params(self.scheduler_url)
-        self.worker_num = worker_num
-        self.local_worker_num = local_worker_num
-        return network_changed
-
     def join_processes(self):
         """
         Join all processes to stop.
@@ -261,11 +249,31 @@ class _ProcessManager:
         so that understandable root cause of exception could be returned.
         """
         has_exception = False
-        for p in self.cgn_processes:
-            p.wait()
-            if p.returncode != 0:
-                has_exception = True
-                logger.error(f"Worker process {p.pid} exit with exception.")
+        success_cgn_processes = set()
+        while True:
+            # Traversal all workers and kill immediately if any exception happens.
+            for p in self.cgn_processes:
+                ret_code = p.poll()
+                if ret_code is None:
+                    # This means the process is still running, poll next process.
+                    continue
+                elif ret_code != 0:
+                    has_exception = True
+                    logger.error(f"Worker process {p.pid} exit with exception.")
+                    break
+                else:
+                    success_cgn_processes.add(p)
+
+            if has_exception:
+                logger.warning("There's worker exits with exception, kill all other workers.")
+                for p in self.cgn_processes:
+                    if p.poll() is None:
+                        p.kill()
+                break
+            elif len(success_cgn_processes) == len(self.cgn_processes):
+                logger.info("All workers successfully exit!")
+                break
+
 
         if self.msn_process:
             self.msn_process.wait()
@@ -335,10 +343,10 @@ class _ProcessManager:
                 time_out_node_log = re.findall(r"node: .* is timed out", scheduler_log)
 
                 # Filter out node ids of the processes which exit abnormally.
-                def node_id_splitter(id):
-                    return re.split(" is timed out", re.split("node: ", id)[1])[0]
-                for id in time_out_node_log:
-                    time_out_node_ids.append(node_id_splitter(id))
+                def node_id_splitter(node_id):
+                    return re.split(" is timed out", re.split("node: ", node_id)[1])[0]
+                for node_id in time_out_node_log:
+                    time_out_node_ids.append(node_id_splitter(node_id))
             logger.error(f"Time out nodes are {time_out_node_ids}")
 
         os.system(f"grep -rn -E 'ERROR|CRITICAL|Traceback|Error' -C 5 {self.log_dir}")

mindspore/parallel/cluster/process_entity/_utils.py

@@ -16,7 +16,6 @@
 import os
 import json
 import socket
-import requests
 import mindspore.log as logger
 
 def _generate_cmd(cmd, cmd_args, output_name):
@@ -25,7 +24,7 @@ def _generate_cmd(cmd, cmd_args, output_name):
     edirecting the output to a log file.
 
     """
-    if cmd not in ['python', 'pytest']:
+    if cmd not in ['python', 'pytest', 'python3']:
         # If user don't set binary file name, defaulty use 'python' to launch the job.
         command = f"python {cmd} {' '.join(cmd_args)} > {output_name} 2>&1 &"
     else:
@@ -99,28 +98,4 @@ def _send_scale_num(url, scale_num):
     Send an HTTP request to a specified URL, informing scale_num.
 
     """
-    try:
-        response = requests.post(url, data={"scale_num": scale_num}, timeout=100)
-        response.raise_for_status()
-        response_data = response.json()
-        response_bool = bool(response_data)
-        return response_bool
-    except requests.exceptions.RequestException:
-        return None
-
-
-def _get_status_and_params(url):
-    """
-    Send an HTTP request to a specified URL to query status and retrieve partial parameters.
-
-    """
-    try:
-        response = requests.get(url, timeout=100)
-        response.raise_for_status()
-        response_data = response.json()
-        network_changed = response_data.get("network_changed")
-        worker_num = response_data.get("worker_num")
-        local_worker_num = response_data.get("local_worker_num")
-        return network_changed, worker_num, local_worker_num
-    except requests.exceptions.RequestException:
-        return None
+    return ""

mindspore/parallel/parameter_broadcast.py

@@ -18,6 +18,8 @@ from __future__ import absolute_import
 __all__ = ["parameter_broadcast"]
 
 import numpy as np
+import mindspore as ms
+from mindspore.communication import get_rank, create_group, get_group_size
 
 
 def parameter_broadcast(net, layout, cur_rank=0, initial_rank=0):
@@ -104,9 +106,6 @@ def parameter_broadcast(net, layout, cur_rank=0, initial_rank=0):
     """
     if not layout:
         return
-    import mindspore as ms
-    from mindspore import Tensor
-    from mindspore.communication import get_rank, create_group, get_group_size
     from mindspore.train._utils import get_parameter_redundancy, remove_param_redundancy
     from mindspore.nn.wrap.cell_wrapper import AllreduceGraph
     origin_parallel_mode = ms.get_auto_parallel_context("parallel_mode")
@@ -143,7 +142,7 @@ def parameter_broadcast(net, layout, cur_rank=0, initial_rank=0):
                 raise ValueError(f"For parameter broadcast, the param: {param} can not be found.")
             real_param = net_param_dict[param]
             if param not in single_params[cur_rank]:
-                real_param.set_data(Tensor(np.zeros(real_param.shape), dtype=real_param.dtype))
+                real_param.set_data(ms.Tensor(np.zeros(real_param.shape), dtype=real_param.dtype))
             allreduce_input.append(real_param)
         if not allreduce_input:
             continue