code: def backward_step(
      input_tensor: Optional[torch.Tensor],
      output_tensor: torch.Tensor,
      output_tensor_grad: Optional[torch.Tensor],
      model_type: ModelType,
      *,
      grad_scaler: Optional[torch.cuda.amp.GradScaler] = None,
      deallocate_pipeline_outputs: bool = False,
  ) -> Union[None, torch.Tensor, Sequence[torch.Tensor]]:
      """Backward step through passed-in output tensor.
  
      If last stage, output_tensor_grad is None, otherwise gradient of loss
      with respect to stage's output tensor.
  
      Returns gradient of loss with respect to input tensor (None if first
      stage).
  
      Args:
          input_tensor:
          output_tensor:
          output_tensor_grad:
      Keyword Arguments:
          grad_scaler:
          deallocate_pipeline_outputs: Experimental.
      Returns:
          input_tensor_grad
      """
  
      # timers = get_timers()
      # timers("backward-compute").start()
  
      # Retain the grad on the input_tensor.
      unwrap_input_tensor_grad = not isinstance(input_tensor, list)
      if unwrap_input_tensor_grad:
          input_tensor = [input_tensor]
  
      input_tensor = [
          inp.get() if isinstance(inp, FutureTensor) else inp
          for inp in input_tensor
      ]
  
      for x in input_tensor:
          if x is not None:
              x.retain_grad()
  
      if not isinstance(output_tensor, list):
          output_tensor = [output_tensor]
  
      output_tensor = [
          out.get() if isinstance(out, FutureTensor) else out
          for out in output_tensor
      ]
  
      if not isinstance(output_tensor_grad, list):
          output_tensor_grad = [output_tensor_grad]
  
      output_tensor_grad = [
          ogr.get() if isinstance(ogr, FutureTensor) else ogr
          for ogr in output_tensor_grad
      ]
  
      # Backward pass.
      if grad_scaler is not None and output_tensor_grad[0] is None:
          output_tensor[0] = grad_scaler.scale(output_tensor[0])
      if deallocate_pipeline_outputs:
          custom_backward(output_tensor[0], output_tensor_grad[0])
      else:
          torch.autograd.backward(
              output_tensor[0], grad_tensors=output_tensor_grad[0]
          )
  
      # Collect the grad of the input_tensor.
      input_tensor_grad = [None]
      if input_tensor is not None:
          input_tensor_grad = []
          for x in input_tensor:
              input_tensor_grad.append(None if x is None else x.grad)
  
      # Handle single skip connection if it exists (encoder_hidden_state in model with encoder and decoder).
      if (
          parallel_state.get_pipeline_model_parallel_world_size() > 1
          and parallel_state.is_pipeline_stage_after_split()
          and model_type == ModelType.encoder_and_decoder
      ):
          if output_tensor_grad[1] is not None:
              # todo (mkozuki): Replace the inplace add with `+= output_tensor_grad[1]`?
              input_tensor_grad[-1].add_(output_tensor_grad[1])
  
      # timers("backward-compute").stop()
      return (
          input_tensor_grad[0]
          if unwrap_input_tensor_grad
          else input_tensor_grad
      )
  
file_path: apex/apex/transformer/pipeline_parallel/schedules/common.py
imports: from typing import Any, Callable, Dict, List, Tuple, Union, Optional, Sequence
  import torch
  from torch.autograd.variable import Variable
  from apex.normalization.fused_layer_norm import FusedLayerNorm
  from apex.transformer import parallel_state
  from apex.transformer.enums import ModelType
  from apex.transformer.pipeline_parallel.p2p_communication import FutureTensor
  from apex.transformer.pipeline_parallel.utils import get_num_microbatches
  from apex.transformer.pipeline_parallel.utils import listify_model
  from apex.transformer.pipeline_parallel.utils import unwrap_model
  from apex.transformer.pipeline_parallel.utils import get_model_type
  from apex.transformer.tensor_parallel.layers import set_defaults_if_not_set_tensor_model_parallel_attributes
  from apex.transformer.log_util import get_transformer_logger
label_desc: Use image node_id 7 for calling a global function with example usage: backward_step(input_tensor, output_tensor, output_tensor_grad, model_type) and returns: unknown