class_docstr: Token-based (input/output is tokens) encoder-decoder model (e.g. T5 Language model.)
code: def __init__(
      self,
      config: ModelParallelConfig,
      encoder_cfg: DictConfig,
      decoder_cfg: DictConfig,
      vocab_size: int,  # TODO: This should eventually go inside encoder_cfg and decoder_cfg when separate enc/dec tokenizers are supported.
      max_position_embeddings,
      num_tokentypes=0,
      parallel_output=True,
      pre_process=True,
      post_process=True,
      fp16_cross_entropy=False,
      precision=16,
      embedding_init_method_std=0.02,
      embedding_dropout=0.1,
      label_smoothing=0.0,
      add_encoder=True,
      add_decoder=True,
      share_token_embeddings=True,
      share_decoder_tokens_head_embeddings=True,
      tokens_head_bias=True,
      hiddens_cfg: DictConfig = None,  # allows for hidden state transformations before the decoder
  ):
      super(MegatronTokenLevelEncoderDecoderModule, self).__init__(
          config=config
      )
  
      self.encoder_cfg = encoder_cfg
      self.decoder_cfg = decoder_cfg
      self.parallel_output = parallel_output
      self.pre_process = pre_process
      self.post_process = post_process
      self.fp16_cross_entropy = fp16_cross_entropy
      self.precision = precision
      self.add_encoder = add_encoder
      self.add_decoder = add_decoder
      self.label_smoothing = label_smoothing
      self.share_token_embeddings = share_token_embeddings
      self.share_decoder_tokens_head_embeddings = (
          share_decoder_tokens_head_embeddings
      )
      self.tokens_head_bias = tokens_head_bias
      self.hiddens_cfg = hiddens_cfg
  
      encoder_kv_channels, decoder_kv_channels = self._validate_config()
  
      encoder, decoder = None, None
      if add_encoder:
          if pre_process:
              self.encoder_embedding = Embedding(
                  config=self.config,
                  hidden_size=encoder_cfg.hidden_size,
                  vocab_size=vocab_size,
                  max_sequence_length=max_position_embeddings,
                  init_method=init_method_normal(
                      embedding_init_method_std
                  ),
                  num_tokentypes=num_tokentypes,
                  embedding_dropout_prob=embedding_dropout,
                  position_embedding_type=encoder_cfg.get(
                      "position_embedding_type", "learned_absolute"
                  ),
              )
              self._encoder_embedding_key = "encoder_embedding"
          if (
              self.encoder_cfg.get(
                  "position_embedding_type", "learned_absolute"
              )
              == "relative"
          ):
              self.encoder_relative_position_embedding = T5RelativePositionEmbedding(
                  init_method=init_method_normal(
                      embedding_init_method_std
                  ),
                  num_attention_heads=encoder_cfg.num_attention_heads,
                  relative_position_num_buckets=encoder_cfg.relative_attention_num_buckets,
                  relative_position_max_distance=encoder_cfg.relative_attention_max_distance,
                  bidirectional=True,
                  layer_type=LayerType.encoder,
              )
              self._encoder_relative_position_embedding_key = (
                  "encoder_relative_position_embedding"
              )
              # Pipeline model parallel rank 0 will have the actual RPE weights. We zero it out on all other ranks and then sync them on setup.
              if parallel_state.get_pipeline_model_parallel_rank() != 0:
                  self.encoder_relative_position_embeddings_weight().data.fill_(
                      0
                  )
                  self.encoder_relative_position_embeddings_weight().shared = (
                      True
                  )
          elif (
              self.encoder_cfg.get(
                  "position_embedding_type", "learned_absolute"
              )
              == "alibi"
          ):
              self.encoder_relative_position_embedding = ALiBiRelativePositionEmbedding(
                  bidirectional=True,
                  num_attention_heads=encoder_cfg.num_attention_heads,
                  layer_type=LayerType.encoder,
                  num_attention_heads_alibi=None,
                  max_seq_len=max_position_embeddings,
              )
              self._encoder_relative_position_embedding_key = (
                  "encoder_alibi_position_embedding"
              )
          elif (
              self.encoder_cfg.get(
                  "position_embedding_type", "learned_absolute"
              )
              == "kerple"
          ):
              self.encoder_relative_position_embedding = KERPLERelativePositionEmbedding(
                  bidirectional=True,
                  num_attention_heads=encoder_cfg.num_attention_heads,
                  layer_type=LayerType.encoder,
                  num_attention_heads_kerple=None,
                  max_seq_len=max_position_embeddings,
              )
              self._encoder_relative_position_embedding_key = (
                  "encoder_kerple_position_embedding"
              )
          else:
              self.encoder_relative_position_embedding = None
  
          if encoder_cfg.get(
              "use_flash_attention", False
          ) and encoder_cfg.get(
              "position_embedding_type", "learned_absolute"
          ) in [
              "relative",
              "kerple",
          ]:
              raise ValueError(
                  "flash-attention not supported with relative or kerple at this point"
              )
  
          encoder = get_encoder_model(
              config=config,
              arch=encoder_cfg.arch,
              hidden_size=encoder_cfg.hidden_size,
              ffn_hidden_size=encoder_cfg.ffn_hidden_size,
              num_layers=encoder_cfg.num_layers,
              num_attention_heads=encoder_cfg.num_attention_heads,
              apply_query_key_layer_scaling=encoder_cfg.get(
                  "apply_query_key_layer_scaling", True
              ),
              kv_channels=encoder_kv_channels,
              init_method=init_method_normal(
                  encoder_cfg.get("init_method_std", 0.02)
              ),
              scaled_init_method=scaled_init_method_normal(
                  encoder_cfg.get("init_method_std", 0.02),
                  encoder_cfg.num_layers,
              ),
              encoder_attn_mask_type=AttnMaskType.padding,
              pre_process=pre_process,
              post_process=post_process,
              init_method_std=encoder_cfg.get("init_method_std", 0.02),
              hidden_dropout=encoder_cfg.get("hidden_dropout", 0.1),
              attention_dropout=encoder_cfg.get(
                  "attention_dropout", 0.1
              ),
              ffn_dropout=encoder_cfg.get("ffn_dropout", 0.0),
              precision=precision,
              fp32_residual_connection=encoder_cfg.get(
                  "fp32_residual_connection", False
              ),
              activations_checkpoint_method=encoder_cfg.get(
                  "activations_checkpoint_method", None
              ),
              activations_checkpoint_num_layers=encoder_cfg.get(
                  "activations_checkpoint_num_layers", 1
              ),
              activations_checkpoint_granularity=encoder_cfg.get(
                  "activations_checkpoint_granularity", None
              ),
              layernorm_epsilon=encoder_cfg.get(
                  "layernorm_epsilon", 1e-5
              ),
              bias_activation_fusion=encoder_cfg.get(
                  "bias_activation_fusion", True
              ),
              bias_dropout_add_fusion=encoder_cfg.get(
                  "bias_dropout_add_fusion", True
              ),
              masked_softmax_fusion=encoder_cfg.get(
                  "masked_softmax_fusion", True
              ),
              persist_layer_norm=encoder_cfg.get(
                  "persist_layer_norm", True
              ),
              openai_gelu=encoder_cfg.get("openai_gelu", False),
              onnx_safe=encoder_cfg.get("onnx_safe", False),
              hidden_steps=encoder_cfg.get("hidden_steps", -1),
              activation=encoder_cfg.get("activation", "gelu"),
              bias=encoder_cfg.get("bias", True),
              normalization=encoder_cfg.get(
                  "normalization", "layernorm"
              ),
              transformer_block_type=encoder_cfg.get(
                  "transformer_block_type", "pre_ln"
              ),
              headscale=encoder_cfg.get("headscale", False),
              parent_model_type=ModelType.encoder_and_decoder,
              num_self_attention_per_cross_attention=encoder_cfg.get(
                  "num_self_attention_per_cross_attention", 1
              ),
              megatron_legacy=encoder_cfg.get("megatron_legacy", False),
              normalize_attention_scores=encoder_cfg.get(
                  "normalize_attention_scores", True
              ),
              num_moe_experts=encoder_cfg.get("num_moe_experts", 1),
              moe_frequency=encoder_cfg.get("moe_frequency", 1),
              moe_dropout=encoder_cfg.get("moe_dropout", 0.0),
              position_embedding_type=encoder_cfg.get(
                  "position_embedding_type", "learned_absolute"
              ),
              use_flash_attention=encoder_cfg.get(
                  "use_flash_attention", False
              ),
          )
  
      if add_decoder:
          # If this is the decoder first stage
          if pre_process:
              # If the encoder also lies on this rank (PP = 1), then just assign embeddings directly.
              if (
                  hasattr(self, "encoder_embedding")
                  and share_token_embeddings
              ):
                  self.decoder_embedding = self.encoder_embedding
              else:
                  # This is the case where PP > 1 and first decoder first stage, or when not sharing embeddings with encoder
                  self.decoder_embedding = Embedding(
                      config=self.config,
                      hidden_size=decoder_cfg.hidden_size,
                      vocab_size=vocab_size,
                      max_sequence_length=max_position_embeddings,
                      init_method=init_method_normal(
                          embedding_init_method_std
                      ),
                      num_tokentypes=num_tokentypes,
                      embedding_dropout_prob=embedding_dropout,
                      position_embedding_type=decoder_cfg.get(
                          "position_embedding_type", "learned_absolute"
                      ),
                  )
                  # We initialize decoder embeddings, but set them to zero since we they're tied with the encoder embeddings.
                  # A later initialize_embedding call will synchronize the embeddings.
                  if share_token_embeddings:
                      self.decoder_embedding.zero_parameters()
  
              self._decoder_embedding_key = "decoder_embedding"
  
          if (
              self.decoder_cfg.get(
                  "position_embedding_type", "learned_absolute"
              )
              == "relative"
          ):
              self.decoder_relative_position_embedding = T5RelativePositionEmbedding(
                  init_method=init_method_normal(
                      embedding_init_method_std
                  ),
                  num_attention_heads=decoder_cfg.num_attention_heads,
                  relative_position_num_buckets=decoder_cfg.relative_attention_num_buckets,
                  relative_position_max_distance=decoder_cfg.relative_attention_max_distance,
                  bidirectional=False,
                  layer_type=LayerType.decoder,
              )
              self._decoder_relative_position_embedding_key = (
                  "decoder_relative_position_embedding"
              )
              # Pipeline model parallel rank == split_rank will have the actual RPE weights. We zero it out on all other ranks and then sync them on setup.
              if (
                  parallel_state.get_pipeline_model_parallel_rank()
                  != parallel_state.get_pipeline_model_parallel_split_rank()
              ):
                  self.decoder_relative_position_embeddings_weight().data.fill_(
                      0
                  )
                  self.decoder_relative_position_embeddings_weight().shared = (
                      True
                  )
  
              if (
                  not self.decoder_cfg.relative_position_bias_self_attention_only
              ):
                  self.decoder_cross_attention_relative_position_embedding = T5RelativePositionEmbedding(
                      init_method=init_method_normal(
                          embedding_init_method_std
                      ),
                      num_attention_heads=decoder_cfg.num_attention_heads,
                      relative_position_num_buckets=decoder_cfg.relative_attention_num_buckets,
                      relative_position_max_distance=decoder_cfg.relative_attention_max_distance,
                      bidirectional=True,
                      layer_type=LayerType.decoder,
                  )
                  self._decoder_cross_attention_relative_position_embedding_key = "decoder_cross_attention_relative_position_embedding"
                  if (
                      parallel_state.get_pipeline_model_parallel_rank()
                      != parallel_state.get_pipeline_model_parallel_split_rank()
                  ):
                      self.decoder_cross_attention_relative_position_embeddings_weight().data.fill_(
                          0
                      )
                      self.decoder_cross_attention_relative_position_embeddings_weight().shared = (
                          True
                      )
  
          elif (
              self.decoder_cfg.get(
                  "position_embedding_type", "learned_absolute"
              )
              == "alibi"
          ):
              self.decoder_relative_position_embedding = ALiBiRelativePositionEmbedding(
                  bidirectional=False,
                  num_attention_heads=decoder_cfg.num_attention_heads,
                  layer_type=LayerType.decoder,
                  num_attention_heads_alibi=None,
                  max_seq_len=max_position_embeddings,
              )
              self._decoder_relative_position_embedding_key = (
                  "decoder_alibi_position_embedding"
              )
          elif (
              self.decoder_cfg.get(
                  "position_embedding_type", "learned_absolute"
              )
              == "kerple"
          ):
              self.decoder_relative_position_embedding = KERPLERelativePositionEmbedding(
                  bidirectional=False,
                  num_attention_heads=decoder_cfg.num_attention_heads,
                  layer_type=LayerType.decoder,
                  num_attention_heads_kerple=None,
                  max_seq_len=max_position_embeddings,
              )
              self._decoder_relative_position_embedding_key = (
                  "decoder_kerple_position_embedding"
              )
          else:
              self.decoder_relative_position_embedding = None
  
          if decoder_cfg.get(
              "use_flash_attention", False
          ) and decoder_cfg.get(
              "position_embedding_type", "learned_absolute"
          ) in [
              "relative",
              "kerple",
          ]:
              raise ValueError(
                  "flash-attention not supported with relative or kerple at this point"
              )
  
          decoder = get_decoder_model(
              config=config,
              arch=decoder_cfg.arch,
              hidden_size=decoder_cfg.hidden_size,
              ffn_hidden_size=decoder_cfg.ffn_hidden_size,
              num_layers=decoder_cfg.num_layers,
              num_attention_heads=decoder_cfg.num_attention_heads,
              apply_query_key_layer_scaling=decoder_cfg.get(
                  "apply_query_key_layer_scaling", True
              ),
              kv_channels=decoder_kv_channels,
              init_method=init_method_normal(
                  decoder_cfg.get("init_method_std", 0.02)
              ),
              scaled_init_method=scaled_init_method_normal(
                  decoder_cfg.get("init_method_std", 0.02),
                  decoder_cfg.num_layers,
              ),
              decoder_attn_mask_type=AttnMaskType.causal,
              pre_process=pre_process,
              post_process=post_process,
              init_method_std=decoder_cfg.get("init_method_std", 0.02),
              hidden_dropout=decoder_cfg.get("hidden_dropout", 0.1),
              attention_dropout=decoder_cfg.get(
                  "attention_dropout", 0.1
              ),
              ffn_dropout=decoder_cfg.get("ffn_dropout", 0.0),
              precision=precision,
              fp32_residual_connection=decoder_cfg.get(
                  "fp32_residual_connection", False
              ),
              activations_checkpoint_method=decoder_cfg.get(
                  "activations_checkpoint_method", None
              ),
              activations_checkpoint_num_layers=decoder_cfg.get(
                  "activations_checkpoint_num_layers", 1
              ),
              activations_checkpoint_granularity=decoder_cfg.get(
                  "activations_checkpoint_granularity", None
              ),
              layernorm_epsilon=decoder_cfg.get(
                  "layernorm_epsilon", 1e-5
              ),
              bias_activation_fusion=decoder_cfg.get(
                  "bias_activation_fusion", True
              ),
              bias_dropout_add_fusion=decoder_cfg.get(
                  "bias_dropout_add_fusion", True
              ),
              masked_softmax_fusion=decoder_cfg.get(
                  "masked_softmax_fusion", True
              ),
              persist_layer_norm=decoder_cfg.get(
                  "persist_layer_norm", True
              ),
              openai_gelu=decoder_cfg.get("openai_gelu", False),
              onnx_safe=decoder_cfg.get("onnx_safe", False),
              hidden_steps=decoder_cfg.get("hidden_steps", -1),
              activation=decoder_cfg.get("activation", "gelu"),
              bias=decoder_cfg.get("bias", True),
              normalization=decoder_cfg.get(
                  "normalization", "layernorm"
              ),
              transformer_block_type=decoder_cfg.get(
                  "transformer_block_type", "pre_ln"
              ),
              headscale=decoder_cfg.get("headscale", False),
              parent_model_type=ModelType.encoder_and_decoder,
              megatron_legacy=decoder_cfg.get("megatron_legacy", False),
              normalize_attention_scores=decoder_cfg.get(
                  "normalize_attention_scores", True
              ),
              num_moe_experts=decoder_cfg.get("num_moe_experts", 1),
              moe_frequency=decoder_cfg.get("moe_frequency", 1),
              moe_dropout=decoder_cfg.get("moe_dropout", 0.0),
              position_embedding_type=decoder_cfg.get(
                  "position_embedding_type", "learned_absolute"
              ),
              use_flash_attention=decoder_cfg.get(
                  "use_flash_attention", False
              ),
          )
  
      hiddens_module = get_hiddens_module(
          hiddens_cfg, model_parallel_cfg=config
      )
      self.enc_dec_model = MegatronTransformerEncoderDecoderModule(
          config=config,
          encoder=encoder,
          decoder=decoder,
          hidden_steps=encoder_cfg.get("hidden_steps", -1),
          hiddens_module=hiddens_module,
      )
      self._enc_dec_model_key = "enc_dec_model"
  
      if self.share_token_embeddings:
          # This is only relevant for PP > 1.
          self.initialize_word_embeddings(
              init_method=init_method_normal(embedding_init_method_std),
              vocab_size=vocab_size,
              hidden_size=encoder_cfg.hidden_size,
          )
  
      if add_decoder and post_process:
          if share_decoder_tokens_head_embeddings:
              self.tokens_head = MegatronTokenLevelHead(
                  self.word_embeddings_weight().size(0),
                  parallel_output,
                  bias=tokens_head_bias,
              )
          else:
              self.tokens_head = tensor_parallel.ColumnParallelLinear(
                  input_size=decoder_cfg.hidden_size,
                  output_size=vocab_size,
                  config=config,
                  bias=tokens_head_bias,
                  gather_output=not self.parallel_output,
                  init_method=init_method_normal(
                      decoder_cfg.init_method_std
                  ),
              )
  
          self._tokens_head_key = "tokens_head"
      self.set_accepted_adapter_types(
          [PromptEncoderAdapterConfig._target_]
      )
  
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 1 to create a new MegatronTokenLevelEncoderDecoderModule object from inherited base classes: MegatronModule, adapter_mixins with example: obj = MegatronTokenLevelEncoderDecoderModule(config, encoder_cfg, decoder_cfg, vocab_size, max_position_embeddings, num_tokentypes, parallel_output, pre_process, post_process, fp16_cross_entropy, precision, embedding_init_method_std, embedding_dropout, label_smoothing, add_encoder, add_decoder, share_token_embeddings, share_decoder_tokens_head_embeddings, tokens_head_bias, hiddens_cfg)