class_docstr: Token-based (input/output is tokens) retrieval encoder-decoder model
class_docstr_tok: ["Token-based","(","input\/output","is","tokens",")","retrieval","encoder-decoder","model"]
code: def __init__(
      self,
      config: ModelParallelConfig,
      vocab_size,
      hidden_size,
      max_position_embeddings,
      num_attention_heads,
      ffn_hidden_size,
      apply_query_key_layer_scaling=True,
      kv_channels=None,
      num_tokentypes=0,
      parallel_output=True,
      pre_process=True,
      post_process=True,
      init_method_std=0.02,
      fp16_cross_entropy=False,
      megatron_amp_O2=False,
      hidden_dropout=0.1,
      attention_dropout=0.1,
      precision=16,
      fp32_residual_connection=False,
      activations_checkpoint_method=None,
      activations_checkpoint_num_layers=1,
      layernorm_epsilon=1e-5,
      persist_layer_norm=False,
      bias_gelu_fusion=True,
      bias_dropout_add_fusion=True,
      masked_softmax_fusion=True,
      openai_gelu=False,
      activation="gelu",
      onnx_safe=False,
      bias=True,
      normalization="layernorm",
      headscale=False,
      transformer_block_type="pre_ln",
      hidden_steps=-1,
      add_encoder=True,
      add_decoder=True,
      chunk_size=64,
      enc_num_layers=4,  # total number of encoder layers
      dec_num_layers=6,  # total number of decoder layers
      enc_cross_attention=[3],  # layer numbers for cross attention
      dec_cross_attention=[
          3,
          5,
      ],  # layer numbers for chunked cross attention
      add_position_embedding=False,
      tokenizer=None,  # tokenizer
      normalize_attention_scores=True,
      activations_checkpoint_granularity=None,
      megatron_lm_compatible=False,
      version=1,
  ):
      super(
          MegatronRetrievalTokenLevelEncoderDecoderModule, self
      ).__init__()
      if megatron_lm_compatible:
          assert (
              apply_query_key_layer_scaling
          ), "megatron lm compatible model has to set apply_query_key_layer_scaling"
          assert (
              add_position_embedding
          ), "megatron lm compatible model has to set add_position_embedding"
      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.add_abs_position_embedding = add_position_embedding  # whether use absolute position embedding
      self.tokenizer = tokenizer
      self.eod_id = tokenizer.eos_id
      self.transformer_block_type = transformer_block_type
      self.num_chunked_cross_attention = len(dec_cross_attention)
      self.megatron_lm_compatible = megatron_lm_compatible
  
      if kv_channels is None:
          assert (
              hidden_size % num_attention_heads == 0
          ), "hidden_size must be divisible by num_attention_heads if kv_channels is None"
          kv_channels = hidden_size // num_attention_heads
  
      if pre_process:
          self.encoder_embedding = Embedding(
              config=config,
              hidden_size=hidden_size,
              vocab_size=vocab_size,
              max_sequence_length=max_position_embeddings,
              init_method=init_method_normal(init_method_std),
              num_tokentypes=num_tokentypes,
              embedding_dropout_prob=hidden_dropout,
              position_embedding_type="learned_absolute"
              if add_position_embedding
              else "",
              transpose_batch_sequence=False,
          )
          self._embedding_key = "embedding"
  
      encoder_init = init_method_normal(init_method_std)
      encoder_scaled_init = scaled_init_method_normal(
          init_method_std, dec_num_layers
      )
      pre_decoder_init = init_method_normal(init_method_std)
      pre_decoder_scaled_init = scaled_init_method_normal(
          init_method_std, dec_num_layers
      )
      post_decoder_init = init_method_normal(init_method_std)
      post_decoder_scaled_init = scaled_init_method_normal(
          init_method_std, dec_num_layers
      )
  
      if add_encoder:
          enc_layer_types = []
          for i in range(enc_num_layers):
              if i in enc_cross_attention:
                  enc_layer_types.append(LayerType.retrieval_encoder)
              else:
                  enc_layer_types.append(LayerType.encoder)
  
          self.encoder = get_encoder_model(
              config=config,
              arch="retro",
              hidden_size=hidden_size,
              ffn_hidden_size=ffn_hidden_size,
              num_layers=enc_num_layers,
              num_attention_heads=num_attention_heads,
              apply_query_key_layer_scaling=apply_query_key_layer_scaling,
              kv_channels=kv_channels,
              init_method=encoder_init,
              scaled_init_method=encoder_scaled_init,
              pre_process=pre_process,
              post_process=False
              if megatron_lm_compatible
              else post_process,  # megatron lm model has no final layer_norm
              init_method_std=init_method_std,
              hidden_dropout=hidden_dropout,
              attention_dropout=attention_dropout,
              precision=precision,
              fp32_residual_connection=fp32_residual_connection,
              activations_checkpoint_method=activations_checkpoint_method,
              activations_checkpoint_num_layers=activations_checkpoint_num_layers,
              activations_checkpoint_granularity=activations_checkpoint_granularity,
              layernorm_epsilon=layernorm_epsilon,
              bias_activation_fusion=bias_gelu_fusion,
              bias_dropout_add_fusion=bias_dropout_add_fusion,
              masked_softmax_fusion=masked_softmax_fusion,
              persist_layer_norm=persist_layer_norm,
              openai_gelu=openai_gelu,
              onnx_safe=onnx_safe,
              hidden_steps=hidden_steps,
              activation=activation,
              bias=bias,
              normalization=normalization,
              transformer_block_type=transformer_block_type,
              headscale=headscale,
              parent_model_type=ModelType.encoder_and_decoder,
              layer_type=enc_layer_types,
              chunk_size=chunk_size,
              layer_number_offset=0,
              normalize_attention_scores=normalize_attention_scores,
              turn_off_rop=megatron_lm_compatible,
              version=version,
          )
          self._encoder_key = "encoder"
  
      if add_decoder:
          pre_decoder_num_layers = min(dec_cross_attention)
          pre_decoder_layer_types = []
          for i in range(pre_decoder_num_layers):
              pre_decoder_layer_types.append(LayerType.encoder)
          pre_decoder_layer_types.append(LayerType.decoder_pre_mlp)
  
          post_decoder_num_layers = (
              dec_num_layers - pre_decoder_num_layers
          )
          post_decoder_layer_types = []
          # the first layer in post decoder has to be chunked cross attention without self attention
          assert pre_decoder_num_layers in dec_cross_attention
          for i in range(post_decoder_num_layers):
              if i == 0:
                  post_decoder_layer_types.append(
                      LayerType.retrieval_decoder_after_self_attn
                  )
              elif i + pre_decoder_num_layers in dec_cross_attention:
                  post_decoder_layer_types.append(
                      LayerType.retrieval_decoder
                  )
              else:
                  post_decoder_layer_types.append(LayerType.encoder)
  
          # it is used to process the inputs for encoder to use as context (H in the paper)
          self.pre_decoder = get_decoder_model(
              config=config,
              arch="retro",
              hidden_size=hidden_size,
              ffn_hidden_size=ffn_hidden_size,
              num_layers=pre_decoder_num_layers + 1,
              num_attention_heads=num_attention_heads,
              apply_query_key_layer_scaling=apply_query_key_layer_scaling,
              kv_channels=kv_channels,
              init_method=pre_decoder_init,
              scaled_init_method=pre_decoder_scaled_init,
              pre_process=pre_process,
              post_process=False,  # no need for post process
              init_method_std=init_method_std,
              hidden_dropout=hidden_dropout,
              attention_dropout=attention_dropout,
              precision=precision,
              fp32_residual_connection=fp32_residual_connection,
              activations_checkpoint_method=activations_checkpoint_method,
              activations_checkpoint_num_layers=activations_checkpoint_num_layers,
              activations_checkpoint_granularity=activations_checkpoint_granularity,
              layernorm_epsilon=layernorm_epsilon,
              bias_activation_fusion=bias_gelu_fusion,
              bias_dropout_add_fusion=bias_dropout_add_fusion,
              masked_softmax_fusion=masked_softmax_fusion,
              persist_layer_norm=persist_layer_norm,
              openai_gelu=openai_gelu,
              onnx_safe=onnx_safe,
              hidden_steps=hidden_steps,
              activation=activation,
              bias=bias,
              normalization=normalization,
              transformer_block_type=transformer_block_type,
              headscale=headscale,
              parent_model_type=ModelType.encoder_and_decoder,
              layer_type=pre_decoder_layer_types,
              chunk_size=chunk_size,
              layer_number_offset=0,
              normalize_attention_scores=normalize_attention_scores,
              turn_off_rop=megatron_lm_compatible,
              version=version,
          )
  
          # it is where the chunked cross attention happens
          self.post_decoder = get_decoder_model(
              config=config,
              arch="retro",
              hidden_size=hidden_size,
              ffn_hidden_size=ffn_hidden_size,
              num_layers=post_decoder_num_layers,
              num_attention_heads=num_attention_heads,
              apply_query_key_layer_scaling=apply_query_key_layer_scaling,
              kv_channels=kv_channels,
              init_method=post_decoder_init,
              scaled_init_method=post_decoder_scaled_init,
              pre_process=False,  # directly take the pre_decoder output, skip preprocess
              post_process=post_process,
              init_method_std=init_method_std,
              hidden_dropout=hidden_dropout,
              attention_dropout=attention_dropout,
              precision=precision,
              fp32_residual_connection=fp32_residual_connection,
              activations_checkpoint_method=activations_checkpoint_method,
              activations_checkpoint_num_layers=activations_checkpoint_num_layers,
              activations_checkpoint_granularity=activations_checkpoint_granularity,
              layernorm_epsilon=layernorm_epsilon,
              bias_activation_fusion=bias_gelu_fusion,
              bias_dropout_add_fusion=bias_dropout_add_fusion,
              masked_softmax_fusion=masked_softmax_fusion,
              persist_layer_norm=persist_layer_norm,
              openai_gelu=openai_gelu,
              onnx_safe=onnx_safe,
              hidden_steps=hidden_steps,
              activation=activation,
              bias=bias,
              normalization=normalization,
              headscale=headscale,
              transformer_block_type=transformer_block_type,
              parent_model_type=ModelType.encoder_and_decoder,
              layer_type=post_decoder_layer_types,
              chunk_size=chunk_size,
              layer_number_offset=pre_decoder_num_layers + 1,
              normalize_attention_scores=normalize_attention_scores,
              turn_off_rop=megatron_lm_compatible,
              version=version,
          )
          self._pre_decoder_key = "pre_decoder"
          self._post_decoder_key = "post_decoder"
  
      self.initialize_word_embeddings(
          init_method=init_method_normal(init_method_std),
          vocab_size=vocab_size,
          hidden_size=hidden_size,
      )
  
      if add_decoder and post_process:
          self.tokens_head = MuReadout(
              self.word_embeddings_weight().size(0), parallel_output
          )
          self._tokens_head_key = "tokens_head"
  
file_path: NeMo/nemo/collections/nlp/modules/common/megatron/retrieval_token_level_encoder_decoder.py
imports: import torch
  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_encoders import get_encoder_model
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.mup.layer import MuReadout
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, init_method_normal, scaled_init_method_normal
label_desc: Use image node_id 1 to create a new MegatronRetrievalTokenLevelEncoderDecoderModule object from inherited base classes: MegatronModule with example: obj = MegatronRetrievalTokenLevelEncoderDecoderModule(config, vocab_size, hidden_size, max_position_embeddings, num_attention_heads, ffn_hidden_size, apply_query_key_layer_scaling, kv_channels, num_tokentypes, parallel_output, pre_process, post_process, init_method_std, fp16_cross_entropy, megatron_amp_O2, hidden_dropout, attention_dropout, precision, fp32_residual_connection, activations_checkpoint_method, activations_checkpoint_num_layers, layernorm_epsilon, persist_layer_norm, bias_gelu_fusion, bias_dropout_add_fusion, masked_softmax_fusion, openai_gelu, activation, onnx_safe, bias, normalization, headscale, transformer_block_type, hidden_steps, add_encoder, add_decoder, chunk_size, enc_num_layers, dec_num_layers, enc_cross_attention, dec_cross_attention, add_position_embedding, tokenizer, normalize_attention_scores, activations_checkpoint_granularity, megatron_lm_compatible, version)