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__()
  
      self.module = model
      self._config = config
  
      self._get_model_config_generate(
          config
      )  # keep for weird backward compatibility
  
      if hasattr(self.module, "config"):
          DSPolicy.hf_model_config = self.module.config
  
      # 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 = []
  
      # This is a hack to remove the prepare_mask function on HF side for BLOOM architecture
      self.remove_mask_prepare_for_bloom()
  
      if 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"
  
      if config.checkpoint and not config.replace_with_kernel_inject:
          self._load_checkpoint(config.checkpoint)
  
      # 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)
  
      # retain this from the old conditional argument being passed to apply_injection_policy()
      if not config.replace_with_kernel_inject:
          config.checkpoint = None
  
      if self.injection_dict:
          for (
              client_module,
              injection_policy,
          ) in self.injection_dict.items():
              # 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
              self._apply_injection_policy(config, client_module)
      elif config.replace_method == "auto":
          self._apply_injection_policy(config)
  
      device = torch.cuda.current_device()
      self.module.to(device)
  
      if config.tensor_parallel.tp_size > 1:
          _rng_state = torch.cuda.get_rng_state().to(
              torch.cuda.current_device()
          )
          dist.broadcast(_rng_state, 0)
          torch.cuda.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"
  
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.replace_module import replace_transformer_layer, generic_injection
  from ..comm.comm import init_distributed
  from ..pipe import PipelineModule
  from ..moe.utils import has_moe_layers
  from ..runtime.zero import GatheredParameters
  from ..module_inject import LinearAllreduce, LinearLayer, Normalize, ReplaceWithTensorSlicing
  from ..module_inject.replace_policy import DSPolicy
  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)