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

Rather than during the first generation.
2025-06-05 21:59:24 +02:00 · 2021-11-25 11:44:41 -05:00
parent 978dc486a5
commit f8bcc3411b
2 changed files with 36 additions and 32 deletions
--- a/aiserver.py
+++ b/aiserver.py
@ -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