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: Backward step through the passed-in output tensor. If it is the last stage, the
  output_obj_grad is None, otherwise it is the gradients with respect to stage's output tensor.
  Returns the gradients with respect to the input tensor (None if first stage).
  This is a helper function and can be ignored by users.
  
  Args:
      engine (colossalai.legacy.engine.Engine): Colossalai engine for training and inference.
      input_obj (Union[:class:`torch.Tensor`, List[:class:`torch.Tensor`]]): input tensor for this pipeline stage.
      output_obj (Union[:class:`torch.Tensor`, List[:class:`torch.Tensor`]]): output tensor for this pipeline stage.
      output_obj_grad (Union[:class:`torch.Tensor`, List[:class:`torch.Tensor`]]): gradient of output tensor for this pipeline stage.
  
  Returns:
      Union[:class:`torch.Tensor`, List[:class:`torch.Tensor`]]: gradient of input tensor.
docstr_tok: ["Backward","step","through","the","passed-in","output","tensor",".","If","it","is","the","last","stage",",","the","output_obj_grad","is","None",",","otherwise","it","is","the","gradients","with","respect","to","stage","'s","output","tensor",".","Returns","the","gradients","with","respect","to","the","input","tensor","(","None","if","first","stage",")",".","This","is","a","helper","function","and","can","be","ignored","by","users",".","Args",":","engine","(","colossalai.legacy.engine.Engine",")",":","Colossalai","engine","for","training","and","inference",".","input_obj","(","Union","[",":","class",":","`","torch.Tensor","`",",","List","[",":","class",":","`","torch.Tensor","`","]","]",")",":","input","tensor","for","this","pipeline","stage",".","output_obj","(","Union","[",":","class",":","`","torch.Tensor","`",",","List","[",":","class",":","`","torch.Tensor","`","]","]",")",":","output","tensor","for","this","pipeline","stage",".","output_obj_grad","(","Union","[",":","class",":","`","torch.Tensor","`",",","List","[",":","class",":","`","torch.Tensor","`","]","]",")",":","gradient","of","output","tensor","for","this","pipeline","stage",".","Returns",":","Union","[",":","class",":","`","torch.Tensor","`",",","List","[",":","class",":","`","torch.Tensor","`","]","]",":","gradient","of","input","tensor","."]
code: def _backward_step(
      self, engine, input_obj, output_obj, output_obj_grad
  ):
      """Backward step through the passed-in output tensor. If it is the last stage, the
      output_obj_grad is None, otherwise it is the gradients with respect to stage's output tensor.
      Returns the gradients with respect to the input tensor (None if first stage).
      This is a helper function and can be ignored by users.
  
      Args:
          engine (colossalai.legacy.engine.Engine): Colossalai engine for training and inference.
          input_obj (Union[:class:`torch.Tensor`, List[:class:`torch.Tensor`]]): input tensor for this pipeline stage.
          output_obj (Union[:class:`torch.Tensor`, List[:class:`torch.Tensor`]]): output tensor for this pipeline stage.
          output_obj_grad (Union[:class:`torch.Tensor`, List[:class:`torch.Tensor`]]): gradient of output tensor for this pipeline stage.
  
      Returns:
          Union[:class:`torch.Tensor`, List[:class:`torch.Tensor`]]: gradient of input tensor.
      """
  
      # Retain the grad on the input_obj.
      if input_obj is not None:
          if isinstance(input_obj, torch.Tensor):
              input_obj.retain_grad()
          else:
              for in_tensor in input_obj:
                  if in_tensor is not None:
                      in_tensor.retain_grad()
      # Backward pass.
      if output_obj_grad is None:
          engine.backward(output_obj)
      else:
          engine.backward_by_grad(output_obj, output_obj_grad)
  
      # Collect the grad of the input_obj.
      input_obj_grad = None
      if input_obj is not None:
          if isinstance(input_obj, torch.Tensor):
              input_obj_grad = input_obj.grad
          else:
              input_obj_grad = []
              for in_tensor in input_obj:
                  input_obj_grad.append(in_tensor.grad)
  
      return input_obj_grad
  
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 10 for calling the PipelineSchedule obj's underlying member method code with example usage: obj._backward_step(engine, input_obj, output_obj, output_obj_grad) and returns: input_obj_grad