rixty
/
KoboldAI-Client
mirror of https://github.com/KoboldAI/KoboldAI-Client.git
1
1
Fork 0
Code Issues Projects Releases Wiki Activity

Handle -1's in prompt_tuner.py breakmodel_gpulayers

This commit is contained in:
vfbd 2022-08-22 21:25:07 -04:00
parent 09750acfa0
commit b60d14e3bf
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
prompt_tuner.py
View File

@ -692,7 +692,7 @@ class TrainerBase(abc.ABC):
if breakmodel_gpulayers is None: if breakmodel_gpulayers is None:
breakmodel_gpulayers = [] breakmodel_gpulayers = []
if breakmodel_primary_device is None: if breakmodel_primary_device is None:
breakmodel_primary_device = 0 if breakmodel_gpulayers else "cpu" breakmodel_primary_device = 0 if sum(x if x >= 0 else 0 for x in breakmodel_gpulayers) else "cpu"
if self.data.params is not None and "max_batch_size" not in self.data.params: if self.data.params is not None and "max_batch_size" not in self.data.params:
self.data.params["max_batch_size"] = 2048 self.data.params["max_batch_size"] = 2048
@ -730,6 +730,8 @@ class TrainerBase(abc.ABC):
model_config = self._get_model_config() model_config = self._get_model_config()
n_layers = utils.num_layers(model_config) n_layers = utils.num_layers(model_config)
breakmodel_gpulayers = [x if x >= 0 else n_layers for x in breakmodel_gpulayers]
convert_to_float16 = True convert_to_float16 = True
hascuda = torch.cuda.is_available() hascuda = torch.cuda.is_available()
usegpu = hascuda and not breakmodel_disklayers and len(breakmodel_gpulayers) == 1 and breakmodel_gpulayers[0] == n_layers usegpu = hascuda and not breakmodel_disklayers and len(breakmodel_gpulayers) == 1 and breakmodel_gpulayers[0] == n_layers