mindspore 2.6.0__cp311-cp311-win_amd64.whl → 2.7.0__cp311-cp311-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,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

@@ -42,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.
@@ -235,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__)

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,12 +26,10 @@ 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.
+    Slice MiniBatch into finer-grained MicroBatch for use in pipeline-parallel training,
+    and specify the segment info.
 
     Note:
         micro_size must be greater or equal to pipeline stages.
@@ -38,6 +38,8 @@ class PipelineCell(Cell):
         network (Cell): The target network to wrap.
         micro_size (int): MicroBatch size.
         stage_config (dict, optional): The stage configuration for each cell's execution in pipeline parallel.
+        segment_config (dict, optional): The segment configuration for each cell's execution in pipeline parallel.
+            Default ``None``.
 
     Supported Platforms:
         ``Ascend``
@@ -49,7 +51,7 @@ class PipelineCell(Cell):
         >>> net = LeNet5()
         >>> net = nn.PipelineCell(net, 4, stage_config={"cell_name_0": 0, "cell_name_1": 1})
     """
-    def __init__(self, network, micro_size, stage_config=None):
+    def __init__(self, network, micro_size, stage_config=None, segment_config=None):
         super(PipelineCell, self).__init__(auto_prefix=False)
         self.network = network
         self.micro_inputs = nn.CellList()
@@ -101,15 +103,46 @@ 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):
+        if segment_config is None:
+            return
+        self._config_segment(segment_config)
+
+
+    def _config_segment(self, segment_config):
+        """
+        Config segment num for cell.
+        """
+        config_dict = segment_config.copy()
+
+        for cell_name, cell in self.network.cells_and_names():
+            if cell_name in segment_config:
+                setattr(cell, "pipeline_segment", segment_config[cell_name])
+                del config_dict[cell_name]
+        if str(self.network) in segment_config:
+            setattr(self.network, "pipeline_segment", segment_config[str(self.network)])
+            del config_dict[str(self.network)]
+        # if there are any config elements left, print them
+        if config_dict:
+            for config_cell_name, config_segment_num in config_dict.items():
+                logger.error("pipeline_cell segment_config set pipeline_segment fail!")
+                logger.warning("config cell name:" + str(config_cell_name) +
+                               " config segment num:" + str(config_segment_num))
+            logger.warning("network:" + str(self.network))
+            logger.warning("cell name available:")
+            for cell_name, _ in self.network.cells_and_names():
+                logger.warning(cell_name)
+            raise KeyError("For 'PipelineCell', the argument 'segment_config' : {} is not "
+                           "found in 'network' : {}".format(config_dict, self.network))
+
+
+    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)
@@ -120,7 +153,8 @@ class PipelineCell(Cell):
 
 class Pipeline(PipelineCell):
     """
-    Specify the number of micro_batch for pipeline parallelism and the division rules for stage.
+    Specify the number of micro_batch for pipeline parallelism and the division rules for stage,
+    and specify the segment info.
 
     Note:
         micro_size must be greater or equal to pipeline stages.
@@ -129,6 +163,8 @@ class Pipeline(PipelineCell):
         network (Cell): The target network to wrap.
         micro_size (int): MicroBatch size.
         stage_config (dict, optional): Stage configuration for cell's execution in pipeline parallel. Default ``None``.
+        segment_config (dict, optional): The segment configuration for each cell's execution in pipeline parallel.
+            Default ``None``.
 
     Raises:
         TypeError: The type of `net` is not cell.
@@ -197,10 +233,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 +287,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)