code: def __init__(
      self,
      config,
      mp_group=None,
      q_scales=None,
      q_groups=1,
      merge_count=1,
      qkv_merging=False,
  ):
      super(DeepSpeedSelfAttention, self).__init__()
      self.config = config
      data_type = (
          torch.int8
          if config.q_int8
          else torch.half
          if config.fp16
          else torch.float
      )
      data_type_fp = torch.half if config.fp16 else torch.float
      self.config.layer_id = DeepSpeedSelfAttention.num_layers
      DeepSpeedSelfAttention.num_layers = (
          DeepSpeedSelfAttention.num_layers + 1
      )
      device = (
          torch.cuda.current_device()
      )  # if config.bigscience_bloom else 'cpu'
      qkv_size_per_partition = (
          self.config.hidden_size // self.config.mp_size
      ) * 3
      self.attn_qkvw = nn.Parameter(
          torch.empty(
              self.config.hidden_size,
              qkv_size_per_partition,
              dtype=data_type,
              device=device,
          ),
          requires_grad=False,
      )
      self.attn_qkvb = nn.Parameter(
          torch.empty(
              qkv_size_per_partition, dtype=data_type_fp, device=device
          ),
          requires_grad=False,
      )
      out_size_per_partition = (
          self.config.hidden_size // self.config.mp_size
      )
      self.attn_ow = nn.Parameter(
          torch.empty(
              out_size_per_partition,
              self.config.hidden_size,
              dtype=data_type,
              device=device,
          ),
          requires_grad=False,
      )
  
      self.attn_ob = nn.Parameter(
          torch.empty(
              self.config.hidden_size, dtype=data_type_fp, device=device
          ),
          requires_grad=False,
      )
  
      self.num_attention_heads_per_partition = (
          self.config.heads // self.config.mp_size
      )
      self.hidden_size_per_partition = (
          self.config.hidden_size // self.config.mp_size
      )
      self.hidden_size_per_attention_head = (
          self.config.hidden_size // self.config.heads
      )
  
      self.mp_group = mp_group
  
      # used for quantization
      self.q_scales = q_scales
      self.q_groups = q_groups
      self.merge_count = int(math.log2(merge_count))
  
      self.norm_factor = math.sqrt(
          math.sqrt(self.config.hidden_size // self.config.heads)
      )
      self.qkv_merging = qkv_merging
  
      if self.config.scale_attn_by_inverse_layer_idx is True:
          self.norm_factor *= math.sqrt(self.config.layer_id + 1)
          # https://github.com/huggingface/transformers/blob/v4.24.0/src/transformers/models/gpt2/modeling_gpt2.py#L191
  
      global inference_cuda_module
      if inference_cuda_module is None:
          builder = op_builder.InferenceBuilder()
          inference_cuda_module = builder.load()
  
      self.score_context_func = (
          inference_cuda_module.softmax_context_fp32
          if (not config.fp16)
          else inference_cuda_module.softmax_context_fp16
      )
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/deepspeed/ops/transformer/inference/ds_attention.py
imports: import math
  import torch
  from torch.autograd import Function
  from ...None import op_builder
  import torch.nn
  from deepspeed import comm
label_desc: Use image node_id 1 to create a new DeepSpeedSelfAttention object from inherited base classes: nn with example: obj = DeepSpeedSelfAttention(config, mp_group, q_scales, q_groups, merge_count, qkv_merging)