code: def forward_step(
      forward_step_func: FwdStepFunc,
      batch: Optional[Batch],
      model: torch.nn.Module,
      input_tensor: Optional[Union[torch.Tensor, List[torch.Tensor]]],
      losses_reduced: List[torch.Tensor],
      dtype: torch.dtype,
      disable_autocast: bool = False,
      checkpoint_activations_micro_batch: Optional[bool] = None,
  ) -> Union[torch.Tensor, Sequence[torch.Tensor]]:
      """Forward step for passed-in model.
  
      If first stage, input tensor is obtained from batch, otherwise passed-in input_tensor is used.
  
      Returns output tensor.
  
      Args:
          forward_step_func: Model specific function. This takes a minibatch and model as its arguments and
              returns the model's output and the loss function.
          batch: minibatch
          model: unwrappable model
          input_tensor:
          losses_reduced:
          dtype:
          disable_autocast:
          checkpoint_activations_micro_batch:
  
      Returns:
          output_tensor
      """
      # timers = get_timers()
      # timers("forward-compute").start()
      unwrapped_model = unwrap_model(model)
      model_type = get_model_type(unwrapped_model)
      # NOTE (mkozuki): The passed `model` is expected to implement `set_input_tensor`.
      # See https://github.com/NVIDIA/Megatron-LM/blob/5ac5571ba0265af4c491ee0af1508ca7589450c6/megatron/model/transformer.py#L679  # NOQA
      # for the details of `set_input_tensor`.
      unwrap_output_tensor = not isinstance(input_tensor, list)
      if unwrap_output_tensor:
          input_tensor = [input_tensor]
  
      input_tensor = [
          inp.get() if isinstance(inp, FutureTensor) else inp
          for inp in input_tensor
      ]
  
      unwrapped_model.set_input_tensor(input_tensor)
      with torch.cuda.amp.autocast(
          enabled=not disable_autocast
          and dtype in (torch.half, torch.bfloat16),
          dtype=dtype,
      ):
          if checkpoint_activations_micro_batch is None:
              output_tensor, loss_func = forward_step_func(batch, model)
          else:
              output_tensor, loss_func = forward_step_func(
                  batch, model, checkpoint_activations_micro_batch
              )
          if parallel_state.is_pipeline_last_stage():
              output_tensor = loss_func(output_tensor)
              loss, loss_reduced = output_tensor
              output_tensor = loss / get_num_microbatches()
              losses_reduced.append(loss_reduced)
      # timers("forward-compute").stop()
  
      # If T5 model (or other model with encoder and decoder)
      # and in decoder stack, then send encoder_hidden_state
      # downstream as well.
      if (
          parallel_state.is_pipeline_stage_after_split()
          and model_type == ModelType.encoder_and_decoder
      ):
          return [output_tensor, input_tensor[-1]]
      if unwrap_output_tensor:
          return output_tensor
      return [output_tensor]
  
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 6 for calling a global function with example usage: forward_step(forward_step_func, batch, model, input_tensor, losses_reduced, dtype, disable_autocast, checkpoint_activations_micro_batch) and returns: list, list, output_tensor