diff --git a/pyproject.toml b/pyproject.toml index 05d428e..f4a20ae 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/heretic/main.py b/src/heretic/main.py index 8ee49f4..26221cf 100644 --- a/src/heretic/main.py +++ b/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) diff --git a/src/heretic/utils.py b/src/heretic/utils.py index 036063a..7e20f89 100644 --- a/src/heretic/utils.py +++ b/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), diff --git a/uv.lock b/uv.lock index 60bd3cb..17a961f 100644 --- a/uv.lock +++ b/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" },