docstr: Args:
      model: torch.nn.Module
      config: DeepSpeedInferenceConfig
docstr_tok: ["Args",":","model",":","torch.nn.Module","config",":","DeepSpeedInferenceConfig"]
code: def __init__(self, model, config):
      """
      Args:
          model: torch.nn.Module
          config: DeepSpeedInferenceConfig
      """
      global DS_INFERENCE_ENABLED
      DS_INFERENCE_ENABLED = True
  
      super().__init__()
  
      # Have to import here because inference_module is a global, but python
      # globals only work at the module level and will not be updated unless
      # we import it each time we init a new inference engine.
      from ..model_implementations.transformers.ds_transformer import (
          inference_module,
      )
  
      if inference_module is not None:
          self.destroy()
  
      self.module = model
      self._config = config
  
      self._get_model_config_generate(
          config
      )  # keep for weird backward compatibility
  
      # patch model generate with ours if model uses it
      if hasattr(self.module, "generate"):
          self.generate = self._generate
  
      if hasattr(self.module, "config"):
          TransformerPolicy.hf_model_config = self.module.config
  
      if (
          config.dtype == torch.half
          and not get_accelerator().is_fp16_supported()
      ):
          raise ValueError("Type fp16 is not supported.")
  
      # todo: keep this self.injection_dict because we don't use to change config.injection_policy API
      # todo: this will get changed when Molly's PR on auto injection dict is merged
      self.injection_dict = config.injection_policy
  
      # todo: refactor the mp_group and mp_size related in the next refactor
      self.mp_group = config.tensor_parallel.tp_group
      self.mpu = config.tensor_parallel.mpu
  
      # self._validate_args(self.mpu, config.replace_with_kernel_inject)
      self.quantize_merge_count = 1
      self.quantization_scales = None
  
      # these are not needed in the config as we are creating them ourselves in the inference engine
      self.ep_group = None  # config.moe.ep_group
      self.expert_mp_group = None  # config.moe.ep_mp_group
  
      self.cuda_graph_created = False
      self.checkpoint_engine = TorchCheckpointEngine()
      quantization_setting = None
      self._init_quantization_setting(
          quantization_setting
      )  # todo: update with the new quant config for weight quant
      self.model_profile_enabled = False
      self._model_times = []
  
      if not self.injection_dict and config.replace_with_kernel_inject:
          # This is a hack to remove the prepare_mask function on HF side for BLOOM architecture
          self.remove_mask_prepare_for_bloom()
  
      if self.injection_dict or not config.replace_with_kernel_inject:
          # This is a hack to redefine the alibi func due to TP
          if config.tensor_parallel.tp_size > 1:
              self.build_alibi_tensor()
              self.build_attn_bias()
  
      if (
          get_accelerator().device_name() == "cuda"
          and config.enable_cuda_graph
      ):
          assert pkg_version.parse(
              torch.__version__
          ) >= pkg_version.parse(
              "1.10"
          ), "If you want to use cuda graph, please upgrade torch to at least v1.10"
  
      # convert model to intended dtype
      if config.dtype:
          self._convert_to_dtype(config)
  
      if self.mpu:
          config.tensor_parallel.tp_size = dist.get_world_size(
              group=self.mpu.get_model_parallel_group()
          )
          self.mp_group = self.mpu.get_model_parallel_group()
      elif config.tensor_parallel.tp_size > 1:
          self._create_model_parallel_group(config)
          config.tensor_parallel.tp_group = self.mp_group
  
      if isinstance(self.module, torch.nn.Module):
          moe, _ = has_moe_layers(self.module)
      else:
          moe = False
  
      if moe and dist.get_world_size() > 1:
          self._create_ep_parallel_group(config.moe.moe_experts)
  
      # We only support three modes: 1) user specified policy for tensor-parallelism, 2) kernel injection (replace_with_kernel_inject), and 3) automatic tensor parallelism if tp_size > 1.
      if self.injection_dict:
          # 1. User specified Tensor Parallelism
          assert (
              not config.replace_with_kernel_inject
          ), "Cannot use both user specified injection policy and kernel injection"
          for (
              client_module,
              injection_policy,
          ) in self.injection_dict.items():
              assert issubclass(
                  client_module, torch.nn.Module
              ), f"{client_module} is not a subclass of torch.nn.Module"
  
              # construct the tuple and pass that instead of a string or dict.
              if isinstance(injection_policy, str):
                  config.injection_policy_tuple = (injection_policy,)
              else:
                  config.injection_policy_tuple = injection_policy
  
              layer_names = [
                  name for name, _ in self.module.named_modules()
              ]
              for policy in config.injection_policy_tuple:
                  if not any(
                      name.endswith(policy) for name in layer_names
                  ):
                      raise ValueError(
                          f"Injection policy layer'{policy}' not valid."
                      )
  
              self._apply_injection_policy(config, client_module)
      else:
          if config.replace_with_kernel_inject:
              # 2. DeepSpeed Kernel Injection
              self._apply_injection_policy(config)
          elif config.tensor_parallel.tp_size > 1:
              # 3. Automatic Tensor Parallelism
              parser_dict = AutoTP.tp_parser(model)
              print("AutoTP: ", parser_dict)
              for client_module, injection_policy in parser_dict:
                  if isinstance(injection_policy, str):
                      config.injection_policy_tuple = (
                          injection_policy,
                      )
                  else:
                      config.injection_policy_tuple = injection_policy
                  self._apply_injection_policy(config, client_module)
  
      device = get_accelerator().current_device_name()
      # NOTE: This check assumes a Hugging Face hierarchy for the device type i.e. module.device.type
      is_meta_device = (
          hasattr(self.module, "device")
          and self.module.device.type == "meta"
      )
      if is_meta_device:
          self.module.to_empty(device=device)
      else:
          self.module.to(device)
  
      if config.tensor_parallel.tp_size > 1:
          _rng_state = (
              get_accelerator()
              .get_rng_state()
              .to(get_accelerator().current_device_name())
          )
          dist.broadcast(_rng_state, 0)
          get_accelerator().set_rng_state(_rng_state.cpu())
  
      if config.tensor_parallel.tp_size > 1:
          assert (
              not config.enable_cuda_graph
          ), "Cuda graph is not supported for model parallelism"
  
      # Check if local CUDA graphs can be created in replacement modules
      self.local_cuda_graph = self._local_cuda_graph_used(self.module)
  
imports: import torch
  import time
  import os
  from deepspeed import comm
  from deepspeed.utils.logging import log_dist
  from torch.nn.modules import Module
  from packaging import version
  from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine
  from deepspeed.utils.timer import SynchronizedWallClockTimer
  from ..runtime.state_dict_factory import SDLoaderFactory
  from ..runtime.weight_quantizer import WeightQuantization
  from ..module_inject import replace_transformer_layer, generic_injection
  from ..comm.comm import init_distributed
  from ..pipe import PipelineModule
  from ..moe.utils import has_moe_layers
  from ..module_inject import LinearAllreduce, LinearLayer, Normalize, ReplaceWithTensorSlicing
  from deepspeed.accelerator import get_accelerator
  from ..module_inject.policy import TransformerPolicy
  from ..module_inject.auto_tp import AutoTP
  from ..module_inject.replace_policy import generic_policies
  from ..module_inject.auto_tp_model_utils import build_bloom_alibi_tensor, build_mpt_atten_bias_tensor, build_mpt_alibi_tensor, get_alibi_mask
  from ..ops.transformer.inference.ds_attention import DeepSpeedSelfAttention
  from ..model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference
  from torch import nn
label_desc: Use image node_id 1 to create a new InferenceEngine object from inherited base classes: Module with example: obj = InferenceEngine(model, config)