class_docstr: Region where selective activation recomputation is applied.
  This region is memory intensive but less compute intensive which
  makes activation checkpointing more efficient for LLMs (20B+).
  See Reducing Activation Recomputation in Large Transformer Models: https://arxiv.org/abs/2205.05198 for more details.
  
  We use the following notation:
   h: hidden size
   n: number of attention heads
   p: number of tensor model parallel partitions
   b: batch size
   s: sequence length
class_docstr_tok: ["Region","where","selective","activation","recomputation","is","applied",".","This","region","is","memory","intensive","but","less","compute","intensive","which","makes","activation","checkpointing","more","efficient","for","LLMs","(","20B+",")",".","See","Reducing","Activation","Recomputation","in","Large","Transformer","Models",":","https",":","\/\/arxiv.org\/abs\/2205.05198","for","more","details",".","We","use","the","following","notation",":","h",":","hidden","size","n",":","number","of","attention","heads","p",":","number","of","tensor","model","parallel","partitions","b",":","batch","size","s",":","sequence","length"]
code: def __init__(
      self,
      config: TransformerConfig,
      layer_number: int,
      attn_mask_type: AttnMaskType,
      attention_type: str,
      attention_dropout: float = None,
  ):
      super().__init__(config=config)
  
      self.config: TransformerConfig = config
  
      assert (
          self.config.context_parallel_size == 1
      ), "Context parallelism is only supported by TEDotProductAttention!"
  
      assert (
          self.config.window_size is None
      ), "Sliding Window Attention is only supported by TEDotProductAttention!"
  
      self.layer_number = max(1, layer_number)
      self.attn_mask_type = attn_mask_type
      self.attention_type = attention_type  # unused for now
  
      projection_size = (
          self.config.kv_channels * self.config.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 = divide(
          projection_size, world_size
      )
      self.hidden_size_per_attention_head = divide(
          projection_size, config.num_attention_heads
      )
      self.num_attention_heads_per_partition = divide(
          self.config.num_attention_heads, world_size
      )
      self.num_query_groups_per_partition = divide(
          self.config.num_query_groups, world_size
      )
  
      coeff = None
      self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)
      if self.config.apply_query_key_layer_scaling:
          coeff = self.layer_number
          self.norm_factor *= coeff
  
      self.scale_mask_softmax = FusedScaleMaskSoftmax(
          input_in_fp16=self.config.fp16,
          input_in_bf16=self.config.bf16,
          attn_mask_type=self.attn_mask_type,
          scaled_masked_softmax_fusion=self.config.masked_softmax_fusion,
          mask_func=attention_mask_func,
          softmax_in_fp32=self.config.attention_softmax_in_fp32,
          scale=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 = torch.nn.Dropout(
          self.config.attention_dropout
          if attention_dropout is None
          else attention_dropout
      )
  
file_path: megatron-lm/megatron/core/transformer/dot_product_attention.py
imports: import math
  import torch
  from torch import Tensor
  from megatron.core import parallel_state, tensor_parallel
  from megatron.core.fusions.fused_softmax import FusedScaleMaskSoftmax
  from megatron.core.transformer.enums import AttnMaskType
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.transformer.utils import attention_mask_func
  from megatron.core.utils import divide
label_desc: Use image node_id 1 to create a new DotProductAttention object from inherited base classes: MegatronModule with example: obj = DotProductAttention(config, layer_number, attn_mask_type, attention_type, attention_dropout)