class_docstr: PipelineStageManager is a helper class to manage pipeline stages.
  
  Args:
      pg_mesh (ProcessGroupMesh): Process group mesh.
      pipeline_axis (int): The axis along which the pipeline is constructed.
      is_virtual (bool): Whether to use circle p2p communication, it will make the first and last stage communicate with each other.
  
  Attributes:
      num_stages (int): Number of stages in the pipeline.
      stage (int): The current stage.
class_docstr_tok: ["PipelineStageManager","is","a","helper","class","to","manage","pipeline","stages",".","Args",":","pg_mesh","(","ProcessGroupMesh",")",":","Process","group","mesh",".","pipeline_axis","(","int",")",":","The","axis","along","which","the","pipeline","is","constructed",".","is_virtual","(","bool",")",":","Whether","to","use","circle","p2p","communication",",","it","will","make","the","first","and","last","stage","communicate","with","each","other",".","Attributes",":","num_stages","(","int",")",":","Number","of","stages","in","the","pipeline",".","stage","(","int",")",":","The","current","stage","."]
code: def __init__(
      self,
      pg_mesh: ProcessGroupMesh,
      pipeline_axis: int,
      enable_interleave: bool = False,
      num_model_chunks: int = 1,
  ) -> None:
      assert (
          enable_interleave or num_model_chunks == 1
      ), "num_model_chunks must be 1 when enable_interleave is False"
  
      self.pg_mesh = pg_mesh
      self.pipeline_axis = pipeline_axis
      self.prev_rank: Optional[Tuple[int, ...]] = None
      self.next_rank: Optional[Tuple[int, ...]] = None
      self.p2p_groups: Dict[Tuple[int, int], ProcessGroup] = {}
  
      # init prev and next coord
      coord = self.pg_mesh.coordinate()
      # the prev rank of rank0 is the last rank
      prev_coord = (
          coord[: self.pipeline_axis]
          + (coord[self.pipeline_axis] - 1,)
          + coord[self.pipeline_axis + 1 :]
      )
      self.prev_rank = self.pg_mesh.ravel(
          prev_coord, self.pg_mesh.shape, mode="wrap"
      )
      # the next rank of the last rank is rank0
      next_coord = (
          coord[: self.pipeline_axis]
          + (coord[self.pipeline_axis] + 1,)
          + coord[self.pipeline_axis + 1 :]
      )
      self.next_rank = self.pg_mesh.ravel(
          next_coord, self.pg_mesh.shape, mode="wrap"
      )
  
      # init p2p process groups
      stages = list(range(self.num_stages))
      for prev, cur in zip(stages[:-1], stages[1:]):
          group = self.pg_mesh.get_group_along_axis(
              self.pipeline_axis, [prev, cur]
          )
          if self.stage in [prev, cur]:
              ranks_in_group = self.pg_mesh.get_ranks_in_group(group)
              self.p2p_groups[tuple(ranks_in_group)] = group
  
      self.is_interleave = enable_interleave
      if enable_interleave:
          # use circle p2p communication
          # add the process group of the first rank and the last rank
          group = self.pg_mesh.get_group_along_axis(
              self.pipeline_axis, [stages[0], stages[-1]]
          )
          if self.stage in [stages[0], stages[-1]]:
              ranks_in_group = self.pg_mesh.get_ranks_in_group(group)
              self.p2p_groups[tuple(ranks_in_group)] = group
  
          # for interleaved pipeline parallel, each device is responsible for multiple chunk of layers
          self.num_model_chunks: int = num_model_chunks
  
          # for shardformer, hold stage indices of model
          self.stage_indices: List[Tuple[int, int]]
          # for shardformer, hold model chunk id
          self.model_chunk_id: Optional[int] = None
  
file_path: colossalai/colossalai/pipeline/stage_manager.py
imports: import contextlib
  from typing import Dict, List, Optional, Tuple
  import torch.distributed
  from torch.distributed import ProcessGroup
  from colossalai.cluster import ProcessGroupMesh
label_desc: Use image node_id 1 to create a new PipelineStageManager object  with example: obj = PipelineStageManager(pg_mesh, pipeline_axis, enable_interleave, num_model_chunks)