Only allow --breakmodel if it's supported

This commit is contained in:
Gnome Ann 2021-08-20 10:52:57 -04:00
parent eef0db8dee
commit 7717168676
1 changed files with 3 additions and 3 deletions

View File

@ -205,7 +205,7 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
if(args.cpu): if(args.cpu):
vars.usegpu = False vars.usegpu = False
vars.breakmodel = False vars.breakmodel = False
if(args.breakmodel): if(vars.bmsupported and args.breakmodel):
vars.usegpu = False vars.usegpu = False
vars.breakmodel = True vars.breakmodel = True
elif(vars.hascuda): elif(vars.hascuda):
@ -387,7 +387,7 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
model.transformer.wte.to(breakmodel.gpu_device) model.transformer.wte.to(breakmodel.gpu_device)
model.transformer.ln_f.to(breakmodel.gpu_device) model.transformer.ln_f.to(breakmodel.gpu_device)
gc.collect() gc.collect()
if(args.breakmodel): if(vars.bmsupported and args.breakmodel):
breakmodel.ram_blocks = max(0, min(n_layers, args.breakmodel)) breakmodel.ram_blocks = max(0, min(n_layers, args.breakmodel))
else: else:
print(colors.CYAN + "\nHow many layers would you like to put into system RAM?") print(colors.CYAN + "\nHow many layers would you like to put into system RAM?")
@ -435,7 +435,7 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
model.transformer.wte.to(breakmodel.gpu_device) model.transformer.wte.to(breakmodel.gpu_device)
model.transformer.ln_f.to(breakmodel.gpu_device) model.transformer.ln_f.to(breakmodel.gpu_device)
gc.collect() gc.collect()
if(args.breakmodel): if(vars.bmsupported and args.breakmodel):
breakmodel.ram_blocks = max(0, min(n_layers, args.breakmodel)) breakmodel.ram_blocks = max(0, min(n_layers, args.breakmodel))
else: else:
print(colors.CYAN + "\nHow many layers would you like to put into system RAM?") print(colors.CYAN + "\nHow many layers would you like to put into system RAM?")