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"]
docstr: Runs non-interleaved 1F1B schedule, with communication between pipeline stages.
  Returns a tuple with losses if the last stage, an empty tuple otherwise.
  
  Args:
      engine (colossalai.legacy.engine.Engine): Colossalai engine for training and inference.
      data_iter (Iterable): Dataloader as the form of an iterator, obtained by calling iter(dataloader).
      forward_only (bool, optional):
          Whether run forward step only. Default is false. If true, no backward will be run.
      return_loss (bool, optional): Whether returns the loss value. Default is true.
      return_output_label (bool, optional): If False, the output and label won't be returned.
  
  Returns:
      Tuple[:class:`torch.Tensor`]: A tuple of (output, label, loss), loss and label could be None.
docstr_tok: ["Runs","non-interleaved","1F1B","schedule",",","with","communication","between","pipeline","stages",".","Returns","a","tuple","with","losses","if","the","last","stage",",","an","empty","tuple","otherwise",".","Args",":","engine","(","colossalai.legacy.engine.Engine",")",":","Colossalai","engine","for","training","and","inference",".","data_iter","(","Iterable",")",":","Dataloader","as","the","form","of","an","iterator",",","obtained","by","calling","iter","(","dataloader",")",".","forward_only","(","bool",",","optional",")",":","Whether","run","forward","step","only",".","Default","is","false",".","If","true",",","no","backward","will","be","run",".","return_loss","(","bool",",","optional",")",":","Whether","returns","the","loss","value",".","Default","is","true",".","return_output_label","(","bool",",","optional",")",":","If","False",",","the","output","and","label","wo","n't","be","returned",".","Returns",":","Tuple","[",":","class",":","`","torch.Tensor","`","]",":","A","tuple","of","(","output",",","label",",","loss",")",",","loss","and","label","could","be","None","."]
code: def forward_backward_step(
      self,
      engine,
      data_iter,
      forward_only=False,
      return_loss=True,
      return_output_label=True,
  ):
      """Runs non-interleaved 1F1B schedule, with communication between pipeline stages.
      Returns a tuple with losses if the last stage, an empty tuple otherwise.
  
      Args:
          engine (colossalai.legacy.engine.Engine): Colossalai engine for training and inference.
          data_iter (Iterable): Dataloader as the form of an iterator, obtained by calling iter(dataloader).
          forward_only (bool, optional):
              Whether run forward step only. Default is false. If true, no backward will be run.
          return_loss (bool, optional): Whether returns the loss value. Default is true.
          return_output_label (bool, optional): If False, the output and label won't be returned.
  
      Returns:
          Tuple[:class:`torch.Tensor`]: A tuple of (output, label, loss), loss and label could be None.
      """
  
      assert (
          forward_only or return_loss
      ), "The argument 'return_loss' has to be True when 'forward_only' is False, but got False."
      self.load_batch(data_iter)
      num_warmup_microbatches = (
          gpc.get_world_size(ParallelMode.PIPELINE)
          - gpc.get_local_rank(ParallelMode.PIPELINE)
          - 1
      )
      num_warmup_microbatches = min(
          num_warmup_microbatches, self.num_microbatches
      )
      num_microbatches_remaining = (
          self.num_microbatches - num_warmup_microbatches
      )
  
      # Input, output tensors only need to be saved when doing backward passes
      input_objs = None
      output_objs = None
      if not forward_only:
          input_objs = []
          output_objs = []
      return_tensors = []
      if return_loss and gpc.is_pipeline_last_stage(
          ignore_virtual=True
      ):
          accum_loss = torch.zeros(1, device=get_current_device())
      else:
          accum_loss = None
      # Used for tensor meta information communication
      ft_shapes = self.tensor_shape
      bt_shapes = None
      fs_checker = self.tensor_shape is None
  
      # Run warmup forward passes.
      for i in range(num_warmup_microbatches):
          if not gpc.is_first_rank(ParallelMode.PIPELINE):
              ft_shapes = comm.recv_obj_meta(ft_shapes)
          input_obj = comm.recv_forward(
              ft_shapes,
              dtype=self.dtype,
              scatter_gather_tensors=self.scatter_gather_tensors,
          )
          output_obj = self._forward_step(
              engine,
              input_obj,
              return_tensors,
              return_output_label=return_output_label,
              accum_loss=accum_loss,
          )
          if not gpc.is_last_rank(ParallelMode.PIPELINE):
              if isinstance(output_obj, torch.Tensor):
                  bt_shapes = output_obj.shape
              else:
                  bt_shapes = []
                  for out_tensor in output_obj:
                      bt_shapes.append(out_tensor.shape)
              fs_checker = comm.send_obj_meta(output_obj, fs_checker)
          comm.send_forward(
              output_obj,
              scatter_gather_tensors=self.scatter_gather_tensors,
          )
  
          if not forward_only:
              input_objs.append(input_obj)
              output_objs.append(output_obj)
  
      # Before running 1F1B, need to receive first forward tensor.
      # If all microbatches are run in warmup / cooldown phase, then no need to
      # receive this tensor here.
      if num_microbatches_remaining > 0:
          if not gpc.is_first_rank(ParallelMode.PIPELINE):
              ft_shapes = comm.recv_obj_meta(ft_shapes)
          input_obj = comm.recv_forward(
              ft_shapes,
              dtype=self.dtype,
              scatter_gather_tensors=self.scatter_gather_tensors,
          )
  
      # Run 1F1B in steady state.
      for i in range(num_microbatches_remaining):
          last_iteration = i == (num_microbatches_remaining - 1)
  
          output_obj = self._forward_step(
              engine,
              input_obj,
              return_tensors,
              return_output_label=return_output_label,
              accum_loss=accum_loss,
          )
          if forward_only:
              comm.send_forward(
                  output_obj,
                  scatter_gather_tensors=self.scatter_gather_tensors,
              )
  
              if not last_iteration:
                  input_obj = comm.recv_forward(
                      ft_shapes,
                      dtype=self.dtype,
                      scatter_gather_tensors=self.scatter_gather_tensors,
                  )
  
          else:
              output_obj_grad = comm.send_forward_recv_backward(
                  output_obj,
                  bt_shapes,
                  dtype=self.dtype,
                  scatter_gather_tensors=self.scatter_gather_tensors,
              )
  
              # Add input_obj and output_obj to end of list.
              input_objs.append(input_obj)
              output_objs.append(output_obj)
  
              # Pop output_obj and output_obj from the start of the list for
              # the backward pass.
              input_obj = input_objs.pop(0)
              output_obj = output_objs.pop(0)
  
              input_obj_grad = self._backward_step(
                  engine, input_obj, output_obj, output_obj_grad
              )
  
              if last_iteration:
                  input_obj = None
                  comm.send_backward(
                      input_obj_grad,
                      scatter_gather_tensors=self.scatter_gather_tensors,
                  )
              else:
                  input_obj = comm.send_backward_recv_forward(
                      input_obj_grad,
                      ft_shapes,
                      dtype=self.dtype,
                      scatter_gather_tensors=self.scatter_gather_tensors,
                  )
  
      # Run cooldown backward passes.
      if not forward_only:
          for i in range(num_warmup_microbatches):
              input_obj = input_objs.pop(0)
              output_obj = output_objs.pop(0)
  
              output_obj_grad = comm.recv_backward(
                  bt_shapes,
                  dtype=self.dtype,
                  scatter_gather_tensors=self.scatter_gather_tensors,
              )
  
              input_obj_grad = self._backward_step(
                  engine, input_obj, output_obj, output_obj_grad
              )
  
              comm.send_backward(
                  input_obj_grad,
                  scatter_gather_tensors=self.scatter_gather_tensors,
              )
  
      if len(return_tensors) > 0:
          output, label = pack_return_tensors(return_tensors)
          return output, label, accum_loss
      else:
          return None, None, accum_loss
  
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 11 for calling the PipelineSchedule obj's underlying member method code with example usage: obj.forward_backward_step(engine, data_iter, forward_only, return_loss, return_output_label) and returns: output,  label,  accum_loss, None,  None,  accum_loss