Updated re PR review

Author: TwoAbove

helpers/arguments.py

@@ -1428,11 +1428,11 @@ def parse_args(input_args=None):
         ),
     )
     parser.add_argument(
-        "--cuda_clear_cache_steps",
+        "--accelerator_cache_clear_interval",
         default=None,
         type=int,
         help=(
-            "Clear the CUDA cache every X steps. This can help prevent memory leaks, but may slow down training."
+            "Clear the cache from VRAM every X steps. This can help prevent memory leaks, but may slow down training."
         ),
     )
 

helpers/caching/memory.py

@@ -0,0 +1,12 @@
+def reclaim_memory():
+    import gc
+    import torch
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+
+    if torch.backends.mps.is_available():
+        torch.mps.empty_cache()
+
+    gc.collect()

train_sd21.py

@@ -21,6 +21,7 @@
 
 from pathlib import Path
 from helpers.arguments import parse_args
+from helpers.caching.memory import reclaim_memory
 from helpers.legacy.validation import prepare_validation_prompt_list
 from helpers.training.validation import Validation
 from helpers.training.state_tracker import StateTracker
@@ -109,15 +110,6 @@
 check_min_version("0.27.0.dev0")
 
 
-def garbage_collection():
-    import gc
-
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.ipc_collect()
-    gc.collect()
-
-
 SCHEDULER_NAME_MAP = {
     "euler": EulerDiscreteScheduler,
     "euler-a": EulerAncestralDiscreteScheduler,
@@ -885,7 +877,7 @@ def main():
         for _, backend in StateTracker.get_data_backends().items():
             if "vaecache" in backend:
                 backend["vaecache"].vae = None
-        garbage_collection()
+        reclaim_memory()
         memory_after_unload = torch.cuda.memory_allocated() / 1024**3
         memory_saved = memory_after_unload - memory_before_unload
         logger.info(
@@ -1570,7 +1562,7 @@ def main():
             )
 
             del text_encoder_lora_layers
-            garbage_collection()
+            reclaim_memory()
 
         if args.use_ema:
             ema_unet.copy_to(unet.parameters())

train_sdxl.py

@@ -22,6 +22,7 @@
 
 from pathlib import Path
 from helpers.arguments import parse_args
+from helpers.caching.memory import reclaim_memory
 from helpers.legacy.validation import prepare_validation_prompt_list
 from helpers.training.validation import Validation
 from helpers.training.state_tracker import StateTracker
@@ -219,15 +220,6 @@ def get_tokenizers(args):
     return tokenizer_1, tokenizer_2, tokenizer_3
 
 
-import gc
-
-
-def garbage_collection():
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.ipc_collect()
-    gc.collect()
-
 def main():
     StateTracker.set_model_type("sdxl")
     args = parse_args()
@@ -742,7 +734,7 @@ def main():
         text_encoder_2 = None
         text_encoder_3 = None
         text_encoders = []
-        garbage_collection()
+        reclaim_memory()
         memory_after_unload = torch.cuda.memory_allocated() / 1024**3
         memory_saved = memory_after_unload - memory_before_unload
         logger.info(
@@ -1182,7 +1174,7 @@ def main():
         for _, backend in StateTracker.get_data_backends().items():
             if "vaecache" in backend:
                 backend["vaecache"].vae = None
-        garbage_collection()
+        reclaim_memory()
         memory_after_unload = torch.cuda.memory_allocated() / 1024**3
         memory_saved = memory_after_unload - memory_before_unload
         logger.info(
@@ -1945,7 +1937,7 @@ def main():
                                 )
 
                 if global_step % args.cuda_clear_cache == 0:
-                    garbage_collection()
+                    reclaim_memory()
 
             logs = {
                 "step_loss": loss.detach().item(),
@@ -2042,7 +2034,7 @@ def main():
             del transformer
             del text_encoder_lora_layers
             del text_encoder_2_lora_layers
-            garbage_collection()
+            reclaim_memory()
         elif args.use_ema:
             if unet is not None:
                 ema_unet.copy_to(unet.parameters())