class_docstr: A single transformer layer.
  
  Transformer layer takes input with size [s, b, h] and returns an
  output of the same size.
class_docstr_tok: ["A","single","transformer","layer",".","Transformer","layer","takes","input","with","size","[","s",",","b",",","h","]","and","returns","an","output","of","the","same","size","."]
code: def __init__(
      self,
      config: ModelParallelConfig,
      init_method,
      output_layer_init_method,
      layer_number,
      hidden_size,
      ffn_hidden_size,
      num_attention_heads,
      layer_type=LayerType.encoder,
      self_attn_mask_type=AttnMaskType.padding,
      fp32_residual_connection=False,
      precision=16,
      apply_query_key_layer_scaling=False,
      kv_channels=None,
      layernorm_epsilon=1e-5,
      hidden_dropout=0.1,
      persist_layer_norm=False,
      bias_activation_fusion=True,
      bias_dropout_add_fusion=True,
      masked_softmax_fusion=True,
      openai_gelu=False,
      onnx_safe=False,
      attention_dropout=0.1,
      ffn_dropout=0.0,
      activation="gelu",
      megatron_legacy=False,
      bias=True,
      chunk_size=64,
      normalization="layernorm",
      transformer_block_type="pre_ln",
      position_embedding_type="learned_absolute",
      multi_query_attention=False,
      headscale=False,
      activations_checkpoint_granularity=None,
      normalize_attention_scores=True,
      num_moe_experts=1,
      moe_frequency=1,
      moe_dropout=0.0,
      use_flash_attention=False,
  ):
      super(ParallelTransformerLayer_, self).__init__(config=config)
  
      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
  
      self.layer_number = layer_number
      self.layer_type = layer_type
      self.bias = bias
      self.transformer_block_type = transformer_block_type
      self.position_embedding_type = position_embedding_type
  
      self.set_accepted_adapter_types(
          [
              LinearAdapterConfig._target_,
              ParallelLinearAdapterConfig._target_,
              ParallelLinearAdapterWeightTyingConfig._target_,
          ]
      )
  
      if not bias and bias_dropout_add_fusion:
          raise ValueError(
              "bias_dropout_add_fusion=True requires bias=True, found bias=False. Either set both to True or both to False."
          )
  
      # the low_precision_layernorm does not require a bias term, whereas layernorm1p from apex
      # does require a bias, so it cannot be used for bias-less low precision LN such as in MPT-7B
      if normalization not in [
          "layernorm",
          "layernorm1p",
          "rmsnorm",
          "low_precision_layernorm",
      ]:
          raise ValueError(
              f'normalization must be "layernorm", "layernorm1p" or "rmsnorm", found {normalization}'
          )
  
      if transformer_block_type not in [
          "pre_ln",
          "post_ln",
          "normformer",
      ]:
          raise ValueError(
              f'transformer_block_type must be either "pre_ln" or "post_ln" or "normformer", found {transformer_block_type}'
          )
  
      self.fp32_residual_connection = fp32_residual_connection  # if true move residual connections to fp32
      self.hidden_dropout = hidden_dropout
      self.attention_dropout = attention_dropout
      self.bias_dropout_add_fusion = (
          bias_dropout_add_fusion  # if true, enable bias dropout fusion
      )
  
      # Self attention.
      # retrieval_decoder_after_self_attn skips the self attention
      if self.layer_type != LayerType.retrieval_decoder_after_self_attn:
          # Layernorm on the input data.
          if normalization == "layernorm":
              self.input_layernorm = get_layer_norm(
                  hidden_size,
                  layernorm_epsilon,
                  persist_layer_norm,
                  config.sequence_parallel,
              )
          elif normalization == "layernorm1p":
              self.input_layernorm = LayerNorm1P(
                  hidden_size,
                  layernorm_epsilon,
                  sequence_parallel_enabled=config.sequence_parallel,
              )
          elif normalization == "low_precision_layernorm":
              self.input_layernorm = LPLayerNorm(
                  hidden_size, layernorm_epsilon
              )
          else:
              self.input_layernorm = MixedFusedRMSNorm(
                  hidden_size, layernorm_epsilon
              )
          # for architectures such as MPT, there is no bias term even on the layernorms
          # this code allows us to remove the bias terms from the layernorm module
          # so that we can support MPT. However, certain apex-based LNs don't support
          # removing bias, so we also have to check for that
          if not bias and normalization not in [
              "layernorm",
              "layernorm1p",
          ]:
              remove_bias_from_layernorm(self.input_layernorm)
  
          self.self_attention = ParallelAttention(
              config=config,
              init_method=init_method,
              output_layer_init_method=output_layer_init_method,
              layer_number=layer_number,
              num_attention_heads=num_attention_heads,
              hidden_size=hidden_size,
              attention_type=AttnType.self_attn,
              attn_mask_type=self_attn_mask_type,
              precision=precision,
              apply_query_key_layer_scaling=apply_query_key_layer_scaling,
              kv_channels=kv_channels,
              masked_softmax_fusion=masked_softmax_fusion,
              attention_dropout=attention_dropout,
              multi_query_attention=multi_query_attention,
              layer_type=layer_type,
              megatron_legacy=megatron_legacy,
              bias=bias,
              headscale=headscale,
              position_embedding_type=position_embedding_type,
              normalize_attention_scores=normalize_attention_scores,
              use_flash_attention=use_flash_attention,
          )
  
          if transformer_block_type == "normformer":
              if normalization == "layernorm":
                  self.post_attention_normformer_norm = get_layer_norm(
                      hidden_size, layernorm_epsilon, persist_layer_norm
                  )
              else:
                  self.post_attention_normformer_norm = (
                      MixedFusedRMSNorm(hidden_size, layernorm_epsilon)
                  )
  
          if (
              self.layer_type != LayerType.decoder_pre_mlp
              or self.transformer_block_type != "post_ln"
          ):
              #  the post_attention_layernorm is used for layermorm after mlp
              # don't need it for decoder_pre_mlp and post_ln
              if normalization == "layernorm":
                  self.post_attention_layernorm = get_layer_norm(
                      hidden_size,
                      layernorm_epsilon,
                      persist_layer_norm,
                      config.sequence_parallel,
                  )
              elif normalization == "layernorm1p":
                  self.post_attention_layernorm = LayerNorm1P(
                      hidden_size,
                      layernorm_epsilon,
                      sequence_parallel_enabled=config.sequence_parallel,
                  )
              elif normalization == "low_precision_layernorm":
                  self.post_attention_layernorm = LPLayerNorm(
                      hidden_size, layernorm_epsilon
                  )
              else:
                  self.post_attention_layernorm = MixedFusedRMSNorm(
                      hidden_size, layernorm_epsilon
                  )
              if not bias and normalization not in [
                  "layernorm",
                  "layernorm1p",
              ]:
                  remove_bias_from_layernorm(
                      self.post_attention_layernorm
                  )
  
      if self.layer_type == LayerType.decoder_pre_mlp:
          # skip MLP and cross attention
          return
  
      # the post_attention_layernorm is used for layermorm after mlp
      # need it for post_ln
      if (
          self.layer_type == LayerType.retrieval_decoder_after_self_attn
          and self.transformer_block_type == "post_ln"
      ):
          # Layernorm on the attention output
          if normalization == "layernorm":
              self.post_attention_layernorm = get_layer_norm(
                  hidden_size,
                  layernorm_epsilon,
                  persist_layer_norm,
                  config.sequence_parallel,
              )
          elif normalization == "layernorm1p":
              self.post_attention_layernorm = LayerNorm1P(
                  hidden_size,
                  layernorm_epsilon,
                  sequence_parallel_enabled=config.sequence_parallel,
              )
          elif normalization == "low_precision_layernorm":
              self.post_attention_layernorm = LPLayerNorm(
                  hidden_size, layernorm_epsilon
              )
          else:
              self.post_attention_layernorm = MixedFusedRMSNorm(
                  hidden_size, layernorm_epsilon
              )
          if not bias and normalization not in [
              "layernorm",
              "layernorm1p",
          ]:
              remove_bias_from_layernorm(self.post_attention_layernorm)
  
      if (
          self.layer_type == LayerType.decoder
          or self.layer_type == LayerType.retrieval_encoder
      ):
          self.inter_attention = ParallelAttention(
              config=config,
              init_method=init_method,
              output_layer_init_method=output_layer_init_method,
              layer_number=layer_number,
              num_attention_heads=num_attention_heads,
              hidden_size=hidden_size,
              attention_type=AttnType.cross_attn,
              attn_mask_type=AttnMaskType.padding,
              precision=precision,
              apply_query_key_layer_scaling=apply_query_key_layer_scaling,
              kv_channels=kv_channels,
              multi_query_attention=multi_query_attention,
              masked_softmax_fusion=masked_softmax_fusion,
              attention_dropout=attention_dropout,
              megatron_legacy=megatron_legacy,
              bias=bias,
              headscale=headscale,
              normalize_attention_scores=normalize_attention_scores,
          )
          # Normformer normalization
          if transformer_block_type == "normformer":
              if normalization == "layernorm":
                  self.post_inter_attention_normformer_norm = (
                      get_layer_norm(
                          hidden_size,
                          layernorm_epsilon,
                          persist_layer_norm,
                          config.sequence_parallel,
                      )
                  )
              elif normalization == "layernorm1p":
                  self.post_inter_attention_normformer_norm = LayerNorm1P(
                      hidden_size,
                      layernorm_epsilon,
                      sequence_parallel_enabled=config.sequence_parallel,
                  )
              else:
                  self.post_inter_attention_normformer_norm = (
                      MixedFusedRMSNorm(hidden_size, layernorm_epsilon)
                  )
  
          # Layernorm on the attention output.
          if normalization == "layernorm":
              self.post_inter_attention_layernorm = get_layer_norm(
                  hidden_size,
                  layernorm_epsilon,
                  persist_layer_norm,
                  config.sequence_parallel,
              )
          elif normalization == "layernorm1p":
              self.post_inter_attention_layernorm = LayerNorm1P(
                  hidden_size,
                  layernorm_epsilon,
                  sequence_parallel_enabled=config.sequence_parallel,
              )
          else:
              self.post_inter_attention_layernorm = MixedFusedRMSNorm(
                  hidden_size, layernorm_epsilon
              )
      elif (
          self.layer_type == LayerType.retrieval_decoder
          or self.layer_type
          == LayerType.retrieval_decoder_after_self_attn
      ):
          self.inter_attention = ParallelChunkedCrossAttention(
              config=config,
              init_method=init_method,
              output_layer_init_method=output_layer_init_method,
              layer_number=layer_number,
              num_attention_heads=num_attention_heads,
              hidden_size=hidden_size,
              precision=precision,
              apply_query_key_layer_scaling=apply_query_key_layer_scaling,
              kv_channels=kv_channels,
              masked_softmax_fusion=masked_softmax_fusion,
              attention_dropout=attention_dropout,
              megatron_legacy=megatron_legacy,
              chunk_size=chunk_size,
              bias=bias,
              headscale=headscale,
          )
          # Normformer normalization
          if transformer_block_type == "normformer":
              if normalization == "layernorm":
                  self.post_inter_attention_normformer_norm = (
                      get_layer_norm(
                          hidden_size,
                          layernorm_epsilon,
                          persist_layer_norm,
                          config.sequence_parallel,
                      )
                  )
              elif normalization == "layernorm1p":
                  self.post_inter_attention_normformer_norm = LayerNorm1P(
                      hidden_size,
                      layernorm_epsilon,
                      sequence_parallel_enabled=config.sequence_parallel,
                  )
              else:
                  self.post_inter_attention_normformer_norm = (
                      MixedFusedRMSNorm(hidden_size, layernorm_epsilon)
                  )
  
          # Layernorm on the attention output.
          if normalization == "layernorm":
              self.post_inter_attention_layernorm = get_layer_norm(
                  hidden_size,
                  layernorm_epsilon,
                  persist_layer_norm,
                  config.sequence_parallel,
              )
          elif normalization == "layernorm1p":
              self.post_inter_attention_layernorm = LayerNorm1P(
                  hidden_size,
                  layernorm_epsilon,
                  sequence_parallel_enabled=config.sequence_parallel,
              )
          else:
              self.post_inter_attention_layernorm = MixedFusedRMSNorm(
                  hidden_size, layernorm_epsilon
              )
  
      # MLP
      if num_moe_experts > 1 and self.layer_number % moe_frequency == 0:
          self.mlp = SwitchMLP(
              config=config,
              num_experts=num_moe_experts,
              init_method=init_method,
              output_layer_init_method=output_layer_init_method,
              hidden_size=hidden_size,
              ffn_hidden_size=ffn_hidden_size,
              bias_activation_fusion=bias_activation_fusion,
              openai_gelu=openai_gelu,
              onnx_safe=onnx_safe,
              activation=activation,
              bias=bias,
              transformer_block_type=transformer_block_type,
              normalization=normalization,
              layernorm_epsilon=layernorm_epsilon,
              persist_layer_norm=persist_layer_norm,
              dropout=moe_dropout,
          )
      else:
          self.mlp = ParallelMLP(
              config=config,
              init_method=init_method,
              output_layer_init_method=output_layer_init_method,
              hidden_size=hidden_size,
              ffn_hidden_size=ffn_hidden_size,
              bias_activation_fusion=bias_activation_fusion,
              openai_gelu=openai_gelu,
              onnx_safe=onnx_safe,
              activation=activation,
              bias=bias,
              transformer_block_type=transformer_block_type,
              normalization=normalization,
              layernorm_epsilon=layernorm_epsilon,
              persist_layer_norm=persist_layer_norm,
              dropout=ffn_dropout,
          )
  
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 1 to create a new ParallelTransformerLayer_ object from inherited base classes: MegatronModule, adapter_mixins with example: obj = ParallelTransformerLayer_(config, init_method, output_layer_init_method, layer_number, hidden_size, ffn_hidden_size, num_attention_heads, layer_type, _attn_mask_type, fp32_residual_connection, precision, apply_query_key_layer_scaling, kv_channels, layernorm_epsilon, hidden_dropout, persist_layer_norm, bias_activation_fusion, bias_dropout_add_fusion, masked_softmax_fusion, openai_gelu, onnx_safe, attention_dropout, ffn_dropout, activation, megatron_legacy, bias, chunk_size, normalization, transformer_block_type, position_embedding_type, multi_query_attention, headscale, activations_checkpoint_granularity, normalize_attention_scores, num_moe_experts, moe_frequency, moe_dropout, use_flash_attention)