From c8759331620101b2a1f48752704d27acfb9ac8a6 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sun, 12 Jul 2026 17:48:16 +0200 Subject: [PATCH] Replace peak RAM/VRAM debug prints with a PeakMemoryRecorder context manager Reworks the always-on, import-time-started peak-memory instrumentation into a scoped context manager alongside TorchMemoryRecorder/TorchProfiler in profiling_util.py. Peak RAM is polled on a daemon thread only for the duration of the with-block (Linux via /proc RssAnon, Windows via psutil private bytes); peak VRAM uses CUDA's native counters, reset per block. Co-Authored-By: Claude Sonnet 5 --- modules/trainer/GenericTrainer.py | 7 ++-- modules/ui/TrainUIController.py | 5 +++ modules/util/profiling_util.py | 63 +++++++++++++++++++++++++++++++ 3 files changed, 72 insertions(+), 3 deletions(-) diff --git a/modules/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index dd17ad76e..d6337c744 100644 --- a/modules/trainer/GenericTrainer.py +++ b/modules/trainer/GenericTrainer.py @@ -30,7 +30,7 @@ from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.TimeUnit import TimeUnit from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.profiling_util import TorchMemoryRecorder, TorchProfiler +from modules.util.profiling_util import PeakMemoryRecorder, TorchMemoryRecorder, TorchProfiler from modules.util.time_util import get_string_timestamp from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -201,8 +201,9 @@ def __enqueue_sample_during_training(self, fun: Callable): self.sample_queue.append(fun) def __execute_sample_during_training(self): - for fun in self.sample_queue: - fun() + with PeakMemoryRecorder("sampling", enabled=False): + for fun in self.sample_queue: + fun() self.sample_queue = [] def __sample_loop( diff --git a/modules/ui/TrainUIController.py b/modules/ui/TrainUIController.py index 72623def4..8f4844055 100644 --- a/modules/ui/TrainUIController.py +++ b/modules/ui/TrainUIController.py @@ -19,6 +19,7 @@ from modules.util.callbacks.TrainCallbacks import TrainCallbacks from modules.util.commands.TrainCommands import TrainCommands from modules.util.config.TrainConfig import TrainConfig +from modules.util.profiling_util import PeakMemoryRecorder from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress from modules.util.ui.validation import flush_and_validate_all @@ -214,6 +215,10 @@ def generate_debug_package(self, zip_path: Path): self.view.on_update_status(f"Error generating debug package: {e}") def __training_thread_function(self): + with PeakMemoryRecorder("training run", enabled=False): + self.__training_thread_function_impl() + + def __training_thread_function_impl(self): error_caught = False self.training_callbacks = TrainCallbacks( diff --git a/modules/util/profiling_util.py b/modules/util/profiling_util.py index 01333b4e2..fd5c1eb9f 100644 --- a/modules/util/profiling_util.py +++ b/modules/util/profiling_util.py @@ -1,8 +1,71 @@ import platform +import threading import torch +def _private_ram_bytes() -> int: + # Private, non-file-backed resident memory. Excludes mmap'd file pages -- notably the safetensors weights + # that safe_open maps during load/materialize, which are reclaimable page cache and don't reflect real + # memory pressure. This matches what system monitors report as "used" memory, unlike ru_maxrss / VmHWM + # (Linux) or working-set (Windows), which count the mmap'd file cache too and so overstate usage on the + # disk-offload path. + if platform.system() == 'Linux': + # RssAnon from /proc; reported in KiB. + with open("/proc/self/status") as f: + for line in f: + if line.startswith("RssAnon:"): + return int(line.split()[1]) * 1024 + raise RuntimeError("RssAnon not found in /proc/self/status") + if platform.system() == 'Windows': + # psutil's pmem.private is process-private bytes, i.e. Windows' analogue of anon RSS. + import psutil + return psutil.Process().memory_info().private + raise RuntimeError(f"peak RAM tracking not supported on {platform.system()}") + + +class PeakMemoryRecorder: + # The OS only tracks a peak for total RSS / working set, not for private/anon RSS, so peak RAM is found by + # polling it on a daemon thread for the duration of the with-block. Peak VRAM needs no polling: the CUDA + # caching allocator already tracks it, reset_peak_memory_stats() scopes it to this block. + def __init__(self, label: str, enabled: bool = True, poll_interval: float = 0.1): + self.label = label + self.enabled = enabled + self.ram_enabled = enabled and platform.system() in ('Linux', 'Windows') + self.poll_interval = poll_interval + self._peak_ram = 0 + self._stop_event = threading.Event() + self._thread = None + + def __enter__(self): + if self.ram_enabled: + self._peak_ram = _private_ram_bytes() + self._stop_event.clear() + self._thread = threading.Thread(target=self._poll_ram, daemon=True) + self._thread.start() + if self.enabled and torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + return self + + def _poll_ram(self): + while not self._stop_event.is_set(): + self._peak_ram = max(self._peak_ram, _private_ram_bytes()) + self._stop_event.wait(self.poll_interval) + + def __exit__(self, exc_type, exc_val, exc_tb): + if self.ram_enabled: + self._stop_event.set() + self._thread.join() + # fold in one last sample in case the peak happened between the last poll and thread stop + self._peak_ram = max(self._peak_ram, _private_ram_bytes()) + print(f"[peak_ram] {self.label}: {self._peak_ram / 1024 ** 3:.2f} GiB") + + if self.enabled and torch.cuda.is_available(): + allocated = torch.cuda.max_memory_allocated() / 1024 ** 3 + reserved = torch.cuda.max_memory_reserved() / 1024 ** 3 + print(f"[peak_vram] {self.label}: {allocated:.2f} GiB (reserved {reserved:.2f} GiB)") + + class TorchMemoryRecorder: def __init__(self, filename: str = "memory.pickle", enabled: bool = True): self.filename = filename