Actually use args.breakmodel_layers

This commit is contained in:
Gnome Ann 2021-08-20 20:50:03 -04:00
parent e00d9c4362
commit a8bbfab87a
1 changed files with 4 additions and 4 deletions

View File

@ -390,8 +390,8 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
if(not hasattr(model.config, 'rotary') or not model.config.rotary): if(not hasattr(model.config, 'rotary') or not model.config.rotary):
model.transformer.wpe.to(breakmodel.gpu_device) model.transformer.wpe.to(breakmodel.gpu_device)
gc.collect() gc.collect()
if(vars.bmsupported and args.breakmodel): if(args.breakmodel_layers is not None):
breakmodel.ram_blocks = max(0, min(n_layers, args.breakmodel)) breakmodel.ram_blocks = max(0, min(n_layers, args.breakmodel_layers))
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?")
print("The more of them you put into system RAM, the slower it will run,") print("The more of them you put into system RAM, the slower it will run,")
@ -441,8 +441,8 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
if(not hasattr(model.config, 'rotary') or not model.config.rotary): if(not hasattr(model.config, 'rotary') or not model.config.rotary):
model.transformer.wpe.to(breakmodel.gpu_device) model.transformer.wpe.to(breakmodel.gpu_device)
gc.collect() gc.collect()
if(vars.bmsupported and args.breakmodel): if(args.breakmodel_layers is not None):
breakmodel.ram_blocks = max(0, min(n_layers, args.breakmodel)) breakmodel.ram_blocks = max(0, min(n_layers, args.breakmodel_layers))
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?")
print("The more of them you put into system RAM, the slower it will run,") print("The more of them you put into system RAM, the slower it will run,")