Added experimental t5 offloading to CPU.

muse_maskgit_pytorch/trainers/maskgit_trainer.py

@@ -58,6 +58,7 @@ def __init__(
         clear_previous_experiments=False,
         validation_image_scale: float = 1.0,
         only_save_last_checkpoint=False,
+        t5_offloading=False,
         args=None,
     ):
         super().__init__(
@@ -88,6 +89,8 @@ def __init__(
         # we are going to use them later to save them to a config file.
         self.args = args
 
+        self.t5_offloading = t5_offloading
+
         # maskgit
         maskgit.vae.requires_grad_(False)
         maskgit.transformer.t5.requires_grad_(False)
@@ -161,7 +164,7 @@ def train(self):
                 if not text_embeds:
                     with torch.no_grad():
                         text_embeds = t5_encode_text_from_encoded(
-                            input_ids, attn_mask, self.model.transformer.t5, self.accelerator.device
+                            input_ids, attn_mask, self.model.transformer.t5, self.accelerator.device if not self.t5_offloading else 'cpu'
                         )
 
                 with self.accelerator.accumulate(self.model), self.accelerator.autocast():

train_muse_maskgit.py

@@ -453,6 +453,12 @@ def decompress_pickle(file):
     default="",
     help="The path to save or load embeds",
 )
+parser.add_argument(
+    "--t5_offloading",
+    action="store_true",
+    default=False,
+    help="Wheter to offload the t5 model to cpu instad of loading it on GPU. Should help with loading bigger models but will affect performance.",
+)
 
 
 @dataclass
@@ -528,6 +534,7 @@ class Arguments:
     attention_type: str = "flash"
     precompute: bool = False
     precompute_path: str = ""
+    t5_offloading = False
 
 
 def main():
@@ -984,6 +991,7 @@ def main():
         validation_image_scale=args.validation_image_scale,
         only_save_last_checkpoint=args.only_save_last_checkpoint,
         num_epochs=args.num_epochs,
+        t5_offloading=args.t5_offloading,
         args=args,
     )