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

Lazy loader should convert model tensors to float16 before moving them

This commit is contained in:
Gnome Ann 2022-03-05 11:31:34 -05:00
parent 5674516f0c
commit 86ac562b0c
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
aiserver.py
View File

@ -1028,7 +1028,7 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly", "TPUMeshTransforme
# Lazy loader
import torch_lazy_loader
def get_lazy_load_callback(n_layers):
def get_lazy_load_callback(n_layers, convert_to_float16=True):
if not vars.lazy_load:
return
@ -1072,7 +1072,10 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly", "TPUMeshTransforme
f.seek(model_dict[key].seek_offset - current_offset, 1)
device = device_map[key]
#print(f"Transferring <{key}> to {'(CPU)' if device == 'cpu' else '[device ' + str(device) + ']'} ... ", end="", flush=True)
model_dict[key] = model_dict[key].materialize(f, map_location=torch.device(device))
model_dict[key] = model_dict[key].materialize(f, map_location="cpu")
if convert_to_float16 and model_dict[key].dtype is torch.float32:
model_dict[key] = model_dict[key].to(torch.float16)
model_dict[key] = model_dict[key].to(device)
#print("OK", flush=True)
finally:
if isinstance(f, zipfile.ZipExtFile):