| 38 | |
| 39 | |
| 40 | class InferenceEngine(Module): |
| 41 | inference_mp_group = None |
| 42 | inference_ep_group = None |
| 43 | expert_mp_group = None |
| 44 | |
| 45 | def __init__(self, model, config): |
| 46 | """ |
| 47 | Args: |
| 48 | model: torch.nn.Module |
| 49 | config: DeepSpeedInferenceConfig |
| 50 | """ |
| 51 | global DS_INFERENCE_ENABLED |
| 52 | DS_INFERENCE_ENABLED = True |
| 53 | |
| 54 | super().__init__() |
| 55 | if DeepSpeedTransformerInference.workspace is not None: |
| 56 | self.destroy() |
| 57 | |
| 58 | self.module = model |
| 59 | self._config = config |
| 60 | |
| 61 | self._get_model_config_generate(config) # keep for weird backward compatibility |
| 62 | |
| 63 | # patch model generate with ours if model uses it |
| 64 | if hasattr(self.module, "generate"): |
| 65 | self.generate = self._generate |
| 66 | |
| 67 | if hasattr(self.module, "config"): |
| 68 | TransformerPolicy.hf_model_config = self.module.config |
| 69 | |
| 70 | if config.dtype not in get_accelerator().supported_dtypes(): |
| 71 | raise ValueError( |
| 72 | f"Data type {config.dtype} is not supported by {get_accelerator().device_name()} accelerator") |
| 73 | |
| 74 | # todo: keep this self.injection_dict because we don't use to change config.injection_policy API |
| 75 | # todo: this will get changed when Molly's PR on auto injection dict is merged |
| 76 | self.injection_dict = config.injection_policy |
| 77 | |
| 78 | # todo: refactor the mp_group and mp_size related in the next refactor |
| 79 | self.mp_group = config.tensor_parallel.tp_group |
| 80 | self.mpu = config.tensor_parallel.mpu |
| 81 | |
| 82 | self.quantize_merge_count = 1 |
| 83 | self.quantization_scales = None |
| 84 | |
| 85 | # these are not needed in the config as we are creating them ourselves in the inference engine |
| 86 | self.ep_group = None # config.moe.ep_group |
| 87 | self.expert_mp_group = None # config.moe.ep_mp_group |
| 88 | |
| 89 | self.cuda_graph_created = False |
| 90 | self.checkpoint_engine = TorchCheckpointEngine() |
| 91 | quantization_setting = None |
| 92 | self._init_quantization_setting( |
| 93 | quantization_setting) # todo: update with the new quant config for weight quant |
| 94 | self.model_profile_enabled = False |
| 95 | self._model_times = [] |
| 96 | |
| 97 | if not self.injection_dict and config.replace_with_kernel_inject: |