rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-02-27 08:57:46 +01:00
Code Issues Projects Releases Wiki Activity

breakmodel fix for models without lm_head

Browse Source
This commit is contained in:
Gnome Ann 2021-08-20 19:32:18 -04:00
parent 8bfcf86a8b
commit e00d9c4362
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
aiserver.py
View File

@ -383,9 +383,10 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
breakmodel.total_blocks = n_layers breakmodel.total_blocks = n_layers
model.half().to('cpu') model.half().to('cpu')
gc.collect() gc.collect()
model.lm_head.to(breakmodel.gpu_device)
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)
if(hasattr(model), 'lm_head'):
model.lm_head.to(breakmodel.gpu_device)
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()
@ -433,9 +434,10 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
breakmodel.total_blocks = n_layers breakmodel.total_blocks = n_layers
model.half().to('cpu') model.half().to('cpu')
gc.collect() gc.collect()
model.lm_head.to(breakmodel.gpu_device)
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)
if(hasattr(model), 'lm_head'):
model.lm_head.to(breakmodel.gpu_device)
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()