T5 second GPU offload

2024-04-10 01:54:08 +02:00
parent 5e3886c094
commit 92323c62e7
4 changed files with 16 additions and 2 deletions
@@ -33,7 +33,9 @@ class T5v11Model(torch.nn.Module):
            else:
                if dtype: model_args["torch_dtype"] = dtype
                self.bnb = False
-                # TODO: custom device map?
+                # second GPU offload hack part 2
+                if device.startswith("cuda"):
+                    model_args["device_map"] = device
            print(f"Loading T5 from '{textmodel_path}'")
            self.transformer = T5EncoderModel.from_pretrained(textmodel_path, **model_args)
        else: