Get LoRA script to work for single gpus

core/__init__.py

@@ -112,7 +112,7 @@ def setup_extras_post(self, extras: Extras, models: Models, optimizers: Optimize
 
     # perform the training here
     @abstractmethod
-    def train(self, data: Data, extras: Extras, models: Models, optimizers: Optimizers, schedulers: Schedulers):
+    def train(self, data: Data, extras: Extras, models: Models, optimizers: Optimizers, schedulers: Schedulers, single_gpu: bool=False):
         raise NotImplementedError("This method needs to be overriden")
     # ------------
 
@@ -357,7 +357,7 @@ def __call__(self, single_gpu=False):
         # TRAIN
         if self.is_main_node:
             print("**TRAINING STARTING...**")
-        self.train(data, extras, models, optimizers, schedulers)
+        self.train(data, extras, models, optimizers, schedulers, single_gpu)
 
         if single_gpu is False:
             barrier()

train/base.py

@@ -239,7 +239,7 @@ def decode_latents(self, latents: torch.Tensor, batch: dict, models: Models, ext
         raise NotImplementedError("This method needs to be overriden")
 
     def train(self, data: WarpCore.Data, extras: WarpCore.Extras, models: Models, optimizers: Optimizers,
-              schedulers: WarpCore.Schedulers):
+            schedulers: WarpCore.Schedulers, single_gpu: bool=False):
         start_iter = self.info.iter + 1
         max_iters = self.config.updates * self.config.grad_accum_steps
         if self.is_main_node:
@@ -304,13 +304,14 @@ def train(self, data: WarpCore.Data, extras: WarpCore.Extras, models: Models, op
                             'bucket_ranges': extras.gdf.loss_weight.bucket_ranges.tolist(),
                             'bucket_losses': extras.gdf.loss_weight.bucket_losses.tolist(),
                         }
-                    self.save_checkpoints(models, optimizers)
+                    self.save_checkpoints(models, optimizers, single_gpu=single_gpu)
                     if self.is_main_node:
                         create_folder_if_necessary(f'{self.config.output_path}/{self.config.experiment_id}/')
                     self.sample(models, data, extras)
 
-    def save_checkpoints(self, models: Models, optimizers: Optimizers, suffix=None):
-        barrier()
+    def save_checkpoints(self, models: Models, optimizers: Optimizers, suffix=None, single_gpu=False):
+        if single_gpu:
+            barrier()
         suffix = '' if suffix is None else suffix
         self.save_info(self.info, suffix=suffix)
         models_dict = models.to_dict()

train/train_c_lora.py

@@ -320,11 +320,11 @@ def decode_latents(self, latents: torch.Tensor, batch: dict, models: Models, ext
 
 if __name__ == '__main__':
     print("Launching Script")
+    single_gpu = bool(sys.argv[2])
     warpcore = WurstCore(
         config_file_path=sys.argv[1] if len(sys.argv) > 1 else None,
-        device=torch.device(int(os.environ.get("SLURM_LOCALID")))
+        device=torch.device(int(os.environ.get("SLURM_LOCALID")) if not single_gpu else 0)
     )
     warpcore.fsdp_defaults['sharding_strategy'] = ShardingStrategy.NO_SHARD
-
     # RUN TRAINING
-    warpcore()
+    warpcore(single_gpu=single_gpu)