class_docstr: Token-based (input/output is tokens) encoder-decoder model (e.g. T5 Language model.)
docstr: Return value is per token / per dimension (i.e., non collapsed loss value)
docstr_tok: ["Return","value","is","per","token","\/","per","dimension","(","i.e.",",","non","collapsed","loss","value",")"]
code: def forward(
      self,
      enc_input_ids=None,
      enc_attn_mask=None,
      dec_input_ids=None,
      dec_attn_mask=None,
      token_type_ids=None,
      labels=None,
      batch_data=None,  # additional data to be passed to hiddens module
      enc_output=None,  # Result of running the entire encoder
      enc_output_attn_mask=None,
      enc_input=None,  # Result of running encoder embedding only
      output_enc_hidden_only=False,
  ):
      """
      Return value is per token / per dimension (i.e., non collapsed loss value)
      """
      (
          encoder_self_attention_relative_position_bias,
          decoder_self_attention_relative_position_bias,
          decoder_cross_attention_relative_position_bias,
      ) = (None, None, None)
  
      if enc_input is not None and enc_output is not None:
          raise ValueError(
              """Both enc_input and enc_output are not None.
                  You should only be passing one of them.
                  enc_input is the result of the encoder embedding layer
                  enc_output is the result of running the entire transformer encoder."""
          )
  
      # In order of precedence, we use enc_output, enc_input, and then enc_input_ids to determine the encoder sequence length.
      if enc_output is not None:
          # If enc_output is provided in `batch_for_pipeline`, we need to transpose it from [B x S x H] -> [S x B x H].
          enc_output = enc_output.transpose(0, 1)
          enc_seq_length = enc_output.size(0)
      elif enc_input is not None:
          # If enc_input is provided, we need to transpose it from [B x S x H] -> [S x B x H].
          enc_input = enc_input.transpose(0, 1)
          enc_seq_length = enc_input.size(0)
      # Only need to run encoder embedding and position ids if enc_input or enc_output is not provided.
      elif enc_input_ids is not None:
          enc_seq_length = enc_input_ids.size(1)
          if self.pre_process and self.add_encoder:
              # We don't need position ids for RPE, because the embedding layer does not have position embeddings.
              if self.encoder_relative_position_embedding is None:
                  enc_position_ids = build_position_ids(enc_input_ids)
              else:
                  enc_position_ids = None
              enc_input = self.encoder_embedding(
                  enc_input_ids,
                  enc_position_ids,
                  token_type_ids=token_type_ids,
              )
              if self.is_adapter_available():
                  _sq, _bs, _hs = enc_input.size()
                  ptuning_adapter = self.get_adapter_module(
                      AdapterName.PTUNING_ADAPTER
                  )
                  v = ptuning_adapter.virtual_tokens
                  if (
                      ptuning_adapter and _sq >= v
                  ):  # The sequence should be longer the v to insert virtual embeddings.
                      virtual_embeddings = ptuning_adapter(_bs)
                      enc_input = enc_input[
                          v:, :, :
                      ]  # the first v tokens are pads so that they can be swapped out with virtual embeddings.
                      enc_input = torch.concat(
                          [virtual_embeddings, enc_input], dim=0
                      )
          else:
              enc_input = None
      else:
          # This should only happen with PP > 1 for enc-dec prompt learning models
          enc_seq_length = enc_attn_mask.size(1)
  
      if (
          self.add_encoder
          and self.encoder_relative_position_embedding is not None
      ):
          encoder_self_attention_relative_position_bias = (
              self.encoder_relative_position_embedding(
                  query_seq_length=enc_seq_length,
                  key_seq_length=enc_seq_length,
              )
          )
  
      if output_enc_hidden_only:
          # When pipeline parallel > 1 we need to make sure encoder exist (will be missing in decoder)
          if (
              enc_output is None
              and self.enc_dec_model.encoder is not None
          ):
              enc_output = self.enc_dec_model.encode(
                  enc_input=enc_input,
                  enc_attn_mask=enc_attn_mask,
                  enc_layer_past=None,
                  enc_get_key_value=False,
                  enc_self_attention_relative_position_bias=encoder_self_attention_relative_position_bias,
                  batch_data=batch_data,
              )
          else:
              enc_output = self.enc_dec_model.encoder_hidden_state
  
          return enc_output
      else:
          if enc_output_attn_mask is None:
              enc_output_attn_mask = enc_attn_mask
  
          if self.pre_process and self.add_decoder:
              # We don't need position ids for RPE, because the embedding layer does not have position embeddings.
              if self.decoder_relative_position_embedding is None:
                  dec_position_ids = build_position_ids(dec_input_ids)
              else:
                  dec_position_ids = None
              dec_input = self.decoder_embedding(
                  dec_input_ids,
                  dec_position_ids,
                  token_type_ids=token_type_ids,
              )
          else:
              # Note: This is when the decoder itself is split across PP ranks.
              dec_input = None
  
          if (
              self.add_decoder
              and self.decoder_relative_position_embedding is not None
          ):
              decoder_self_attention_relative_position_bias = (
                  self.decoder_relative_position_embedding(
                      query_seq_length=dec_input_ids.size(1),
                      key_seq_length=dec_input_ids.size(1),
                  )
              )
              if (
                  not self.decoder_cfg.relative_position_bias_self_attention_only
              ):
                  decoder_cross_attention_relative_position_bias = self.decoder_cross_attention_relative_position_embedding(
                      query_seq_length=dec_input_ids.size(1),
                      key_seq_length=enc_seq_length,
                  )
              else:
                  decoder_cross_attention_relative_position_bias = None
  
          output = self.enc_dec_model(
              enc_input=enc_input,
              enc_attn_mask=enc_attn_mask,
              dec_input=dec_input,
              dec_attn_mask=dec_attn_mask,
              enc_layer_past=None,
              enc_get_key_value=False,
              enc_output=enc_output,
              enc_output_attn_mask=enc_output_attn_mask,
              dec_layer_past=None,
              dec_get_key_value=False,
              enc_self_attention_relative_position_bias=encoder_self_attention_relative_position_bias,
              dec_self_attention_relative_position_bias=decoder_self_attention_relative_position_bias,
              dec_cross_attention_relative_position_bias=decoder_cross_attention_relative_position_bias,
              batch_data=batch_data,
          )
  
          if self.post_process and self.add_decoder:
              (
                  dec_output,
                  enc_output,
              ) = output  # [s, b, h], enc_output might be a dict if hiddens_module is used
              # project decoder output to vocabulary-size dimensions
              if self.share_decoder_tokens_head_embeddings:
                  token_logits = self.tokens_head(
                      dec_output, self.word_embeddings_weight()
                  )
              else:
                  token_logits = self.tokens_head(dec_output)[0]
  
              if labels is not None:
                  # compute loss here
                  # [b, s] -> [s, b]
                  labels = labels.transpose(0, 1).contiguous()
  
                  # Set label smoothing to 0 if in eval mode.
                  label_smoothing = (
                      self.label_smoothing if self.training else 0.0
                  )
  
                  # tensor_parallel.vocab_parallel_cross_entropy performs log_softmax and return log p(x_i|z) per token i
                  if self.fp16_cross_entropy:
                      assert token_logits.dtype == torch.half
                      tokens_loss = vocab_parallel_cross_entropy(
                          token_logits, labels, label_smoothing
                      )
                  else:
                      tokens_loss = vocab_parallel_cross_entropy(
                          token_logits.float(), labels, label_smoothing
                      )
  
                  # [s, b] -> [b, s]
                  tokens_loss = tokens_loss.transpose(0, 1).contiguous()
  
                  # check if hiddens is used
                  if self.hiddens_cfg is not None:
                      loss_dict = self.enc_dec_model.hiddens_module.apply_loss_transforms(
                          outputs=enc_output,
                          batch_data=batch_data,
                      )
                      loss_dict["tokens_loss"] = tokens_loss
                      # We need to store default output in a known key, so that we can mimic default behaviour
                      loss_dict["output"] = tokens_loss
                      return loss_dict
                  else:
                      return tokens_loss
              else:
                  # else return token logits (and hiddens if needed)
                  # [s, b, h] -> [b, s, h]
                  token_logits = token_logits.transpose(
                      0, 1
                  ).contiguous()
                  if self.hiddens_cfg is not None:
                      # return all hiddens and token logits
                      hiddens_dict = enc_output
                      hiddens_dict["token_logits"] = token_logits
                      # We need to store default output in a known key, so that we can mimic default behaviour
                      hiddens_dict["output"] = token_logits
                      return hiddens_dict
                  else:
                      return token_logits
  
          elif self.add_decoder and not self.add_encoder:
              decoder_output, _ = output
              return decoder_output
          else:
              encoder_output = output
              return encoder_output
  
imports: import torch
  from omegaconf import DictConfig
  from nemo.collections.nlp.modules.common.megatron.adapters.parallel_adapters import AdapterName, PromptEncoderAdapterConfig
  from nemo.collections.nlp.modules.common.megatron.hiddens import get_hiddens_module
  from nemo.collections.nlp.modules.common.megatron.language_model import Embedding
  from nemo.collections.nlp.modules.common.megatron.layer_type import LayerType
  from nemo.collections.nlp.modules.common.megatron.megatron_decoders import get_decoder_model
  from nemo.collections.nlp.modules.common.megatron.megatron_encoder_decoder import MegatronTransformerEncoderDecoderModule
  from nemo.collections.nlp.modules.common.megatron.megatron_encoders import get_encoder_model
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.position_embedding import ALiBiRelativePositionEmbedding, KERPLERelativePositionEmbedding, T5RelativePositionEmbedding
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, build_position_ids, init_method_normal, parallel_lm_logits, scaled_init_method_normal
  from nemo.collections.nlp.modules.common.megatron.vocab_parallel_cross_entropy import vocab_parallel_cross_entropy
  from nemo.core.classes.mixins import adapter_mixins
label_desc: Use image node_id 7 for calling the MegatronTokenLevelEncoderDecoderModule obj's underlying member method code with example usage: obj.forward(enc_input_ids, enc_attn_mask, dec_input_ids, dec_attn_mask, token_type_ids, labels, batch_data, enc_output, enc_output_attn_mask, enc_input, output_enc_hidden_only) and returns: enc_output, decoder_output, encoder_output, loss_dict, tokens_loss, hiddens_dict, token_logits