mindspore 2.6.0rc1__cp310-cp310-win_amd64.whl → 2.7.0rc1__cp310-cp310-win_amd64.whl

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,7 +176,7 @@ 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)
         elif os.getenv("MS_SIMULATION_LEVEL"):
@@ -205,15 +212,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 +277,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 +330,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 +357,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 +443,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 +459,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 +517,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 +538,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,9 +115,10 @@ 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",

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
@@ -59,8 +58,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 +219,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 +235,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__)

mindspore/parallel/nn/parallel_cell_wrapper.py

@@ -17,6 +17,8 @@
 from __future__ import absolute_import
 from __future__ import division
 
+__all__ = ['PipelineCell', 'Pipeline', 'MicroBatchInterleaved', 'GradAccumulation']
+
 from mindspore import nn
 from mindspore.ops import operations as P
 from mindspore.nn.cell import Cell
@@ -24,9 +26,6 @@ from mindspore.nn.wrap.cell_wrapper import _MicroBatch
 from mindspore import log as logger
 
 
-__all__ = ['PipelineCell', 'Pipeline', 'MicroBatchInterleaved', 'GradAccumulation']
-
-
 class PipelineCell(Cell):
     """
     Slice MiniBatch into finer-grained MicroBatch for use in pipeline-parallel training.
@@ -101,15 +100,15 @@ class PipelineCell(Cell):
                                    " config stage num:" + str(config_stage_num))
                 logger.warning("network:" + str(self.network))
                 logger.warning("cell name available:")
-                for cell_name, cell in self.network.cells_and_names():
+                for cell_name, _ in self.network.cells_and_names():
                     logger.warning(cell_name)
                 raise KeyError("For 'PipelineCell', the argument 'stage_config' : {} is not "
                                "found in 'network' : {}".format(config_dict, network))
 
-    def construct(self, *inputs):
+    def construct(self, *args, **kwargs):
         ret = None
         for i in range(self.micro_size):
-            micro_input = self.micro_inputs[i](i, *inputs)
+            micro_input = self.micro_inputs[i](i, *args, **kwargs)
             output = self.network(*micro_input)
             if ret is not None:
                 ret = self.add_list[i](ret, output)
@@ -197,10 +196,10 @@ class MicroBatchInterleaved(Cell):
             self.interleave_inputs.append(interleave_data)
         self._get_attr_from_cell(network)
 
-    def construct(self, *inputs):
+    def construct(self, *args, **kwargs):
         output = 0.0
         for i in range(self.interleave_num):
-            interleave_input = self.interleave_inputs[i](i, *inputs)
+            interleave_input = self.interleave_inputs[i](i, *args, **kwargs)
             output = self.add(output, self.network(*interleave_input))
         return output
 
@@ -251,10 +250,10 @@ class GradAccumulation(Cell):
             self.add_list.append(self.add)
         self._get_attr_from_cell(network)
 
-    def construct(self, *inputs):
+    def construct(self, *args, **kwargs):
         ret = None
         for i in range(self.micro_size):
-            micro_input = self.micro_inputs[i](i, *inputs)
+            micro_input = self.micro_inputs[i](i, *args, **kwargs)
             output = self.network(*micro_input)
             if ret is not None:
                 ret = self.add_list[i](ret, output)

mindspore/parallel/nn/parallel_grad_reducer.py

@@ -15,6 +15,8 @@
 """parallel serialization"""
 from __future__ import absolute_import
 
+__all__ = ['PipelineGradReducer']
+
 from mindspore import context
 from mindspore.nn.cell import Cell
 from mindspore.ops import functional as F, composite as C, operations as P
@@ -25,8 +27,6 @@ from mindspore.common.parameter import Parameter
 from mindspore.nn.layer import Identity
 from mindspore.parallel._utils import _get_enable_parallel_optimizer
 
-__all__ = ['PipelineGradReducer']
-
 
 grad_scale = C.MultitypeFuncGraph("grad_scale")
 shard_grad_scale = C.MultitypeFuncGraph("shard_grad_scale")
@@ -81,6 +81,7 @@ class PipelineGradReducer(Cell):
         >>> import mindspore as ms
         >>> from mindspore import nn, ops, Tensor
         >>> from mindspore.communication import init
+        >>> from mindspore.parallel.auto_parallel import AutoParallel
         >>>
         >>> ms.set_context(mode=ms.GRAPH_MODE)
         >>> ms.reset_auto_parallel_context()
@@ -113,7 +114,7 @@ class PipelineGradReducer(Cell):
         >>> net.layer3.pipeline_stage = 1
         >>> loss_fn = nn.CrossEntropyLoss()
         >>> optimizer = nn.SGD(net.trainable_params(), 1e-2)
-        >>> net_with_loss = nn.Pipeline(nn.WithLossCell(net, loss_fn), 2)
+        >>> net_with_loss = nn.PipelineCell(nn.WithLossCell(net, loss_fn), 2)
         >>> net_with_loss.set_train()
         >>> def forward_fn(inputs, target):
         ...     loss = net_with_loss(inputs, target)
@@ -135,7 +136,7 @@ class PipelineGradReducer(Cell):
         >>> label = Tensor(np.ones([size, out_features]).astype(np.float32))
         >>> loss, _ = train_one_step(inputs, label)
         >>> print(loss)
-        46.36721
+        46.304886
     """
     def __init__(self, parameters, scale_sense=1.0, opt_shard=None):
         super(PipelineGradReducer, self).__init__(auto_prefix=False)
@@ -151,10 +152,10 @@ class PipelineGradReducer(Cell):
             self.opt_shard = opt_shard
 
     @jit
-    def construct(self, grads):
+    def construct(self, *args, **kwargs):
         new_grads = None
         if self.opt_shard:
-            grads = self.grad_reducer(grads)
+            grads = self.grad_reducer(*args, **kwargs)
             new_grads = self.hyper_map(F.partial(shard_grad_scale, self.scale_sense * self.degree),
                                        grads, self.accu_grads)
         else:

mindspore/parallel/shard.py

@@ -121,8 +121,7 @@ class Layout:
         >>> layout0 = layout("dp", "mp")
         >>> print(layout0.to_dict())
         {"device_matrix": (2, 2, 2), "tensor_map": (2, 0), "interleaved_parallel": False,
-        'alias_name': {'dp', 'sp', 'mp'}, "rank_list": [0, 1, 2, 3, 4, 1, 6, 7]}
-        >>> # Total device num is 4, but split the tensor in local device into two copies.
+        'alias_name': {'dp', 'sp', 'mp'}, "rank_list": [0, 1, 2, 3, 4, 5, 6, 7]}
         >>> layout = Layout((2, 2, 2), ("dp", "sp", "interleaved_parallel"))
         >>> layout1 = layout(("dp", "interleaved_parallel"), "sp")
     """
@@ -408,8 +407,8 @@ def shard(fn, in_strategy, out_strategy=None, parameter_plan=None, device="Ascen
     The parallel strategies of remaining operators are derived from the strategy specified by the input and output.
 
     Note:
-        - If shard is called, the parallel mode in `set_auto_parallel_context` (parallel_mode) will be set to
-          "auto_parallel" and the search mode (search_mode) to "sharding_propagation".
+        - It is valid only in semi auto parallel or auto parallel mode.
+          In other parallel modes, strategies set here will be ignored.
         - If the input contain Parameter, its strategy should be set in `in_strategy`.
         - This method currently does not support dynamic shapes.