docstr: A helper schedule class for pipeline parallelism running environment.
  It uses interleaved 1F1B strategy. Other properties are similar as
  :class:`NonPipelineSchedule`.
  
  Args:
      num_microbatches (int): The number of microbatches.
      num_model_chunks (int): The number of model chunks.
      data_process_func (Callable, optional):
          The preprocessing function which receives a batch of data, and it will be executed in `load_batch`.
      tensor_shape (torch.Size, optional): Specified shape in pipeline communication.
      scatter_gather_tensors (bool, optional):
          If set to `True`, communication will be reduced over pipeline when using 1D tensor parallelization.
docstr_tok: ["A","helper","schedule","class","for","pipeline","parallelism","running","environment",".","It","uses","interleaved","1F1B","strategy",".","Other","properties","are","similar","as",":","class",":","`","NonPipelineSchedule","`",".","Args",":","num_microbatches","(","int",")",":","The","number","of","microbatches",".","num_model_chunks","(","int",")",":","The","number","of","model","chunks",".","data_process_func","(","Callable",",","optional",")",":","The","preprocessing","function","which","receives","a","batch","of","data",",","and","it","will","be","executed","in","`","load_batch","`",".","tensor_shape","(","torch.Size",",","optional",")",":","Specified","shape","in","pipeline","communication",".","scatter_gather_tensors","(","bool",",","optional",")",":","If","set","to","`","True","`",",","communication","will","be","reduced","over","pipeline","when","using","1D","tensor","parallelization","."]
code: def __init__(
      self,
      num_microbatches: int,
      num_model_chunks: int,
      data_process_func: Callable = None,
      tensor_shape: Union[torch.Size, List[int], Tuple[int]] = None,
      scatter_gather_tensors: bool = False,
  ):
      """A helper schedule class for pipeline parallelism running environment.
      It uses interleaved 1F1B strategy. Other properties are similar as
      :class:`NonPipelineSchedule`.
  
      Args:
          num_microbatches (int): The number of microbatches.
          num_model_chunks (int): The number of model chunks.
          data_process_func (Callable, optional):
              The preprocessing function which receives a batch of data, and it will be executed in `load_batch`.
          tensor_shape (torch.Size, optional): Specified shape in pipeline communication.
          scatter_gather_tensors (bool, optional):
              If set to `True`, communication will be reduced over pipeline when using 1D tensor parallelization.
      """
      assert (
          num_microbatches % gpc.get_world_size(ParallelMode.PIPELINE)
          == 0
      ), "num_microbatches must be an integer multiple of pipeline parallel world size"
      assert (
          isinstance(num_model_chunks, int) and num_model_chunks > 0
      ), f"expected num_model_chunks to be an integer and larger than 0, but got {num_model_chunks}"
      super().__init__(
          num_microbatches,
          data_process_func=data_process_func,
          tensor_shape=tensor_shape,
          scatter_gather_tensors=scatter_gather_tensors,
      )
      gpc.set_virtual_pipeline_parallel_size(num_model_chunks)
      gpc.set_virtual_pipeline_parallel_rank(0)
      self.num_model_chunks = num_model_chunks
  
imports: import inspect
  from typing import Callable, List, Tuple, Union
  import torch.cuda
  import colossalai.legacy.communication
  from colossalai.legacy.amp.naive_amp import NaiveAMPModel
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from colossalai.legacy.utils import switch_virtual_pipeline_parallel_rank
  from colossalai.logging import get_dist_logger
  from colossalai.utils.device import get_current_device
  from ._base_schedule import BaseSchedule
label_desc: Use image node_id 1 to create a new InterleavedPipelineSchedule object from inherited base classes: PipelineSchedule with example: obj = InterleavedPipelineSchedule(num_microbatches, num_model_chunks, data_process_func, tensor_shape, scatter_gather_tensors)