In breakmodel mode, move layers to GPU as soon as model loads

Rather than during the first generation.
This commit is contained in:
Gnome Ann
2021-11-25 11:44:41 -05:00
parent 978dc486a5
commit f8bcc3411b
2 changed files with 36 additions and 32 deletions

View File

@ -303,6 +303,7 @@ def device_config(model):
gc.collect()
GPTNeoModel.forward = breakmodel.new_forward
generator = model.generate
breakmodel.move_hidden_layers(model.transformer)
#==================================================================#
# Startup