From 5f224e1366f94fd4c431e9698483b1036a34b2f4 Mon Sep 17 00:00:00 2001 From: somebody Date: Wed, 21 Jun 2023 14:13:14 -0500 Subject: [PATCH] Restore choice of lazyload or not --- modeling/inference_models/hf_torch.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/modeling/inference_models/hf_torch.py b/modeling/inference_models/hf_torch.py index c1bcdf0b..10b0fa3c 100644 --- a/modeling/inference_models/hf_torch.py +++ b/modeling/inference_models/hf_torch.py @@ -101,8 +101,9 @@ class HFTorchInferenceModel(HFInferenceModel): ret = super().set_input_parameters(parameters) # Hook onto input param setting for setting breakmodel stuff - self.breakmodel_config.gpu_blocks = self.layers - self.breakmodel_config.disk_blocks = self.disk_layers + if self.breakmodel: + self.breakmodel_config.gpu_blocks = self.layers + self.breakmodel_config.disk_blocks = self.disk_layers return ret @@ -303,20 +304,19 @@ class HFTorchInferenceModel(HFInferenceModel): # Try to determine model type from either AutoModel or falling back to legacy try: - with lazy_loader.use_lazy_load(dematerialized_modules=True): - metamodel = AutoModelForCausalLM.from_config(self.model_config) - device_map = self.breakmodel_config.get_device_map(metamodel) + if self.lazy_load: + with lazy_loader.use_lazy_load(dematerialized_modules=True): + metamodel = AutoModelForCausalLM.from_config(self.model_config) + tf_kwargs["device_map"] = self.breakmodel_config.get_device_map(metamodel) + print("Rodger rodger", tf_kwargs) with lazy_loader.use_lazy_load( - enable=True, + enable=self.lazy_load, # DO NOT DEMATERIALIZE MODULES / INIT WEIGHTS EMPTY!!! IT WILL EXPLODE!!!!!!! dematerialized_modules=False, ): - print(device_map) model = AutoModelForCausalLM.from_pretrained( location, - # device_map="auto", - device_map=device_map, offload_folder="accelerate-disk-cache", torch_dtype=torch.float16, **tf_kwargs,