feat: print memory usage during run

src/heretic/main.py

@@ -46,6 +46,7 @@ from .utils import (
     get_trial_parameters,
     load_prompts,
     print,
+    print_memory_usage,
     prompt_password,
     prompt_path,
     prompt_select,
@@ -310,6 +311,8 @@ def run():
             return
 
     model = Model(settings)
+    print()
+    print_memory_usage()
 
     print()
     print(f"Loading good prompts from [bold]{settings.good_prompts.dataset}[/]...")
@@ -543,6 +546,7 @@ def run():
             print(
                 f"[grey50]Estimated remaining time: [bold]{format_duration(remaining_time)}[/][/]"
             )
+        print_memory_usage()
 
         trial.set_user_attr("kl_divergence", kl_divergence)
         trial.set_user_attr("refusals", refusals)

src/heretic/utils.py

@@ -22,6 +22,7 @@ from datasets.config import DATASET_STATE_JSON_FILENAME
 from datasets.download.download_manager import DownloadMode
 from datasets.utils.info_utils import VerificationMode
 from optuna import Trial
+from psutil import Process
 from questionary import Choice, Style
 from rich.console import Console
 
@@ -30,6 +31,23 @@ from .config import DatasetSpecification, Settings
 print = Console(highlight=False).print
 
 
+def print_memory_usage():
+    def p(label: str, size_in_bytes: int):
+        print(f"[grey50]{label}: [bold]{size_in_bytes / (1024**3):.2f} GB[/][/]")
+
+    p("Resident system RAM", Process().memory_info().rss)
+
+    if torch.cuda.is_available():
+        p("Allocated GPU VRAM", torch.cuda.memory_allocated())
+        p("Reserved GPU VRAM", torch.cuda.memory_reserved())
+    elif is_xpu_available():
+        p("Allocated XPU memory", torch.xpu.memory_allocated())
+        p("Reserved XPU memory", torch.xpu.memory_reserved())
+    elif torch.backends.mps.is_available():
+        p("Allocated MPS memory", torch.mps.current_allocated_memory())
+        p("Driver (reserved) MPS memory", torch.mps.driver_allocated_memory())
+
+
 def is_notebook() -> bool:
     # Check for specific environment variables (Colab, Kaggle).
     # This is necessary because when running as a subprocess (e.g. !heretic),