class_docstr: Initialize the MultiHeadAttention.
  
  Static variable:
  
      layer_id: The layer-index counter starting from 0 and incrementing by 1 every time a layer object is instantiated,
      e.g. if a model has 24 transformer layers, layer_id goes from 0 to 23.
  
  Arguments:
      hidden_size: Total dimension of hidden_size.
      nhead: Number of parallel attention heads.
      batch_size: Batch Size for one forward
      max_seq_len: Max length of input sequence
      dropout: Dropout probability
      norm_first: perform LayerNorms before attention
class_docstr_tok: ["Initialize","the","MultiHeadAttention",".","Static","variable",":","layer_id",":","The","layer-index","counter","starting","from","0","and","incrementing","by","1","every","time","a","layer","object","is","instantiated",",","e.g",".","if","a","model","has","24","transformer","layers",",","layer_id","goes","from","0","to","23",".","Arguments",":","hidden_size",":","Total","dimension","of","hidden_size",".","nhead",":","Number","of","parallel","attention","heads",".","batch_size",":","Batch","Size","for","one","forward","max_seq_len",":","Max","length","of","input","sequence","dropout",":","Dropout","probability","norm_first",":","perform","LayerNorms","before","attention"]
code: def reset_parameters(self):
      hs = self.config.hidden_size
  
      nn.init.zeros_(self.out_proj_bias)
  
      nn.init.ones_(self.norm_weight)
      nn.init.zeros_(self.norm_bias)
  
      if self.pg_size > 1:
          rank_in_pg = torch.distributed.get_rank(self.pg)
          attn_qkvw_global = torch.empty(hs * 3, hs)
          attn_qkvb_global = torch.empty(hs * 3)
          nn.init.xavier_uniform_(
              attn_qkvw_global, 1.0 / math.sqrt(2.0)
          )
          bound = self.calc_bound(attn_qkvw_global)
          nn.init.uniform_(attn_qkvb_global, -bound, bound)
  
          attn_qkvw_global = attn_qkvw_global.cuda()
          attn_qkvb_global = attn_qkvb_global.cuda()
          torch.distributed.broadcast(
              attn_qkvw_global, src=0, group=self.pg
          )
          torch.distributed.broadcast(
              attn_qkvb_global, src=0, group=self.pg
          )
          attn_qkvw_global = attn_qkvw_global.cpu()
          attn_qkvb_global = attn_qkvb_global.cpu()
  
          with torch.no_grad():
              self.in_proj_weight.copy_(
                  attn_qkvw_global.view(3, hs, hs)[
                      :,
                      int(hs * rank_in_pg / self.pg_size) : int(
                          hs * (rank_in_pg + 1) / self.pg_size
                      ),
                      :,
                  ]
              )
              self.in_proj_bias.copy_(
                  attn_qkvb_global.view(3, hs)[
                      :,
                      int(hs * rank_in_pg / self.pg_size) : int(
                          hs * (rank_in_pg + 1) / self.pg_size
                      ),
                  ]
              )
  
          attn_ow_global = torch.empty(hs, hs)
          nn.init.xavier_uniform_(attn_ow_global, 1.0)
          attn_ow_global = attn_ow_global.cuda()
          torch.distributed.broadcast(
              attn_ow_global, src=0, group=self.pg
          )
          attn_ow_global = attn_ow_global.cpu()
          with torch.no_grad():
              self.out_proj_weight.copy_(
                  attn_ow_global[
                      :,
                      int(hs * rank_in_pg / self.pg_size) : int(
                          hs * (rank_in_pg + 1) / self.pg_size
                      ),
                  ]
              )
  
      else:
          attn_qkvw = self.in_proj_weight.view(-1, hs)
          nn.init.xavier_uniform_(attn_qkvw, 1.0 / math.sqrt(2.0))
          bound = self.calc_bound(attn_qkvw)
          nn.init.uniform_(self.in_proj_bias, -bound, bound)
  
          nn.init.xavier_uniform_(self.out_proj_weight, 1.0)
  
file_path: ColossalAI/colossalai/kernel/cuda_native/multihead_attention.py
imports: import math
  from dataclasses import dataclass
  import torch
  from torch import nn
  from torch.autograd import Function
label_desc: Use image node_id 3 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.reset_parameters() without return types