rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-06-05 21:59:24 +02:00
Code Issues Projects Releases Wiki Activity

Basic breakmodel ui support

Browse Source 
Seems to work
This commit is contained in:
somebody
2023-06-21 13:57:32 -05:00
parent f326fc07e8
commit 0052ad401a
2 changed files with 83 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
modeling/inference_models/generic_hf_torch/class.py
View File

@@ -59,8 +59,6 @@ class model_backend(HFTorchInferenceModel):
# Also, lazy loader doesn't support GPT-2 models
self.lazy_load = False
# If we're using torch_lazy_loader, we need to get breakmodel config
# early so that it knows where to load the individual model tensors
logger.debug(
"lazy_load: {} hascuda: {} breakmodel: {} nobreakmode: {}".format(
self.lazy_load,
@@ -70,6 +68,16 @@ class model_backend(HFTorchInferenceModel):
)
)
# If we're using torch_lazy_loader, we need to get breakmodel config
# early so that it knows where to load the individual model tensors
if (
self.lazy_load
and utils.koboldai_vars.hascuda
and utils.koboldai_vars.breakmodel
and not utils.koboldai_vars.nobreakmodel
):
self.breakmodel_device_config(self.model_config)
if self.lazy_load:
# torch_lazy_loader.py and low_cpu_mem_usage can't be used at the same time
tf_kwargs.pop("low_cpu_mem_usage", None)