class_docstr: Transformer class.
class_docstr_tok: ["Transformer","class","."]
code: def forward(
      self,
      hidden_states,
      attention_mask,
      layer_past=None,
      get_key_value=False,
      encoder_output=None,
      enc_dec_attn_mask=None,
      set_inference_key_value_memory=False,
      inference_max_sequence_len=None,
      rotary_pos_emb=None,  # list of positional embedding tensors, first one self attention, second one and third one are for cross attention (q, k)
      retrieved_emb=None,  # tensor of retrieved embedding of shape [b, k, r, n, d]
      self_attention_relative_position_bias=None,
      cross_attention_relative_position_bias=None,
      checkpoint_activations_all_layers=None,
  ):
      # Checks.
      if inference_max_sequence_len:
          assert (
              self.activations_checkpoint_method is None
          ), "inference does not work with activation checkpointing"
  
      if layer_past is not None:
          assert get_key_value, (
              "for not None values in layer_past, "
              "expected get_key_value to be set"
          )
      if get_key_value:
          assert self.activations_checkpoint_method is None, (
              "get_key_value does not work with "
              "activation checkpointing"
          )
  
      if not self.pre_process:
          # See set_input_tensor()
          hidden_states = self.input_tensor
  
      # TODO: @Yi Dong, what should this be?
      if retrieved_emb is not None:
          assert len(retrieved_emb.shape) == 5
          # this is retrieval decoder, need special transpose
          encoder_output = rearrange(
              retrieved_emb, "b k r n d -> k r n b d"
          ).contiguous()
  
      """
          is_first_microbatch is an optimization parameter for transformer engine.
          It indicates if the current step in the forward pass is the first in a gradient accumulation cycle.
          If set, FP8 weights are cached and some minor optimizations are applied to fuse_wgrad_accumulation
          """
      from apex.transformer.pipeline_parallel.utils import (
          _GLOBAL_NUM_MICROBATCHES_CALCULATOR,
      )
  
      num_micro_batches = getattr(
          _GLOBAL_NUM_MICROBATCHES_CALCULATOR, "num_micro_batches", 1
      )
  
      if self.sequence_parallel:
          rng_context = (
              tensor_parallel.random.get_cuda_rng_tracker().fork()
          )
      else:
          rng_context = nullcontext()
  
      with rng_context:
          # fp8_autocast will not do anything if TE or FP8 isn't used
          fp8_group = None
          if (
              self.fp8
              and parallel_state.model_parallel_is_initialized()
          ):
              fp8_group = parallel_state.get_amax_reduction_group(
                  with_context_parallel=True
              )
  
          if HAVE_TE:
              # if TE is installed but fp8 is not available then this will do nothing
              fp8_context = fp8_autocast(
                  enabled=self.fp8,
                  fp8_recipe=self.fp8_recipe,
                  fp8_group=fp8_group,
              )
  
          else:
              fp8_context = nullcontext()
  
          with fp8_context:
              if (
                  self.activations_checkpoint_granularity == "full"
                  and self.activations_checkpoint_num_layers > 0
              ):
                  hidden_states = self._checkpointed_forward(
                      hidden_states,
                      attention_mask,
                      encoder_output,
                      enc_dec_attn_mask,
                      rotary_pos_emb,
                      self_attention_relative_position_bias,
                      cross_attention_relative_position_bias,
                      checkpoint_activations_all_layers,
                  )
              else:
                  if get_key_value:
                      presents = []
  
                  if self.transformer_engine:
                      # Pass key value information to TE through inference_params to pre-allocate memory
                      if set_inference_key_value_memory:
                          self.inference_params = type("", (), {})()
                          self.inference_params.max_sequence_len = (
                              inference_max_sequence_len
                          )
                          self.inference_params.max_batch_size = (
                              hidden_states.size(1)
                          )
                          self.inference_params.batch_size_offset = 0
                          self.inference_params.key_value_memory_dict = (
                              {}
                          )
                          self.inference_params.sequence_len_offset = 0
                          self.inference_current_sequence_len = 0
  
                      if self.inference_params != None:
                          self.inference_params.sequence_len_offset = (
                              self.inference_current_sequence_len
                          )
  
                  if self.return_select_layer < 0:
                      assert (
                          parallel_state.get_pipeline_model_parallel_world_size()
                          == 1
                      ), f"##{parallel_state.get_pipeline_model_parallel_world_size}"
                      if self.num_layers + self.return_select_layer < 0:
                          logging.warning(
                              "Returning embeddings states only!"
                          )
                          return hidden_states
  
                  for index in range(self.num_layers):
                      layer = self._get_layer(index)
                      past = None
  
                      if layer_past is not None:
                          past = layer_past[index]
  
                      if (
                          self.activations_checkpoint_granularity
                          == "selective"
                      ):
                          # When pipeline-parallel size > 1 and 'num_micro_batches_with_partial_activation_checkpoints' = int,
                          # pipeline scheduling can force to checkpoint all layers or partial layers in a micro-batch.
                          if (
                              checkpoint_activations_all_layers == True
                              or self.activations_checkpoint_method
                              == "uniform"
                          ):
                              checkpoint_core_attention = True
                          elif (
                              self.activations_checkpoint_method
                              == "block"
                          ):
                              activations_checkpoint_num_layers = (
                                  self.activations_checkpoint_num_layers
                              )
                              # Decrease the number of layers to checkpoint at later pipeline stages
                              if (
                                  self.activations_checkpoint_layers_per_pipeline
                                  is not None
                              ):
                                  activations_checkpoint_num_layers -= int(
                                      parallel_state.get_pipeline_model_parallel_rank()
                                      * self.activations_checkpoint_layers_per_pipeline
                                  )
                              checkpoint_core_attention = (
                                  index
                                  < activations_checkpoint_num_layers
                              )
                      else:
                          checkpoint_core_attention = False
  
                      # Cache FP8 weight and transpose at (1) the first micro-batch in each global-batch
                      # in training, (2) the first micro-batch in each validation and test routine.
                      # The caching happens in TransformerEngine when passing `is_first_microbatch=True`.
                      is_first_microbatch = (
                          self.is_first_train_microbatch
                          and self.training
                      ) or (
                          self.is_prev_microbatch_training
                          and not self.training
                      )
                      if self.transformer_engine:
                          hidden_states = layer(
                              hidden_states,
                              attention_mask,
                              encoder_output=encoder_output,
                              enc_dec_attn_mask=enc_dec_attn_mask,
                              inference_params=self.inference_params,
                              is_first_microbatch=is_first_microbatch,
                              checkpoint_core_attention=checkpoint_core_attention,
                          )
                      else:
                          hidden_states = layer(
                              hidden_states,
                              attention_mask,
                              encoder_output=encoder_output,
                              enc_dec_attn_mask=enc_dec_attn_mask,
                              layer_past=past,
                              get_key_value=get_key_value,
                              set_inference_key_value_memory=set_inference_key_value_memory,
                              inference_max_sequence_len=inference_max_sequence_len,
                              rotary_pos_emb=rotary_pos_emb,
                              self_attention_relative_position_bias=self_attention_relative_position_bias,
                              cross_attention_relative_position_bias=cross_attention_relative_position_bias,
                              checkpoint_core_attention=checkpoint_core_attention,
                          )
  
                      if self.return_select_layer < 0:
                          assert (
                              parallel_state.get_pipeline_model_parallel_world_size()
                              == 1
                          ), f"##{parallel_state.get_pipeline_model_parallel_world_size}"
                          if (
                              index
                              == self.num_layers
                              + self.return_select_layer
                          ):
                              return hidden_states
  
                  # Update current sequence length outside of the loops
                  if self.transformer_engine:
                      self.inference_current_sequence_len += (
                          hidden_states.size(0)
                      )
  
      # Skip counter update for eval and activation checkpointing
      if torch.is_grad_enabled() and self.training:
          self.microbatch_count += 1
          if self.microbatch_count % num_micro_batches == 0:
              self.microbatch_count = 0
              self.is_first_train_microbatch = True
          else:
              self.is_first_train_microbatch = False
      self.is_prev_microbatch_training = self.training
  
      output = hidden_states
  
      # Final layer norm.
      if self.post_process:
          # only apply the final_layernorm for pre-ln
          if self.transformer_block_type != "post_ln":
              output = self.final_layernorm(hidden_states)
  
      if get_key_value:
          output = [output, presents]
  
      return output
  
imports: from contextlib import nullcontext
  from typing import Any, Callable, Optional
  import torch
  import torch.nn
  from einops import rearrange
  from nemo.collections.common.parts.adapter_modules import LinearAdapterConfig
  from nemo.collections.nlp.modules.common.megatron.adapters.parallel_adapters import AdapterName, ParallelLinearAdapterConfig, ParallelLinearAdapterWeightTyingConfig
  from nemo.collections.nlp.modules.common.megatron.attention import ParallelAttention, ParallelChunkedCrossAttention
  from nemo.collections.nlp.modules.common.megatron.fused_bias_dropout_add import bias_dropout_add, bias_dropout_add_fused_inference, bias_dropout_add_fused_train, dropout_add
  from nemo.collections.nlp.modules.common.megatron.fused_layer_norm import get_layer_norm
  from nemo.collections.nlp.modules.common.megatron.layer_norm_1p import LayerNorm1P, LPLayerNorm
  from nemo.collections.nlp.modules.common.megatron.layer_type import LayerType
  from nemo.collections.nlp.modules.common.megatron.mlp import ParallelMLP, SwitchMLP
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults
  from nemo.collections.nlp.parts import utils_funcs
  from nemo.core import adapter_mixins
  from nemo.utils import logging
label_desc: Use image node_id 6 for calling the ParallelTransformer obj's underlying member method code with example usage: obj.forward(hidden_states, attention_mask, layer_past, get_key_value, encoder_output, enc_dec_attn_mask, set_inference_key_value_memory, inference_max_sequence_len, rotary_pos_emb, retrieved_emb, _attention_relative_position_bias, cross_attention_relative_position_bias, checkpoint_activations_all_layers) and returns: output, hidden_states, hidden_states