class_docstr: Initialize the DeepSpeed Transformer Layer.
  Arguments:
      layer_id: The layer index starting from 0, e.g. if model has 24 transformer layers,
          layer_id will be 0,1,2...23 when each layer object is instantiated
      config: An object of DeepSpeedInferenceConfig
class_docstr_tok: ["Initialize","the","DeepSpeed","Transformer","Layer",".","Arguments",":","layer_id",":","The","layer","index","starting","from","0",",","e.g",".","if","model","has","24","transformer","layers",",","layer_id","will","be","0,1,2","...","23","when","each","layer","object","is","instantiated","config",":","An","object","of","DeepSpeedInferenceConfig"]
code: def __init__(
      self,
      config,
  ):
      super(DeepSpeedDiffusersAttention, self).__init__()
  
      self.config = config
      self.config.layer_id = DeepSpeedDiffusersAttention.layer_id
      DeepSpeedDiffusersAttention.layer_id += 1
      device = (
          get_accelerator().current_device_name()
          if config.bigscience_bloom
          else "cpu"
      )
      qkv_size_per_partition = (
          self.config.hidden_size // self.config.mp_size
      ) * 3
  
      data_type = self.config.dtype
      data_type_fp = (
          torch.half
          if self.config.dtype == torch.int8
          else self.config.dtype
      )
      global inference_module
      if inference_module is None:
          builder = InferenceBuilder()
          inference_module = builder.load()
  
      if DeepSpeedDiffusersAttention.layer_id == 1:
          log_dist(
              f"DeepSpeed-Attention config: {self.config.__dict__}", [0]
          )
  
      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_kw = nn.Parameter(
          torch.empty(
              self.config.hidden_size,
              self.config.hidden_size,
              dtype=data_type,
              device=device,
          ),
          requires_grad=False,
      )
      self.attn_vw = nn.Parameter(
          torch.empty(
              self.config.hidden_size,
              self.config.hidden_size,
              dtype=data_type,
              device=device,
          ),
          requires_grad=False,
      )
      self.attn_qw = nn.Parameter(
          torch.empty(
              self.config.hidden_size,
              self.config.hidden_size,
              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.do_out_bias = True
  
      if triton_flash_attn is None:
          load_triton_flash_attn()
      self.triton_flash_attn_kernel = triton_flash_attn()
      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.norm_factor = math.sqrt(
          math.sqrt(self.config.hidden_size // self.config.heads)
      )
  
      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
  
      if self.config.dtype in [torch.float16, torch.int8]:
          self.score_context_func = (
              inference_module.softmax_context_fp16
          )
          self.linear_func = inference_module.linear_layer_fp16
          self.allocate_workspace = (
              inference_module.allocate_workspace_fp16
          )
      else:
          self.score_context_func = (
              inference_module.softmax_context_fp32
          )
          self.linear_func = inference_module.linear_layer_fp32
          self.allocate_workspace = (
              inference_module.allocate_workspace_fp32
          )
  
file_path: DeepSpeed/deepspeed/ops/transformer/inference/diffusers_attention.py
imports: import math
  import torch
  from torch.autograd import Function
  import torch.nn
  from packaging import version
  from deepspeed.utils.logging import log_dist
  from deepspeed.accelerator import get_accelerator
  from deepspeed.ops.op_builder import InferenceBuilder
label_desc: Use image node_id 1 to create a new DeepSpeedDiffusersAttention object from inherited base classes: nn with example: obj = DeepSpeedDiffusersAttention(config)