class_docstr: Model parallel Multi-headed attention.
  This performs the Multi-headed attention over multiple gpus.
  
  See "Megatron-LM: https://arxiv.org/pdf/1909.08053.pdf" for more details.
class_docstr_tok: ["Model","parallel","Multi-headed","attention",".","This","performs","the","Multi-headed","attention","over","multiple","gpus",".","See","``","Megatron-LM",":","https",":","\/\/arxiv.org\/pdf\/1909.08053.pdf","''","for","more","details","."]
code: def __init__(
      self,
      embed_dim,
      num_heads,
      kdim=None,
      vdim=None,
      dropout=0.0,
      bias=True,
      self_attention=False,
      encoder_decoder_attention=False,
  ):
      super().__init__()
      if not has_megatron_submodule:
          raise ImportError(
              "\n\nPlease install the megatron submodule:"
              "\n\n  git submodule update --init "
              "fairseq/model_parallel/megatron"
          )
      self.embed_dim = embed_dim
      self.kdim = kdim if kdim is not None else embed_dim
      self.vdim = vdim if vdim is not None else embed_dim
      self.qkv_same_dim = (
          self.kdim == embed_dim and self.vdim == embed_dim
      )
  
      self.model_parallel_size = get_model_parallel_world_size()
  
      self.num_heads_partition = num_heads // self.model_parallel_size
      assert (
          self.num_heads_partition * self.model_parallel_size
          == num_heads
      ), "Number of heads must be divisible by model parallel size"
  
      self.dropout_module = FairseqDropout(
          dropout, module_name=self.__class__.__name__
      )
      self.head_dim = embed_dim // num_heads
      assert (
          self.head_dim * num_heads == self.embed_dim
      ), "embed_dim must be divisible by num_heads"
      self.scaling = self.head_dim**-0.5
  
      self.self_attention = self_attention
      self.encoder_decoder_attention = encoder_decoder_attention
  
      assert (
          not self.self_attention or self.qkv_same_dim
      ), "Self-attention requires query, key and value to be of the same size"
  
      self.k_proj = ColumnParallelLinear(
          self.kdim, embed_dim, bias=bias, gather_output=False
      )
      self.v_proj = ColumnParallelLinear(
          self.vdim, embed_dim, bias=bias, gather_output=False
      )
      self.q_proj = ColumnParallelLinear(
          embed_dim, embed_dim, bias=bias, gather_output=False
      )
      self.out_proj = RowParallelLinear(
          embed_dim, embed_dim, bias=bias, input_is_parallel=True
      )
  
file_path: fairseq/fairseq/model_parallel/modules/multihead_attention.py
imports: from typing import Dict, Optional, Tuple
  import torch
  import torch.nn.functional
  from torch import Tensor, nn
  from fairseq import utils
  from fairseq.incremental_decoding_utils import with_incremental_state
  from fairseq.modules.fairseq_dropout import FairseqDropout
label_desc: Use image node_id 1 to create a new ModelParallelMultiheadAttention object from inherited base classes: nn with example: obj = ModelParallelMultiheadAttention(embed_dim, num_heads, kdim, vdim, dropout, bias, _attention, encoder_decoder_attention)