code: def __init__(
      self,
      config,
      mp_group=None,
      q_scales=None,
      q_groups=1,
      merge_count=1,
      mlp_extra_grouping=False,
  ):
      super(DeepSpeedMLP, 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
      device = (
          torch.cuda.current_device()
      )  # if config.bigscience_bloom else 'cpu'
      self.attn_nw = nn.Parameter(
          torch.empty(
              self.config.hidden_size, dtype=data_type_fp, device=device
          ),
          requires_grad=False,
      )
      self.attn_nb = nn.Parameter(
          torch.empty(
              self.config.hidden_size, dtype=data_type_fp, device=device
          ),
          requires_grad=False,
      )
      intm_size_per_partition = (
          self.config.intermediate_size // self.config.mp_size
      )
      self.inter_w = nn.Parameter(
          torch.empty(
              self.config.hidden_size,
              intm_size_per_partition,
              dtype=data_type,
              device=device,
          ),
          requires_grad=False,
      )
      self.inter_b = nn.Parameter(
          torch.empty(
              intm_size_per_partition, dtype=data_type_fp, device=device
          ),
          requires_grad=False,
      )
      self.output_w = nn.Parameter(
          torch.empty(
              intm_size_per_partition,
              self.config.hidden_size,
              dtype=data_type,
              device=device,
          ),
          requires_grad=False,
      )
      self.output_b = nn.Parameter(
          torch.empty(
              self.config.hidden_size, dtype=data_type_fp, device=device
          ),
          requires_grad=False,
      )
  
      # used for quantization
      self.q_scales = q_scales
      self.q_groups = q_groups * 2 if mlp_extra_grouping else q_groups
      self.merge_count = int(math.log2(merge_count))
  
      # load the cuda module
      global inference_cuda_module
      if inference_cuda_module is None:
          builder = op_builder.InferenceBuilder()
          inference_cuda_module = builder.load()
  
      self.mp_group = mp_group
      self.mlp_gemm_func = (
          inference_cuda_module.mlp_gemm_fp16
          if config.fp16
          else inference_cuda_module.mlp_gemm_fp32
      )
      self.vector_matmul_func = (
          inference_cuda_module.vector_matmul_fp16
          if config.fp16
          else inference_cuda_module.vector_matmul_fp32
      )
      self.fused_gemm_gelu = (
          inference_cuda_module.fused_gemm_gelu_fp16
          if config.fp16
          else inference_cuda_module.fused_gemm_gelu_fp32
      )
  
      self.bias_residual_func = (
          inference_cuda_module.bias_residual_fp16
          if config.fp16 or config.q_int8
          else inference_cuda_module.bias_residual_fp32
      )
  
      self.residual_add_func = (
          inference_cuda_module.residual_add_bias_fp16
          if config.fp16 or config.q_int8
          else inference_cuda_module.residual_add_bias_fp32
      )
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/deepspeed/ops/transformer/inference/ds_mlp.py
imports: import torch
  from torch.autograd import Function
  from deepspeed.utils.types import ActivationFuncType
  from deepspeed import comm
  import torch.nn
  import math
  from ...None import op_builder
label_desc: Use image node_id 1 to create a new DeepSpeedMLP object from inherited base classes: nn with example: obj = DeepSpeedMLP(config, mp_group, q_scales, q_groups, merge_count, mlp_extra_grouping)