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

mindspore/parallel/checkpoint_transform.py

@@ -582,6 +582,8 @@ def transform_checkpoints(src_checkpoints_dir, dst_checkpoints_dir, ckpt_prefix,
         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.
 
+        This function does not support converting remove_redundancy's checkpoint file.
+
     Args:
         src_checkpoints_dir (str): The source checkpoints directory.
         dst_checkpoints_dir (str): The destination checkpoints directory to save the converted checkpoints.
@@ -924,8 +926,8 @@ def set_op_strategy_config(mode="SAVE", path=""):
 
 def build_searched_strategy(strategy_filename):
     """
-    Extract the sharding strategy for each parameter in the network
-    from the strategy file for distributed inference scenarios.
+    Extract the sharding strategy for each parameter in the network from the strategy file
+    for distributed inference scenarios.
 
     Args:
         strategy_filename (str): Name of strategy file.
@@ -1025,8 +1027,10 @@ def load_distributed_checkpoint(network, checkpoint_filenames=None, predict_stra
         >>> from mindspore.parallel.auto_parallel import AutoParallel
         >>> from mindspore.nn.utils import no_init_parameters
         >>> from mindspore.common.initializer import initializer, One
+        >>> from mindspore.communication.management import get_group_size
         >>>
         >>> step_per_epoch = 4
+        >>> device_num = get_group_size()
         >>>
         >>> # Define the network structure.
         >>> class Net(nn.Cell):
@@ -1070,7 +1074,7 @@ def load_distributed_checkpoint(network, checkpoint_filenames=None, predict_stra
         ...     network = AutoParallel(network, parallel_mode="semi_auto")
         ...     network.save_param_strategy_file(file_path="./train_strategy.ckpt")
         ...     model = ms.Model(network=network, loss_fn=net_loss, optimizer=net_opt)
-        ...     ckpt_config = train.CheckpointConfig(keep_checkpoint_max=1, integrated_save=False)
+        ...     ckpt_config = train.CheckpointConfig(keep_checkpoint_max=1, integrated_save=True)
         ...     global_rank_id = int(os.getenv("RANK_ID"))
         ...     ckpt_path = "./rank_{}_ckpt".format(global_rank_id)
         ...     ckpt_callback = train.ModelCheckpoint(prefix="parallel", directory=ckpt_path, config=ckpt_config)
@@ -1096,10 +1100,10 @@ def load_distributed_checkpoint(network, checkpoint_filenames=None, predict_stra
         >>>
         >>> train_net()
         >>> load_model()
-        [[-7.3259363 -7.497216  -7.398196  ... -7.374962  -7.204874  -7.234935 ]
-        [ 3.362938   3.3535435  3.3832688 ...  3.4263954  3.279045   3.3202887]
+        [[-9.62929535e+00, -9.76258755e+00, -9.70192051e+00 ... -9.67151260e+00, -9.71998310e+00, -9.64571190e+00],
+        [-4.63218540e-01, -4.07317460e-01, -3.78161550e-01 ... -3.95918339e-01, -2.87363172e-01, -3.48693460e-01],
         ...
-        [ 1.6067538  1.6244187  1.5384722 ...  1.5449994  1.6195512  1.6176052]]
+        [-4.28075647e+00, -4.36630344e+00, -4.25664043e+00 ... -4.32012939e+00, -4.30337954e+00, -4.27571440e+00]]
     """
     if format not in ['safetensors', 'ckpt'] or output_format not in ['safetensors', 'ckpt']:
         raise ValueError(
@@ -1161,6 +1165,8 @@ def load_distributed_checkpoint(network, checkpoint_filenames=None, predict_stra
             train_strategy_filename = ms.context.get_auto_parallel_context("strategy_ckpt_load_file")
 
     _train_strategy = build_searched_strategy(train_strategy_filename)
+    if not _train_strategy:
+        return True
     train_strategy = _convert_to_list(_train_strategy)
 
     train_dev_count = 1
@@ -1185,7 +1191,7 @@ def load_distributed_checkpoint(network, checkpoint_filenames=None, predict_stra
     param_not_in_ckpt = []
     for _, param in network.parameters_and_names():
         sliced_params = []
-        if param.name not in rank_list.keys():
+        if param.name not in rank_list:
             param_not_in_strategy.append(param.name)
             continue
         if param.name not in param_total_dict:

mindspore/parallel/cluster/process_entity/_api.py

@@ -22,7 +22,8 @@ import socket
 import psutil
 import mindspore.log as logger
 from ._utils import _generate_cmd_args_list, _generate_cmd_args_list_with_core, _generate_url, \
-    _is_local_ip, _convert_addr_to_ip, _send_scale_num, _get_local_ip
+    _is_local_ip, _convert_addr_to_ip, _send_scale_num, _get_local_ip, _generate_auto_bind_core_strategy, \
+    _generate_bind_core_strategy
 
 
 class _Node:
@@ -79,11 +80,12 @@ class _ComputeGraphNode(_Node):
     Worker node for dynamic networking. Inherits from the Node class.
     """
 
-    def __init__(self, worker_num, sched_host, sched_port, timeout, node_id, args_list, output_file,
+    def __init__(self, worker_num, sched_host, sched_port, timeout, node_id, node_rank, args_list, output_file,
                  tail_worker_log, join, is_simulation):
         super().__init__(worker_num, sched_host, sched_port, timeout, args_list, output_file,
                          tail_worker_log, join, is_simulation)
         self.node_id = node_id
+        self.node_rank = node_rank
 
     def run(self):
         """
@@ -95,6 +97,8 @@ class _ComputeGraphNode(_Node):
         super().run()
         if self.node_id is not None:
             os.environ["MS_NODE_ID"] = str(self.node_id)
+        if self.node_rank is not None:
+            os.environ["MS_NODE_RANK"] = str(self.node_rank)
         # If simulation level is set, environment variable 'MS_ROLE' will not be set.
         if not self.is_simulation:
             os.environ["MS_ROLE"] = "MS_WORKER"
@@ -119,6 +123,9 @@ class _ComputeGraphNode(_Node):
         return subprocess.Popen(['/usr/bin/tail', '-f', self.output_file])
 
     def enable_tail_worker_log(self):
+        """
+        Get valid rank ID for tailing the corresponding worker log.
+        """
         tail_worker_log_list = []
         if self.tail_worker_log != "-1":
             tail_worker_log_list.extend([int(num) for num in self.tail_worker_log.split(',')])
@@ -169,12 +176,15 @@ class _ProcessManager:
 
         self.sim_level = args.sim_level
         self.sim_rank_id = args.sim_rank_id
-        self.is_simulation = (self.sim_level != -1)
+        self.is_simulation = self.sim_level != -1
         if self.is_simulation:
             os.environ["MS_SIMULATION_LEVEL"] = str(self.sim_level)
+            if self.sim_rank_id == -1:
+                self.sim_rank_id = int(os.getenv("RANK_ID", "-1"))
         elif os.getenv("MS_SIMULATION_LEVEL"):
             self.is_simulation = True
-            self.sim_rank_id = int(os.getenv("RANK_ID", "-1"))
+            if self.sim_rank_id == -1:
+                self.sim_rank_id = int(os.getenv("RANK_ID", "-1"))
             if os.getenv("RANK_SIZE"):
                 self.exported_rank_size = os.getenv("RANK_SIZE")
         # If sim_rank_id is set, single worker can be started.
@@ -205,15 +215,24 @@ class _ProcessManager:
             finally:
                 os.umask(origin_mask)
 
+        self.device_to_cpu_map = {}
+        if self.bind_core is True:
+            self.device_to_cpu_map = _generate_auto_bind_core_strategy(self.local_worker_num)
+
         self.proc_rank_map = {}
         self.enable_mindx = False
+        self._check_taskd()
+
+    def _check_taskd(self):
+        """check if enable taskd."""
         tft_env = os.getenv("MS_ENABLE_TFT", "")
-        if ("TTP:1" in tft_env) or ("UCE:1" in tft_env) or ("ARF:1" in tft_env):
+        if any(v in tft_env for v in ('TTP:1', 'UCE:1', 'ARF:1', 'TSP:1', 'RSC:1', 'HCCE:1')):
             try:
                 from taskd.python.framework.agent.ms_mgr.msrun_plugin import MSRunPlugin
                 self.msmgr = MSRunPlugin()
                 self.msmgr.register_callbacks("KILL_WORKER", self.kill_workers)
                 self.msmgr.register_callbacks("START_ALL_WORKER", self.start_all_workers)
+                self.msmgr.register_callbacks("START_WORKER_LIST", self.start_worker_list)
                 self.msmgr.register_callbacks("MONITOR", self.monitor_rank_status)
                 self.enable_mindx = True
                 os.environ["MS_ENABLE_RECOVERY"] = str(1)
@@ -261,6 +280,45 @@ class _ProcessManager:
                               self.is_simulation)
         self.msn_process = msn.run()
 
+    def _start_single_worker(self, local_rank):
+        """
+        Start worker processor
+
+        Args:
+            local_rank: local rank id.
+        """
+        os.environ["DEVICE_ID"] = str(local_rank)
+        node_id, log_name = self._get_node_id_and_log_path(local_rank)
+        if node_id is None:
+            logger.warning(f"Rank ids will be assigned automatically, "
+                           "please use 'grep -rn 'rank id:' command to check each worker log's rank id.")
+        else:
+            # If node_id is generated in '_get_node_id_and_log_path' method, export 'RANK_ID' environment variable.
+            # This is for rank_table method's compatibility consideration.
+            os.environ["RANK_ID"] = str(node_id)
+            print(f"Start worker process with rank id:{node_id}, log file:{log_name}. "
+                  f"Environment variable [RANK_ID={node_id}] is exported.", flush=True)
+        if self.is_simulation and (self.sim_rank_id != -1):
+            # Reset RANK_ID env to sim_rank_id if sim_rank_id is set.
+            os.environ["RANK_ID"] = str(self.sim_rank_id)
+            logger.warning(f"In dryrun case, RANK_ID is assigned to {self.sim_rank_id}.")
+
+        if self.bind_core:
+            affinity_cpu_str = _generate_bind_core_strategy(local_rank, self.device_to_cpu_map, self.bind_core)
+            if affinity_cpu_str is not None:
+                cmd = _generate_cmd_args_list_with_core(self.cmd, self.cmd_args, affinity_cpu_str)
+            else:
+                cmd = _generate_cmd_args_list(self.cmd, self.cmd_args)
+        else:
+            cmd = _generate_cmd_args_list(self.cmd, self.cmd_args)
+        cgn = _ComputeGraphNode(self.worker_num, self.master_addr, self.master_port, self.cluster_time_out,
+                                node_id, self.node_rank, cmd, log_name, self.tail_worker_log, self.join,
+                                self.is_simulation)
+        process, tail_process = cgn.run()
+        self.cgn_processes.append(process)
+        self.tail_cgn_processes.append(tail_process)
+        self.proc_rank_map[local_rank] = process
+
     def start_workers(self):
         """
         Starts the worker nodes.
@@ -275,40 +333,8 @@ class _ProcessManager:
                            "'rank_id' of each process will be assigned after cluster is successfully built.\n"
                            "You can access 'RANK_ID' environment variable after calling "
                            "'mindspore.communication.init()'")
-
         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, "
-                               "please use 'grep -rn 'rank id:' command to check each worker log's rank id.")
-            else:
-                # If node_id is generated in '_get_node_id_and_log_path' method, export 'RANK_ID' environment variable.
-                # This is for rank_table method's compatibility consideration.
-                os.environ["RANK_ID"] = str(node_id)
-                print(f"Start worker process with rank id:{node_id}, log file:{log_name}. "
-                      f"Environment variable [RANK_ID={node_id}] is exported.", flush=True)
-            if self.is_simulation and (self.sim_rank_id != -1):
-                # Reset RANK_ID env to sim_rank_id if sim_rank_id is set.
-                os.environ["RANK_ID"] = str(self.sim_rank_id)
-                logger.warning(f"In dryrun case, RANK_ID is assigned to {self.sim_rank_id}.")
-
-            if self.bind_core:
-                cpu_num = subprocess.getoutput("cat /proc/cpuinfo|grep processor|wc -l")
-                if not cpu_num.isdigit():
-                    raise RuntimeError(f"Got cpu number from '/proc/cpuinfo' is {cpu_num}, failed to bind core.")
-                avg = int(cpu_num) // self.local_worker_num
-                cpu_start = avg * i
-                cpu_end = cpu_start + avg - 1
-                cmd = _generate_cmd_args_list_with_core(self.cmd, self.cmd_args, cpu_start, cpu_end)
-            else:
-                cmd = _generate_cmd_args_list(self.cmd, self.cmd_args)
-            cgn = _ComputeGraphNode(self.worker_num, self.master_addr, self.master_port, self.cluster_time_out,
-                                    node_id, cmd, log_name, self.tail_worker_log, self.join, self.is_simulation)
-            process, tail_process = cgn.run()
-            self.cgn_processes.append(process)
-            self.tail_cgn_processes.append(tail_process)
-            self.proc_rank_map[i] = process
+            self._start_single_worker(i)
 
     def join_processes(self):
         """
@@ -334,7 +360,7 @@ class _ProcessManager:
                     continue
                 elif ret_code != 0:
                     has_exception = True
-                    logger.error(f"Worker process {p.pid} exit with exception.")
+                    logger.error(f"Worker process {p.pid} exit with exception. Error code: {ret_code}.")
                     break
                 else:
                     success_cgn_processes.add(p)
@@ -420,14 +446,9 @@ class _ProcessManager:
         Args:
             NA.
         """
-        for p in self.cgn_processes:
-            if p.poll() is None:
-                p.kill()
+        self.kill_worker_processes()
+        self.kill_tail_log_processes()
         self.cgn_processes.clear()
-
-        for p in self.tail_cgn_processes:
-            if p is not None:
-                p.kill()
         self.tail_cgn_processes.clear()
 
     def kill_single_worker(self, pid):
@@ -441,7 +462,7 @@ class _ProcessManager:
         for i in range(len(self.cgn_processes)):
             p = self.cgn_processes[i]
             if p.pid == pid and p.poll() is None:
-                p.kill()
+                os.killpg(os.getpgid(p.pid), signal.SIGKILL)
                 del self.cgn_processes[i]
                 tail_p = self.tail_cgn_processes[i]
                 if tail_p is not None:
@@ -499,7 +520,8 @@ class _ProcessManager:
                 p_status = p.poll()
                 if (not psutil.pid_exists(p.pid)) and (p_status != 0):
                     p_status = 300
-                return {"pid": p.pid, "status": p_status, "global_rank": global_rank_id}
+                return {"pid": p.pid, "status": p_status, "global_rank": global_rank_id, "local_rank": rank_id,
+                        "node_id": self.node_rank}
             except KeyError:
                 logger.info(f"Process rank {rank_id} has not been initialized.")
                 return {"pid": None, "status": 200, "global_rank": global_rank_id}
@@ -519,7 +541,24 @@ class _ProcessManager:
         self.start_workers()
         worker_status = self.monitor_rank_status([-1])
         for i in range(self.local_worker_num):
-            if worker_status[i]["status"] != None:  # pylint: disable=singleton-comparison
+            if worker_status[i]["status"] is not None:
+                return 1
+        return 0
+
+    def start_worker_list(self, rank_ids):
+        """
+        Start worker processor by rank list.
+
+        Args:
+            rank_ids: worker process's local rank list, which is also device_id.
+        """
+        if not isinstance(rank_ids, list):
+            raise TypeError(f"The type of 'rank_ids' must be a list, but got:{rank_ids}")
+        for idx in rank_ids:
+            self._start_single_worker(idx)
+        worker_status = self.monitor_rank_status(rank_ids)
+        for i in rank_ids:
+            if worker_status[i]["status"] is not None:
                 return 1
         return 0
 

mindspore/parallel/cluster/process_entity/_utils.py

@@ -18,6 +18,8 @@ import json
 import socket
 import ipaddress
 import mindspore.log as logger
+from mindspore.runtime.thread_bind_core import _get_physical_device_id, _get_cpu_available, \
+    _auto_generate_strategy, _equal_distribution_strategy
 
 CURRENT_IP = None
 
@@ -45,19 +47,19 @@ def _generate_cmd_args_list(cmd, cmd_args):
     return [cmd] + cmd_args
 
 
-def _generate_cmd_args_list_with_core(cmd, cmd_args, cpu_start, cpu_end):
+def _generate_cmd_args_list_with_core(cmd, cmd_args, affinity_cpu_str):
     """
     Generates arguments list for 'Popen'. It consists of a binary file name and subsequential arguments.
     """
     # Bind cpu cores to this process.
-    taskset_args = ['taskset'] + ['-c'] + [str(cpu_start) + '-' + str(cpu_end)]
+    taskset_args = ['taskset'] + ['-c'] + [affinity_cpu_str]
     final_cmd = []
     if cmd not in ['python', 'pytest', 'python3']:
         # If user don't set binary file name, defaulty use 'python' to launch the job.
         final_cmd = taskset_args + ['python'] + [cmd] + cmd_args
     else:
         final_cmd = taskset_args + [cmd] + cmd_args
-    logger.info(f"Launch process with command: {' '.join(final_cmd)}")
+    logger.warning(f"Launch process with command: {' '.join(final_cmd)}")
     return final_cmd
 
 
@@ -83,8 +85,8 @@ def _get_local_ip(ip_address):
             CURRENT_IP = s.getsockname()[0]
             s.close()
         except Exception as e:
-            raise RuntimeError(f"Get local ip failed: {e}. Please check whether an accessible address "
-                               "is input by '--master_address'.")
+            raise RuntimeError("Get local ip has failed. Please verify that the accessible address has been "
+                               "specified in the '--master_address' parameter") from e
     return CURRENT_IP
 
 
@@ -124,8 +126,8 @@ def _convert_addr_to_ip(master_addr):
             logger.info(f"Convert input host name:{master_addr} to ip address:{ip_address}.")
             return ip_address
         except socket.gaierror as e:
-            raise RuntimeError(f"DNS resolution failed: {e}. Please check whether a correct host name "
-                               "is input by '--master_address'.")
+            raise RuntimeError("DNS resolution has failed. Please verify that the correct hostname has been "
+                               "specified in the '--master_address' parameter") from e
 
 
 def _send_scale_num(url, scale_num):
@@ -134,3 +136,89 @@ def _send_scale_num(url, scale_num):
 
     """
     return ""
+
+
+def _parse_global_device_to_cpu_map(local_rank_id, physical_device_id, device_to_cpu_map):
+    """
+    Parse the global device_to_cpu_map and return a cpu list for assigned local_rank_id.
+
+    """
+    input_device_id = int(list(device_to_cpu_map.keys())[local_rank_id].replace("device", ""))
+    if physical_device_id != input_device_id:
+        return ""
+    affinity_cpu_list = list(device_to_cpu_map.values())[local_rank_id]
+    affinity_cpu_str = ",".join(affinity_cpu_list)
+    return affinity_cpu_str
+
+
+def _generate_auto_bind_core_strategy(local_worker_num):
+    """
+    Get device to core range assigned for the all processes.
+
+    """
+    simulation_level = os.getenv("MS_SIMULATION_LEVEL", "").strip()
+
+    try:
+        available_cpus = _get_cpu_available()
+    except RuntimeError as e:
+        logger.warning(f"Failed to acquire available cpu info, error: {e} Will not launch process with taskset.")
+        return {}
+
+    if not simulation_level:
+        device_to_cpu_map = _auto_generate_strategy(local_worker_num, available_cpus)
+    else:
+        device_to_cpu_map = _equal_distribution_strategy(local_worker_num, available_cpus)
+
+    return device_to_cpu_map
+
+
+def ranges_to_str(num_list):
+    """
+    Convert a num list to a range string.
+
+    """
+    ranges = []
+    start = num_list[0]
+    for i in range(1, len(num_list)):
+        if num_list[i] != num_list[i-1] + 1:
+            ranges.append((start, num_list[i-1]))
+            start = num_list[i]
+    ranges.append((start, num_list[-1]))
+
+    parts = []
+    for start, end in ranges:
+        if start == end:
+            parts.append(str(start))
+        else:
+            parts.append(f"{start}-{end}")
+    return ",".join(parts)
+
+
+def _generate_bind_core_strategy(local_rank_id, device_to_cpu_map, arg_bind_core):
+    """
+    Get device to core range assigned for the all processes.
+
+    """
+    affinity_cpu_str = ""
+    cpu_list_for_device = []
+    simulation_level = os.getenv("MS_SIMULATION_LEVEL", "").strip()
+
+    try:
+        physical_device_id = _get_physical_device_id(local_rank_id, simulation_level)
+    except RuntimeError as e:
+        logger.warning(f"Failed to acquire device id, error: {e} Will not launch process with taskset.")
+        return None
+
+    if isinstance(arg_bind_core, dict):
+        affinity_cpu_str = _parse_global_device_to_cpu_map(local_rank_id, physical_device_id, arg_bind_core)
+        if not affinity_cpu_str:
+            logger.warning(f"Failed to find physical_device_id[{physical_device_id}] for "
+                           f"process[{local_rank_id}]. Will not launch process with taskset.")
+            return None
+    elif arg_bind_core is True:
+        cpu_list_for_device = device_to_cpu_map.get(physical_device_id, [])
+        if not cpu_list_for_device:
+            return None
+        os.environ["MSRUN_CPU_LIST"] = str(cpu_list_for_device)
+        affinity_cpu_str = ranges_to_str(cpu_list_for_device)
+    return affinity_cpu_str

mindspore/parallel/cluster/run.py

@@ -14,9 +14,47 @@
 # ============================================================================
 """Entrypoint of ms_run"""
 import ast
-from argparse import REMAINDER, ArgumentParser
+import re
+import json
+from argparse import REMAINDER, ArgumentParser, ArgumentTypeError
 from .process_entity import _ProcessManager
 
+
+def parse_and_validate_bind_core(value):
+    """
+    Parse input argument of --bind_core.
+
+    """
+    if value.lower() == "true":
+        return True
+    if value.lower() == "false":
+        return False
+
+    try:
+        value_dict = json.loads(value)
+    except json.JSONDecodeError as e:
+        raise ArgumentTypeError("Failed to parse JSON into a dictionary") from e
+
+    if isinstance(value_dict, dict):
+        range_pattern = re.compile(r'^\d+-\d+$')
+        for device_id, affinity_cpu_list in value_dict.items():
+            if not re.fullmatch(r"device\d+", device_id):
+                raise ArgumentTypeError(f"Key '{device_id}' must be in format 'deviceX' (X ≥ 0).")
+            if not isinstance(affinity_cpu_list, list):
+                raise ArgumentTypeError(f"Value for '{device_id}':{affinity_cpu_list} should be a list, "
+                                        f"but got {type(affinity_cpu_list)}.")
+
+            for cpu_range in affinity_cpu_list:
+                if not isinstance(cpu_range, str):
+                    raise ArgumentTypeError(f"CPU range '{cpu_range}' in '{affinity_cpu_list}' should be a string.")
+                if not range_pattern.match(cpu_range):
+                    raise ArgumentTypeError(f"CPU range '{cpu_range}' in '{affinity_cpu_list}' should be "
+                                            "in format 'cpuidX-cpuidY'.")
+        return value_dict
+
+    raise ArgumentTypeError(f"Type of {value} should be bool or dict, but got {type(value)}.")
+
+
 def get_args():
     """
     Parses and retrieves command-line arguments.
@@ -77,23 +115,26 @@ def get_args():
     parser.add_argument(
         "--bind_core",
         default=False,
-        type=ast.literal_eval,
-        choices=[True, False],
-        help="specifies whether msrun should bind cpu cores to spawned processes."
+        type=parse_and_validate_bind_core,
+        help="specifies whether msrun should bind CPU cores to spawned processes. "
+             "If set to True, msrun will bind core based on the environment automatically, "
+             "and if passed a dict, msrun will bind core based on this dict information."
     )
     parser.add_argument(
         "--sim_level",
         default=-1,
         type=int,
         choices=[0, 1, 2, 3],
-        help="specifies simulation level. When this argument is set, msrun only spawns one process "
-             "but export RANK_SIZE with value worker_num and RANK_ID with value sim_rank_id."
+        help="specifies simulation level. This argument activates dryrun mode, functioning "
+             "equivalently to environment variable 'MS_SIMULATION_LEVEL' while having higher priority."
     )
     parser.add_argument(
         "--sim_rank_id",
         default=-1,
         type=int,
-        help="specifies simulation process's rank id. Only one process is spawned in simulation scenario."
+        help="specifies simulation process's rank id. When this argument is set, only one process "
+             "is spawned on dryrun mode, functioning equivalently to environment variable 'RANK_ID' "
+             "while having higher priority."
     )
     parser.add_argument(
         "--rank_table_file",

mindspore/parallel/function/__init__.py

@@ -16,8 +16,15 @@
 """
 Parallel function operator
 """
+from __future__ import absolute_import
 
-from mindspore.parallel.function.reshard_func import reshard
+from . import (
+    reshard_func
+)
+
+from .reshard_func import (
+    reshard
+)
 
 __all__ = []
 __all__.extend(reshard_func.__all__)

mindspore/parallel/function/reshard_func.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ============================================================================
 """Defines parameter operators with functional form."""
-import mindspore as ms
 from mindspore import context, ops
 from mindspore import log as logger
 from mindspore.ops import operations as P
@@ -43,11 +42,12 @@ def reshard(tensor, layout):
     can check :class:`mindspore.parallel.Layout` for reference.
 
     Note:
-        - In the Graph mode, this function can set the sharding propagation strategy of a tensor.
-          For those tensor do not manually be set, their strategies are decided by the sharding
-          strategy propagation algorithm automatically.
-        - In PyNative mode, you can use this method to arrange tensors in a cell (that is, cells
-          that use Cell.shard/F.shard in PyNative mode) that is executed in parallel in graph mode.
+        In the Graph mode, this function can set the sharding propagation strategy of a tensor.
+        For those tensor do not manually be set, their strategies are decided by the sharding
+        strategy propagation algorithm automatically.
+
+    .. warning::
+        The method is currently not supported in PyNative mode.
 
     Args:
         tensor (Tensor): The tensor to be set the sharding strategy.
@@ -59,8 +59,8 @@ def reshard(tensor, layout):
         Tensor. The mathematically equivalent of the input tensor.
 
     Raises:
-        TypeError: Reshard takes in Tensor type as the first input param, but got: `type(tensor)`.
-        TypeError: Reshard only support type mindspore.parallel.Layout but got: `type(layout)`.
+        TypeError: If the type of input param `tensor` is not mindspore.Tensor.
+        TypeError: If the type of input param `layout` is not mindspore.parallel.Layout.
 
     Supported Platforms:
         ``Ascend``
@@ -220,11 +220,11 @@ def _redistribute(tensor, dst_dtensor_info):
             if not comm_tensor_data_func._current_rank_has_data:
                 new_tensor_shape = tuple([tensor_data.shape[i] // tensor._dtensor_info.sharding_strategy[i]
                                           for i in range(len(tensor.shape))])
-                tensor_data = comm_tensor_data_func.comm_data(ops.zeros(new_tensor_shape, tensor.dtype))
+                tensor_data = ops.zeros(new_tensor_shape, tensor.dtype)
+                _ = comm_tensor_data_func.comm_data(tensor_data)
             else:
-                tensor_data = comm_tensor_data_func.comm_data(tensor)
+                _ = comm_tensor_data_func.comm_data(tensor_data)
             all_reduce_data = True
-    ms.communication.comm_func.barrier()
     if src_layout_info['device_matrix'] == dst_layout_info['device_matrix'] and src_layout_info['tensor_map'] == \
             dst_layout_info['tensor_map']:
         return tensor_data
@@ -236,7 +236,7 @@ def _redistribute(tensor, dst_dtensor_info):
     global REDIST_CELL_CACHE
     redist_cache_key = (f"{src_layout_info['device_matrix']}, {src_layout_info['tensor_map']} -> "
                         f"{dst_layout_info['device_matrix']}, {dst_layout_info['tensor_map']}")
-    if redist_cache_key in REDIST_CELL_CACHE.keys():
+    if redist_cache_key in REDIST_CELL_CACHE:
         logger.debug(f"redist_cache_key is {redist_cache_key}, match cache")
         redist_func = REDIST_CELL_CACHE[redist_cache_key]
     else:

mindspore/parallel/nn/__init__.py

@@ -17,8 +17,21 @@ Interfaces for parallel-related functionality
 """
 from __future__ import absolute_import
 
-from mindspore.parallel.nn.parallel_grad_reducer import PipelineGradReducer
-from mindspore.parallel.nn.parallel_cell_wrapper import PipelineCell, Pipeline, MicroBatchInterleaved, GradAccumulation
+from . import (
+    parallel_grad_reducer,
+    parallel_cell_wrapper
+)
+
+from .parallel_grad_reducer import (
+    PipelineGradReducer
+)
+
+from .parallel_cell_wrapper import (
+    PipelineCell,
+    Pipeline,
+    MicroBatchInterleaved,
+    GradAccumulation
+)
 
 __all__ = []
 __all__.extend(parallel_grad_reducer.__all__)