code: def get_forward_backward_func():
      """Retrieves the appropriate forward_backward function given the
      configuration of parallel_state.
  
      Returns a function that will perform all of the forward and
      backward passes of the model given the pipeline model parallel
      world size and virtual pipeline model parallel world size in the
      global parallel_state.
  
      Note that if using sequence parallelism, the sequence length component of
      the tensor shape is updated to original_sequence_length /
      tensor_model_parallel_world_size.
  
      The function returned takes the following arguments:
  
      forward_step_func (required): A function that takes a data
          iterator and a model as its arguments and return the model's
          forward output and the loss function. The loss function should
          take one torch.Tensor and return a torch.Tensor of loss and a
          dictionary of string -> torch.Tensor.
  
          A third argument, checkpoint_activations_microbatch, indicates
          that the activations for this microbatch should be
          checkpointed. A None value for this argument indicates that
          the default from the configuration should be used. This is
          used when the
          num_microbatches_with_partial_activation_checkpoints is used.
  
          For example:
  
          def loss_func(loss_mask, output_tensor):
              losses = output_tensor.float()
              loss_mask = loss_mask.view(-1).float()
              loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()
  
              # Reduce loss for logging.
              averaged_loss = average_losses_across_data_parallel_group([loss])
  
              return loss, {'lm loss': averaged_loss[0]}
  
          def forward_step(data_iterator, model):
              data, loss_mask = next(data_iterator)
              output = model(data)
              return output, partial(loss_func, loss_mask)
  
  
          forward_backward_func(forward_step_func=forward_step, ...)
  
  
      data_iterator (required): an iterator over the data, will be
          passed as is to forward_step_func. Expected to be a list of
          iterators in the case of interleaved pipeline parallelism.
  
      model (required): the actual model. Expected to be a list of modules in the case of interleaved
          pipeline parallelism. Must be a (potentially wrapped) megatron.core.models.MegatronModule.
  
      num_microbatches (int, required):
          The number of microbatches to go through
  
      seq_length (int, required): Sequence length of the current global batch. If this is a dual-stack
          transformer, this is the encoder's sequence length. This is ignored if variable_seq_lengths
          in the config is True. Otherwise, each microbatch in the current global batch size must use
          this sequence length.
  
      micro_batch_size (int, required): The number of sequences in a microbatch.
  
      decoder_seq_length (int, optional): The sequence length for the decoder in a dual-stack
          transformer. This is ignored for a single-stack transformer.
  
      forward_only (optional, default = False): Perform only the forward step
  
      collect_non_loss_data (optional, bool, default=False): TODO
  
      """
      pipeline_model_parallel_size = (
          parallel_state.get_pipeline_model_parallel_world_size()
      )
      if pipeline_model_parallel_size > 1:
          if (
              parallel_state.get_virtual_pipeline_model_parallel_world_size()
              is not None
          ):
              forward_backward_func = (
                  forward_backward_pipelining_with_interleaving
              )
          else:
              forward_backward_func = (
                  forward_backward_pipelining_without_interleaving
              )
      else:
          forward_backward_func = forward_backward_no_pipelining
      return forward_backward_func
  
imports: import contextlib
  from typing import Callable, Iterator, List, Optional, Union
  import torch
  from torch.autograd.variable import Variable
  from megatron.core import parallel_state
  from megatron.core.enums import ModelType
  from megatron.core.pipeline_parallel import p2p_communication
  from megatron.core.utils import get_attr_wrapped_model, get_model_config, get_model_type
label_desc: Use image node_id 1 for calling a global function with example usage: get_forward_backward_func() and returns: forward_backward_func