class_docstr: Transformer class.
class_docstr_tok: ["Transformer","class","."]
code: def __init__(
      self,
      config: ModelParallelConfig,
      init_method,
      output_layer_init_method,
      num_layers,
      hidden_size,
      ffn_hidden_size,
      num_attention_heads,
      apply_query_key_layer_scaling=False,
      kv_channels=None,
      layer_type=LayerType.encoder,  # it can be a list of types or single type
      self_attn_mask_type=AttnMaskType.padding,
      pre_process=True,
      post_process=True,
      precision=16,
      fp32_residual_connection=False,
      activations_checkpoint_method=None,
      activations_checkpoint_num_layers=None,
      layernorm_epsilon=1e-5,
      hidden_dropout=0.1,
      attention_dropout=0.1,
      ffn_dropout=0.0,
      bias_activation_fusion=True,
      bias_dropout_add_fusion=True,
      masked_softmax_fusion=True,
      persist_layer_norm=False,
      openai_gelu=False,
      onnx_safe=False,
      activation="gelu",
      model_type=ModelType.encoder_or_decoder,
      megatron_legacy=False,
      bias=True,
      chunk_size=64,
      normalization="layernorm",
      transformer_block_type="pre_ln",
      position_embedding_type="learned_absolute",
      headscale=False,
      layer_number_offset=0,  # this is use only for attention norm_factor scaling
      activations_checkpoint_granularity=None,
      activations_checkpoint_layers_per_pipeline=None,
      transformer_engine=False,
      fp8=False,
      fp8_e4m3=False,
      fp8_hybrid=False,
      fp8_margin=0,
      fp8_interval=1,
      fp8_amax_history_len=1024,
      fp8_amax_compute_algo="max",
      reduce_amax=True,
      use_emha=False,
      ub_tp_comm_overlap=False,
      normalize_attention_scores=True,
      multi_query_attention=False,
      num_moe_experts=1,
      moe_frequency=1,
      moe_dropout=0.0,
      use_flash_attention=False,
  ):
      super(ParallelTransformer, 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.fp32_residual_connection = fp32_residual_connection
      self.pre_process = pre_process
      self.post_process = post_process
      self.input_tensor = None
      self.self_attn_mask_type = self_attn_mask_type
      self.model_type = model_type
      self.normalization = normalization
      self.transformer_block_type = transformer_block_type
      self.layer_type = layer_type
      self.position_embedding_type = position_embedding_type
      self.multi_query_attention = multi_query_attention
  
      self.inference_current_sequence_len = 0
      self.inference_params = None
  
      self.activations_checkpoint_method = activations_checkpoint_method
      self.activations_checkpoint_num_layers = (
          activations_checkpoint_num_layers
      )
      self.activations_checkpoint_granularity = (
          activations_checkpoint_granularity
      )
      self.activations_checkpoint_layers_per_pipeline = (
          activations_checkpoint_layers_per_pipeline
      )
  
      if self.activations_checkpoint_granularity:
          if self.activations_checkpoint_granularity == "selective":
              if self.activations_checkpoint_method == "uniform":
                  logging.info(
                      (
                          f"Using uniform activation checkpointing with granularity selective forces all layers to use checkpointing."
                      )
                  )
              elif self.activations_checkpoint_method == "block":
                  logging.info(
                      (
                          f"Using block activation checkpointing with granularity selective forces all layers to use checkpointing."
                      )
                  )
              else:
                  raise ValueError(
                      f'activations_checkpoint_method should be "uniform" or "block" when using granularity selective.'
                  )
              self.activations_checkpoint_num_layers = (
                  num_layers  # forcing all layers
              )
          elif self.activations_checkpoint_granularity == "full":
              if self.activations_checkpoint_method in [
                  "uniform",
                  "block",
              ]:
                  if not self.activations_checkpoint_num_layers:
                      logging.info(
                          (
                              f"Using uniform or block activation checkpointing requires activations_checkpoint_num_layers to be set."
                              f"Got: {self.activations_checkpoint_num_layers}. Setting to 1 by default."
                          )
                      )
                      self.activations_checkpoint_num_layers = (
                          1  # keeping the old default
                      )
              else:
                  raise ValueError(
                      f'activations_checkpoint_method should be "uniform" or "block" when using granularity full.'
                  )
          else:
              raise ValueError(
                  f'activations_checkpoint_granularity should be "selective" or "full".'
              )
  
      self.sequence_parallel = config.sequence_parallel
      self.transformer_engine = transformer_engine
      self.fp8 = fp8
      self.fp8_e4m3 = fp8_e4m3
      self.fp8_hybrid = fp8_hybrid
      self.fp8_margin = fp8_margin
      self.fp8_interval = fp8_interval
      self.fp8_amax_history_len = fp8_amax_history_len
      self.fp8_amax_compute_algo = fp8_amax_compute_algo
      self.reduce_amax = reduce_amax
  
      self.fp8_recipe = None
  
      if self.fp8:
          if self.fp8_e4m3:
              fp8_format = recipe.Format.E4M3
          elif self.fp8_hybrid:
              fp8_format = recipe.Format.HYBRID
          self.fp8_recipe = recipe.DelayedScaling(
              margin=self.fp8_margin,
              interval=self.fp8_interval,
              fp8_format=fp8_format,
              amax_history_len=self.fp8_amax_history_len,
              amax_compute_algo=self.fp8_amax_compute_algo,
              reduce_amax=reduce_amax,
          )
  
      self.is_first_train_microbatch = True  # Is the current micro-batch the first micro-batch in a global-batch in training
      self.is_prev_microbatch_training = (
          True  # Is the previous micro-batch in training mode
      )
      self.microbatch_count = 0  # transformer engine forward needs to know if it is working on the first microbatch
      self.checkpoint_core_attention = (
          activations_checkpoint_granularity == "selective"
      )  # transformer engine forward allows for more granular selective checkpointing
  
      if self.model_type == ModelType.encoder_or_decoder:
          assert (
              num_layers
              % parallel_state.get_pipeline_model_parallel_world_size()
              == 0
          ), "num_layers must be divisible by pipeline_model_parallel_size"
  
      assert (
          moe_frequency <= num_layers
      ), "MoE frequency must be <= number of transformer layers"
      # TODO: Add similar assert for encoder-decoder.
  
      self.num_layers = self.get_num_layers(num_layers)
  
      if (
          self.activations_checkpoint_num_layers is not None
          and self.activations_checkpoint_num_layers > self.num_layers
      ):
          self.activations_checkpoint_num_layers = self.num_layers
  
      # Transformer layers.
      def build_layer(layer_number):
          if isinstance(layer_type, list):
              lt = layer_type[layer_number - 1]
          else:
              lt = layer_type
  
          if self.transformer_engine:
              return AutocastTransformerLayer(
                  hidden_size=hidden_size,
                  ffn_hidden_size=ffn_hidden_size,
                  layernorm_epsilon=layernorm_epsilon,
                  num_attention_heads=num_attention_heads,
                  init_method=init_method,
                  output_layer_init_method=output_layer_init_method,
                  hidden_dropout=hidden_dropout,
                  attention_dropout=attention_dropout,
                  layer_number=layer_number + layer_number_offset,
                  kv_channels=kv_channels,
                  self_attn_mask_type=self_attn_mask_type.name,
                  tp_size=parallel_state.get_tensor_model_parallel_world_size(),
                  params_dtype=config.params_dtype,
                  get_rng_state_tracker=tensor_parallel.random.get_cuda_rng_tracker,
                  fuse_wgrad_accumulation=config.gradient_accumulation_fusion,
                  seq_length=None,  # used for jit warmup
                  micro_batch_size=None,  # used for jit warmup
                  sequence_parallel=config.sequence_parallel,
                  apply_residual_connection_post_layernorm=False,
                  autocast_dtype=precision,
                  use_emha=use_emha,
                  ub_tp_comm_overlap=ub_tp_comm_overlap,
                  zero_centered_gamma=normalization == "layernorm1p",
                  device="cpu"
                  if config.use_cpu_initialization
                  else "cuda",
              )
          else:
              return ParallelTransformerLayer(
                  config=config,
                  init_method=init_method,
                  output_layer_init_method=output_layer_init_method,
                  layer_number=layer_number + layer_number_offset,
                  hidden_size=hidden_size,
                  ffn_hidden_size=ffn_hidden_size,
                  num_attention_heads=num_attention_heads,
                  apply_query_key_layer_scaling=apply_query_key_layer_scaling,
                  kv_channels=kv_channels,
                  layer_type=lt,
                  self_attn_mask_type=self_attn_mask_type,
                  precision=precision,
                  fp32_residual_connection=fp32_residual_connection,
                  layernorm_epsilon=layernorm_epsilon,
                  hidden_dropout=hidden_dropout,
                  attention_dropout=attention_dropout,
                  ffn_dropout=ffn_dropout,
                  bias_activation_fusion=bias_activation_fusion,
                  bias_dropout_add_fusion=bias_dropout_add_fusion,
                  masked_softmax_fusion=masked_softmax_fusion,
                  persist_layer_norm=persist_layer_norm,
                  position_embedding_type=position_embedding_type,
                  openai_gelu=openai_gelu,
                  onnx_safe=onnx_safe,
                  activation=activation,
                  megatron_legacy=megatron_legacy,
                  bias=bias,
                  chunk_size=chunk_size,
                  normalization=normalization,
                  transformer_block_type=transformer_block_type,
                  headscale=headscale,
                  activations_checkpoint_granularity=activations_checkpoint_granularity,
                  normalize_attention_scores=normalize_attention_scores,
                  num_moe_experts=num_moe_experts,
                  moe_frequency=moe_frequency,
                  moe_dropout=moe_dropout,
                  use_flash_attention=use_flash_attention,
              )
  
      if (
          parallel_state.get_virtual_pipeline_model_parallel_world_size()
          is not None
      ):
          assert (
              num_layers
              % parallel_state.get_virtual_pipeline_model_parallel_world_size()
              == 0
          ), (
              "num_layers_per_stage must be divisible by "
              "virtual_pipeline_model_parallel_size"
          )
  
          assert (
              self.model_type.value != 2
          ), f"virtual pipeline parallel currently only supported for GPT"
  
          # Number of layers in each model chunk is the number of layers in the stage,
          # divided by the number of model chunks in a stage.
          self.num_layers = (
              self.num_layers
              // parallel_state.get_virtual_pipeline_model_parallel_world_size()
          )
          # With 8 layers, 2 stages, and 4 model chunks, we want an assignment of
          # layers to stages like (each list is a model chunk):
          # Stage 0: [0]  [2]  [4]  [6]
          # Stage 1: [1]  [3]  [5]  [7]
          # With 8 layers, 2 stages, and 2 virtual stages, we want an assignment of
          # layers to stages like (each list is a model chunk):
          # Stage 0: [0, 1]  [4, 5]
          # Stage 1: [2, 3]  [6, 7]
          offset = parallel_state.get_virtual_pipeline_model_parallel_rank() * (
              num_layers
              // parallel_state.get_virtual_pipeline_model_parallel_world_size()
          ) + (
              parallel_state.get_pipeline_model_parallel_rank()
              * self.num_layers
          )
      else:
          # Each stage gets a contiguous set of layers.
          if (
              self.model_type == ModelType.encoder_and_decoder
              and parallel_state.get_pipeline_model_parallel_world_size()
              > 1
          ):
              pipeline_rank = (
                  parallel_state.get_pipeline_model_parallel_rank()
              )
              if layer_type == LayerType.encoder:
                  offset = pipeline_rank * self.num_layers
              else:
                  num_ranks_in_enc = (
                      parallel_state.get_pipeline_model_parallel_split_rank()
                  )
                  offset = (
                      pipeline_rank - num_ranks_in_enc
                  ) * self.num_layers
          else:
              offset = (
                  parallel_state.get_pipeline_model_parallel_rank()
                  * self.num_layers
              )
  
      self.layers = torch.nn.ModuleList(
          [build_layer(i + 1 + offset) for i in range(self.num_layers)]
      )
  
      if self.post_process and self.transformer_block_type != "post_ln":
          # Final layer norm before output.
          if normalization == "layernorm":
              self.final_layernorm = get_layer_norm(
                  hidden_size,
                  layernorm_epsilon,
                  persist_layer_norm,
                  sequence_parallel=config.sequence_parallel,
              )
          elif normalization == "layernorm1p":
              self.final_layernorm = LayerNorm1P(
                  hidden_size,
                  layernorm_epsilon,
                  sequence_parallel_enabled=config.sequence_parallel,
              )
          elif normalization == "low_precision_layernorm":
              self.final_layernorm = LPLayerNorm(
                  hidden_size, layernorm_epsilon
              )
          else:
              self.final_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.final_layernorm)
  
      # Hacky set up for vision encoder select layer, won't support PP
      # It indicates the layer number of hidden states that we want to return.
      # For example -2 means we skip the last layer in the decoder, and return at -2 layer.
      self.return_select_layer = 0
  
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 ParallelTransformer object from inherited base classes: MegatronModule with example: obj = ParallelTransformer(config, init_method, output_layer_init_method, num_layers, hidden_size, ffn_hidden_size, num_attention_heads, apply_query_key_layer_scaling, kv_channels, layer_type, _attn_mask_type, pre_process, post_process, precision, fp32_residual_connection, activations_checkpoint_method, activations_checkpoint_num_layers, layernorm_epsilon, hidden_dropout, attention_dropout, ffn_dropout, bias_activation_fusion, bias_dropout_add_fusion, masked_softmax_fusion, persist_layer_norm, openai_gelu, onnx_safe, activation, model_type, megatron_legacy, bias, chunk_size, normalization, transformer_block_type, position_embedding_type, headscale, layer_number_offset, activations_checkpoint_granularity, activations_checkpoint_layers_per_pipeline, transformer_engine, fp8, fp8_e4m3, fp8_hybrid, fp8_margin, fp8_interval, fp8_amax_history_len, fp8_amax_compute_algo, reduce_amax, use_emha, ub_tp_comm_overlap, normalize_attention_scores, multi_query_attention, num_moe_experts, moe_frequency, moe_dropout, use_flash_attention)