Actually use args.breakmodel_layers
This commit is contained in:
parent
e00d9c4362
commit
a8bbfab87a
|
@ -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,")
|
||||||
|
|
Loading…
Reference in New Issue