class_docstr: A helper schedule class for pipeline parallelism running environment.
  It uses non-interleaved 1F1B strategy. Other properties are similar as
  :class:`NonPipelineSchedule`.
  
  Args:
      num_microbatches (int): The number of microbatches.
      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.
  
  Example:
  
      # this shows an example of customized data_process_func
      def data_process_func(stage_output, dataloader_output):
          output1, output2 = stage_output
          item1, item2, item3 = dataloader_output
  
          # assume item2 is not needed
          data = (output1, output2, item1)
          label = item3
          return data, label
class_docstr_tok: ["A","helper","schedule","class","for","pipeline","parallelism","running","environment",".","It","uses","non-interleaved","1F1B","strategy",".","Other","properties","are","similar","as",":","class",":","`","NonPipelineSchedule","`",".","Args",":","num_microbatches","(","int",")",":","The","number","of","microbatches",".","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",".","Example",":","#","this","shows","an","example","of","customized","data_process_func","def","data_process_func","(","stage_output",",","dataloader_output",")",":","output1",",","output2","=","stage_output","item1",",","item2",",","item3","=","dataloader_output","#","assume","item2","is","not","needed","data","=","(","output1",",","output2",",","item1",")","label","=","item3","return","data",",","label"]
code: def pre_processing(self, engine):
      from colossalai.legacy.zero import ShardedModelV2
  
      # TODO: remove this after testing new zero with pipeline parallelism
      model = engine.model
      if isinstance(model, NaiveAMPModel):
          self.dtype = torch.half
          model = model.model
      if isinstance(model, ShardedModelV2):
          self.dtype = torch.half
          model = model.module
  
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 5 for calling the PipelineSchedule obj's underlying member method code with example usage: obj.pre_processing(engine) without return types