class_docstr: Region where selective activation recomputation is applied.
  See Figure 3. in Reducing Activation Recomputation in Large Transformer Models
  https://arxiv.org/pdf/2205.05198.pdf for more details.
class_docstr_tok: ["Region","where","selective","activation","recomputation","is","applied",".","See","Figure","3.","in","Reducing","Activation","Recomputation","in","Large","Transformer","Models","https",":","\/\/arxiv.org\/pdf\/2205.05198.pdf","for","more","details","."]
code: def __init__(
      self,
      config: ModelParallelConfig,
      layer_number,
      num_attention_heads,
      hidden_size,
      attention_type=AttnType.self_attn,
      attn_mask_type=AttnMaskType.padding,
      precision=16,
      apply_query_key_layer_scaling=False,
      kv_channels=None,
      masked_softmax_fusion=True,
      attention_dropout=0.1,
      normalize_attention_scores=True,
      multi_query_attention=False,
      position_embedding_type="learned_absolute",
      use_flash_attention=False,
  ):
      super(CoreAttention, self).__init__(config=config)
  
      self.precision = precision
      self.fp16 = False
      self.bf16 = False
      if precision in ["bf16", "bf16-mixed"]:
          self.bf16 = True
      elif precision in [16, "16", "16-mixed"]:
          self.fp16 = True
      self.multi_query_attention = multi_query_attention
      self.position_embedding_type = position_embedding_type
  
      self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
      self.attention_softmax_in_fp32 = False
      if self.apply_query_key_layer_scaling:
          self.attention_softmax_in_fp32 = True
      self.layer_number = max(1, layer_number)
      self.attention_type = attention_type
      self.attn_mask_type = attn_mask_type
      self.sequence_parallel = config.sequence_parallel
      # If True, will scale attention scores by 1 / sqrt(hidden_size_per_attention_head).
      # This arg is been provided mostly to support weight conversion of Huggingface models. (ex: T5v1.1)
      self.normalize_attention_scores = normalize_attention_scores
  
      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
  
      projection_size = kv_channels * num_attention_heads
  
      # Per attention head and per partition values.
      world_size = parallel_state.get_tensor_model_parallel_world_size()
      self.hidden_size_per_partition = safe_divide(
          projection_size, world_size
      )
      self.hidden_size_per_attention_head = safe_divide(
          projection_size, num_attention_heads
      )
      self.num_attention_heads_per_partition = safe_divide(
          num_attention_heads, world_size
      )
      self.num_attention_heads_partition_offset = (
          self.num_attention_heads_per_partition
          * parallel_state.get_tensor_model_parallel_rank()
      )
  
      coeff = None
      self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)
      if self.apply_query_key_layer_scaling:
          coeff = self.layer_number
          self.norm_factor *= coeff
  
      self.scale_mask_softmax = MatchedScaleMaskSoftmax(
          self.fp16,
          self.bf16,
          self.attn_mask_type,
          masked_softmax_fusion,
          attention_mask_func,
          self.attention_softmax_in_fp32,
          coeff,
      )
  
      # Dropout. Note that for a single iteration, this layer will generate
      # different outputs on different number of parallel partitions but
      # on average it should not be partition dependent.
      self.attention_dropout_p = attention_dropout
      self.attention_dropout = torch.nn.Dropout(attention_dropout)
  
      if use_flash_attention:
          self.attn_fn = self.flash_attention
      else:
          self.attn_fn = self.torch_attention
  
      if position_embedding_type.lower() == "xpos":
          self.xpos = XPOSPositionEmbedding(kv_channels)
  
imports: import math
  import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from nemo.collections.nlp.modules.common.megatron.adapters.parallel_adapters import AdapterName, InfusedAdapterConfig, LoraKQVAdapterConfig, LoraKQVAdapterWeightTyingConfig, LoraKVAdapterConfig, LoraQAdapterConfig
  from nemo.collections.nlp.modules.common.megatron.fused_softmax import MatchedScaleMaskSoftmax
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.position_embedding import XPOSPositionEmbedding
  from nemo.collections.nlp.modules.common.megatron.position_embedding.rotary_position_embedding import apply_rotary_pos_emb
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, _cast_if_autocast_enabled, attention_mask_func
  from nemo.core import adapter_mixins
label_desc: Use image node_id 1 to create a new CoreAttention object from inherited base classes: MegatronModule with example: obj = CoreAttention(config, layer_number, num_attention_heads, hidden_size, attention_type, attn_mask_type, precision, apply_query_key_layer_scaling, kv_channels, masked_softmax_fusion, attention_dropout, normalize_attention_scores, multi_query_attention, position_embedding_type, use_flash_attention)