From ce022f06e9a9ef1680888128ba8cd0134e956cf3 Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Thu, 13 Aug 2026 23:31:30 +0700 Subject: [PATCH] Add direct mmap and memory profiling controls --- src/h3_blackwell_runtime/checkpoint.py | 26 ++++++++++++++++++++++++ src/h3_blackwell_runtime/vae_decoder.py | 22 ++++++++++++++++---- tools/direct_t2v_preview.py | 27 +++++++++++++++++++++++++ 3 files changed, 71 insertions(+), 4 deletions(-) diff --git a/src/h3_blackwell_runtime/checkpoint.py b/src/h3_blackwell_runtime/checkpoint.py index b5aee5e..720cd6f 100644 --- a/src/h3_blackwell_runtime/checkpoint.py +++ b/src/h3_blackwell_runtime/checkpoint.py @@ -1,5 +1,6 @@ """Lazy loading for the current Comfy-format H3 safetensors checkpoint.""" +import os from pathlib import Path import torch @@ -13,8 +14,25 @@ class H3Checkpoint: def __init__(self, path: str | Path, device: str | torch.device = "cuda"): self.path = str(path) self.device = str(device) + self._no_mmap_tensors: dict[str, torch.Tensor] | None = None + + def _disable_mmap(self) -> bool: + return os.getenv("H3_DISABLE_MMAP", "").lower() in {"1", "true", "yes", "on"} + + def _all_tensors_no_mmap(self) -> dict[str, torch.Tensor]: + if self._no_mmap_tensors is None: + from safetensors.torch import load + + with open(self.path, "rb") as file: + tensors = load(file.read()) + self._no_mmap_tensors = {name: value.to(self.device) for name, value in tensors.items()} + return self._no_mmap_tensors def tensor(self, name: str, *, dtype: torch.dtype | None = None) -> torch.Tensor: + if self._disable_mmap(): + value = self._all_tensors_no_mmap()[name] + return value.to(dtype=dtype) if dtype is not None else value + from safetensors import safe_open with safe_open(self.path, framework="pt", device=self.device) as checkpoint: @@ -24,6 +42,14 @@ class H3Checkpoint: def nvfp4_linear(self, prefix: str, *, output_dtype=torch.bfloat16) -> Nvfp4Linear: names = ("comfy_quant", "weight", "weight_scale", "weight_scale_2", "bias", "pre_quant_scale") tensors = {} + if self._disable_mmap(): + available_tensors = self._all_tensors_no_mmap() + for suffix in names: + name = f"{prefix}.{suffix}" + if name in available_tensors: + tensors[name] = available_tensors[name] + return load_nvfp4_linear(tensors, prefix, output_dtype=output_dtype) + from safetensors import safe_open with safe_open(self.path, framework="pt", device=self.device) as checkpoint: diff --git a/src/h3_blackwell_runtime/vae_decoder.py b/src/h3_blackwell_runtime/vae_decoder.py index 781dc57..e3a5447 100644 --- a/src/h3_blackwell_runtime/vae_decoder.py +++ b/src/h3_blackwell_runtime/vae_decoder.py @@ -3,6 +3,7 @@ from __future__ import annotations import math +import os from pathlib import Path import torch @@ -164,14 +165,27 @@ class MiniMaxH3VideoVAE(nn.Module): def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", tiling: bool = True) -> "MiniMaxH3VideoVAE": model = cls(device="meta", tiling=tiling) expected = model.state_dict() - with safe_open(str(path), framework="pt", device=str(device)) as checkpoint: - available = set(checkpoint.keys()) + if os.getenv("H3_DISABLE_MMAP", "").lower() in {"1", "true", "yes", "on"}: + from safetensors.torch import load + + with open(path, "rb") as file: + available_weights = load(file.read()) + available = set(available_weights) missing = sorted(set(expected) - available) - shape_errors = [(name, tuple(expected[name].shape), tuple(checkpoint.get_slice(name).get_shape())) for name in expected if name in available and tuple(expected[name].shape) != tuple(checkpoint.get_slice(name).get_shape())] + shape_errors = [(name, tuple(expected[name].shape), tuple(available_weights[name].shape)) for name in expected if name in available and tuple(expected[name].shape) != tuple(available_weights[name].shape)] if missing or shape_errors: details = ([f"missing: {', '.join(missing)}"] if missing else []) + ([f"shape mismatch: {shape_errors}"] if shape_errors else []) raise ValueError("incompatible H3 VAE checkpoint; " + "; ".join(details)) - weights = {name: checkpoint.get_tensor(name) for name in expected} + weights = {name: available_weights[name].to(device) for name in expected} + else: + with safe_open(str(path), framework="pt", device=str(device)) as checkpoint: + available = set(checkpoint.keys()) + missing = sorted(set(expected) - available) + shape_errors = [(name, tuple(expected[name].shape), tuple(checkpoint.get_slice(name).get_shape())) for name in expected if name in available and tuple(expected[name].shape) != tuple(checkpoint.get_slice(name).get_shape())] + if missing or shape_errors: + details = ([f"missing: {', '.join(missing)}"] if missing else []) + ([f"shape mismatch: {shape_errors}"] if shape_errors else []) + raise ValueError("incompatible H3 VAE checkpoint; " + "; ".join(details)) + weights = {name: checkpoint.get_tensor(name) for name in expected} model.load_state_dict(weights, strict=True, assign=True) return model diff --git a/tools/direct_t2v_preview.py b/tools/direct_t2v_preview.py index 5a39c70..2d5b35a 100644 --- a/tools/direct_t2v_preview.py +++ b/tools/direct_t2v_preview.py @@ -1,6 +1,7 @@ """Generate a minimal direct, video-only H3 T2V preview without ComfyUI.""" import argparse +import os from pathlib import Path import subprocess import warnings @@ -29,16 +30,38 @@ parser.add_argument("--steps", type=int, default=12) parser.add_argument("--seed", type=int, default=440204) parser.add_argument("--attention", choices=("sage2", "sdpa", "sage3"), default="sage2") 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() + +def report_memory(stage: str) -> None: + if not args.profile_memory: + return + rss_kb = 0 + try: + with open("/proc/self/status", encoding="utf-8") as file: + for line in file: + if line.startswith("VmRSS:"): + rss_kb = int(line.split()[1]) + break + except OSError: + 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), "disable_mmap": os.getenv("H3_DISABLE_MMAP", "")}, flush=True) + +report_memory("start") checkpoint = H3Checkpoint("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors") conditioner = Qwen3VLPromptConditioner( "/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", "/opt/h3-blackwell-runtime/src/h3_blackwell_runtime/qwen25_tokenizer", ) +report_memory("qwen_loaded") video, audio, frames = random_av_latents(args.width, args.height, args.frames, args.seed) model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend=args.attention).eval() +report_memory("h3_loaded") text = H3TokenRefiner(checkpoint, attention_backend=args.attention)(conditioner(args.prompt)) +report_memory("text_conditioned") model_timesteps = None if args.model_timesteps_capture is not None: model_timesteps = [ @@ -46,10 +69,14 @@ if args.model_timesteps_capture is not None: for index in range(args.steps) ] latent = sample_video_res_multistep(model, H3PromptPacker(checkpoint), text, video, audio, steps=args.steps, model_timesteps=model_timesteps) +report_memory("sampled") vae = MiniMaxH3VideoVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda").eval() +report_memory("vae_loaded") with torch.inference_mode(): pixels = vae.decode(latent.to(next(vae.parameters()).dtype))[:, :, :frames] +report_memory("vae_decoded") pixels = ((pixels[0].permute(1, 2, 3, 0).clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu() +report_memory("pixels_cpu") args.output.parent.mkdir(parents=True, exist_ok=True) raw = args.output.with_suffix(".rgb") pixels.numpy().tofile(raw)