Fix memory leak from sigint signal and deepspeed

src/axolotl/train.py

@@ -36,7 +36,6 @@
 configure_logging()
 LOG = get_logger("axolotl.train")
 
-
 @dataclass
 class TrainDatasetMeta:
     """
@@ -126,15 +125,18 @@ def train(
 
     # In case we want to stop early with ctrl+c, this is a nice to have to save the pretrained model
     if cfg.local_rank == 0:
-
-        def terminate_handler(_, __, model):
-            if cfg.flash_optimum and BetterTransformer:
-                model = BetterTransformer.reverse(model)
-            model.save_pretrained(cfg.output_dir, safe_serialization=safe_serialization)
+        import weakref
+        def terminate_handler(_, __, model_weakref):
+            if model_weakref():
+                _model = model_weakref()
+                if cfg.flash_optimum and BetterTransformer:
+                    _model = BetterTransformer.reverse(_model)
+                _model.save_pretrained(cfg.output_dir, safe_serialization=safe_serialization)
             sys.exit(0)
 
+        _model_weakref = weakref.ref(model)
         signal.signal(
-            signal.SIGINT, lambda signum, frame: terminate_handler(signum, frame, model)
+            signal.SIGINT, lambda signum, frame: terminate_handler(signum, frame, _model_weakref)
         )
 
     badge_markdown = """[<img src="https://raw.githubusercontent.com/OpenAccess-AI-Collective/axolotl/main/image/axolotl-badge-web.png" alt="Built with Axolotl" width="200" height="32"/>](https://github.com/OpenAccess-AI-Collective/axolotl)"""
@@ -216,6 +218,11 @@ def terminate_handler(_, __, model):
         # defensively push to the hub to ensure the model card is updated
         trainer.push_to_hub()
 
+    if cfg.deepspeed:
+        trainer.deepspeed.destroy()
+    trainer.accelerator.free_memory()
+    trainer.model, trainer.optimizer = None, None
+
     return model, tokenizer