Timestamp direct memory profile output

This commit is contained in:
Daniel Maddern 2026-08-13 23:49:41 +07:00
parent bc0cff88ce
commit c3d72d9f1e
3 changed files with 10 additions and 3 deletions

View file

@ -32,7 +32,7 @@ class H3Checkpoint:
with fastsafe_open(filenames=[self.path], nogds=True, device=self._fast_device()) as checkpoint:
self._no_mmap_tensors = {
name: checkpoint.get_tensor(name).clone().detach()
for name in checkpoint.get_keys()
for name in checkpoint.keys()
}
return self._no_mmap_tensors

View file

@ -172,7 +172,7 @@ class MiniMaxH3VideoVAE(nn.Module):
with fastsafe_open(filenames=[str(path)], nogds=True, device=fast_device) as checkpoint:
available_weights = {
name: checkpoint.get_tensor(name).clone().detach()
for name in checkpoint.get_keys()
for name in checkpoint.keys()
}
available = set(available_weights)
missing = sorted(set(expected) - available)

View file

@ -4,6 +4,8 @@ import argparse
import os
from pathlib import Path
import subprocess
import time
from datetime import datetime, timezone
import warnings
warnings.filterwarnings("ignore", message="Found GPU0 NVIDIA GB10 which is of cuda capability 12.1.*", category=UserWarning)
@ -32,11 +34,15 @@ parser.add_argument("--attention", choices=("sage2", "sdpa", "sage3"), default="
parser.add_argument("--model-timesteps-capture", type=Path, help="Directory containing captured input_XX.pt H3 timesteps for strict parity checks.")
parser.add_argument("--profile-memory", action="store_true")
args = parser.parse_args()
started = time.perf_counter()
last_report = started
def report_memory(stage: str) -> None:
global last_report
if not args.profile_memory:
return
now = time.perf_counter()
rss_kb = 0
try:
with open("/proc/self/status", encoding="utf-8") as file:
@ -48,7 +54,8 @@ def report_memory(stage: str) -> None:
pass
cuda_alloc = torch.cuda.memory_allocated() / 1024**3 if torch.cuda.is_available() else 0.0
cuda_reserved = torch.cuda.memory_reserved() / 1024**3 if torch.cuda.is_available() else 0.0
print({"stage": stage, "rss_gb": round(rss_kb / 1024**2, 3), "cuda_alloc_gb": round(cuda_alloc, 3), "cuda_reserved_gb": round(cuda_reserved, 3), "fast_safetensors": os.getenv("H3_FAST_SAFETENSORS", ""), "disable_mmap": os.getenv("H3_DISABLE_MMAP", "")}, flush=True)
print({"stage": stage, "ts": datetime.now(timezone.utc).isoformat(), "epoch_s": round(time.time(), 3), "elapsed_s": round(now - started, 3), "delta_s": round(now - last_report, 3), "rss_gb": round(rss_kb / 1024**2, 3), "cuda_alloc_gb": round(cuda_alloc, 3), "cuda_reserved_gb": round(cuda_reserved, 3), "fast_safetensors": os.getenv("H3_FAST_SAFETENSORS", ""), "disable_mmap": os.getenv("H3_DISABLE_MMAP", "")}, flush=True)
last_report = now
report_memory("start")
checkpoint = H3Checkpoint("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")