feat: print memory usage during run

pyproject.toml

@@ -30,6 +30,7 @@ dependencies = [
     "kernels>=0.11.7",
     "optuna>=4.5.0",
     "peft>=0.14.0",
+    "psutil>=7.1.3",
     "pydantic-settings>=2.10.1",
     "questionary>=2.1.1",
     "rich>=14.1.0",

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),

uv.lock

@@ -868,6 +868,7 @@ dependencies = [
     { name = "kernels" },
     { name = "optuna" },
     { name = "peft" },
+    { name = "psutil" },
     { name = "pydantic-settings" },
     { name = "questionary" },
     { name = "rich" },
@@ -907,6 +908,7 @@ requires-dist = [
     { name = "optuna", specifier = ">=4.5.0" },
     { name = "pacmap", marker = "extra == 'research'", specifier = ">=0.8.0" },
     { name = "peft", specifier = ">=0.14.0" },
+    { name = "psutil", specifier = ">=7.1.3" },
     { name = "pydantic-settings", specifier = ">=2.10.1" },
     { name = "questionary", specifier = ">=2.1.1" },
     { name = "rich", specifier = ">=14.1.0" },