From e9aee731e70718b007ec99856923a8eae1058abd Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Thu, 2 Jul 2026 11:58:46 +0300 Subject: [PATCH] nvfp4: static input-scale calibration (capture once, inject forever) Justified by live data: sampler phase timing shows sampling at 89% of the chain node, all 792 quantized layers on the fast path, and the probe puts dynamic input quantization at ~0.30 ms/layer-call vs ~0.11 ms with a provided scale (amax dominates) - a top-line per-call cost. New input_scales widget on the loader: - "auto-calibrate", first load without a sidecar: forward-pre-hooks record each quantized layer's running input amax (one generation is enough; sidecar under refinements/nvfp4_scales/.json, per-machine and rebuilt in one run whenever absent). - every later load: injects input_scale = amax * 1.25 / (448*6) into the state dict next to each quantized weight - core's loader registers it natively and skips the per-call amax from then on. Margin leaves clip headroom (NVFP4 clips above the scale). Tests: suffix-matched injection across checkpoint prefixes, value math, core Linear registering input_scale and forwarding, sidecar read. Co-Authored-By: Claude Fable 5 --- nvfp4_patch.py | 142 +++++++++++++++++++++++++++++++++++++- tests/test_nvfp4_patch.py | 52 ++++++++++++++ 2 files changed, 193 insertions(+), 1 deletion(-) diff --git a/nvfp4_patch.py b/nvfp4_patch.py index 8a62db7d..30bd349c 100644 --- a/nvfp4_patch.py +++ b/nvfp4_patch.py @@ -38,10 +38,123 @@ from __future__ import annotations import json import logging +import os import re import torch +# --------------------------------------------------------------------------- +# Static input-scale calibration +# --------------------------------------------------------------------------- +# Core quantizes each layer's INPUT dynamically (amax reduction + quantize) on every +# call unless the layer has a static `input_scale` parameter. Live probe on sm_120: +# the dynamic path costs ~0.30 ms/layer-call and a provided scale cuts it to ~0.11 ms +# (amax dominates) - across ~800 quantized layers that is a top-line cost of every +# model call. Calibration is capture-then-freeze: a "calibrate" run records each +# quantized layer's input amax via forward-pre-hooks into a per-checkpoint sidecar; +# later loads inject `input_scale = amax * MARGIN / (448 * 6)` into the state dict, +# where core's loader picks it up natively. MARGIN leaves headroom because NVFP4 +# clips activations above the scale; per the disposable-keys philosophy the sidecar +# is per-machine and rebuilt in one run whenever absent. +CALIBRATION_MARGIN = 1.25 +_F8F4_MAX = 448.0 * 6.0 # F8_E4M3_MAX * F4_E2M1_MAX +INPUT_SCALE_MODES = ["dynamic", "auto-calibrate"] + + +def _scales_path(checkpoint_path): + base = os.path.join(os.path.dirname(os.path.abspath(__file__)), "refinements", "nvfp4_scales") + return os.path.join(base, os.path.basename(checkpoint_path) + ".json") + + +def load_calibration(checkpoint_path): + """{module_path: input_amax} recorded by a previous calibrate run, or None.""" + try: + p = _scales_path(checkpoint_path) + if not os.path.exists(p): + return None + with open(p, "r", encoding="utf-8") as f: + data = json.load(f) + amax = data.get("amax") or {} + return {str(k): float(v) for k, v in amax.items()} or None + except Exception as e: # noqa: BLE001 + logging.warning(f"[FunPack NVFP4] could not read calibration sidecar: {e}") + return None + + +def inject_input_scales(sd, target_weight_keys, calibration, margin=CALIBRATION_MARGIN): + """Add `.input_scale` next to each quantized weight whose module path has a + calibrated amax. Keys are matched by suffix so any checkpoint prefix + ('model.diffusion_model.' or none) works. Returns the number injected.""" + injected = 0 + for k in target_weight_keys: + prefix = k[: -len("weight")] # '...transformer_blocks.3.attn1.to_q.' + module_path = prefix.rstrip(".") + for cal_key, amax in calibration.items(): + if module_path == cal_key or module_path.endswith("." + cal_key): + sd[prefix + "input_scale"] = torch.tensor( + float(amax) * margin / _F8F4_MAX, dtype=torch.float32) + injected += 1 + break + return injected + + +def install_calibration_capture(model, checkpoint_path): + """Forward-pre-hooks on every quantized Linear recording the running input amax. + The sidecar is written on the first full sweep and refreshed periodically; hooks + live for the model instance's lifetime (a calibration run is expected to be + slightly SLOWER - it pays the amax it is trying to eliminate, once).""" + store = {"amax": {}, "writes": 0, "calls": 0} + path = _scales_path(checkpoint_path) + + def _flush(): + try: + os.makedirs(os.path.dirname(path), exist_ok=True) + tmp = path + ".tmp" + with open(tmp, "w", encoding="utf-8") as f: + json.dump({"amax": store["amax"], "margin": CALIBRATION_MARGIN}, f) + os.replace(tmp, path) + store["writes"] += 1 + except Exception as e: # noqa: BLE001 + logging.warning(f"[FunPack NVFP4] calibration flush failed: {e}") + + hooked = 0 + try: + for name, m in model.model.diffusion_model.named_modules(): + if getattr(m, "quant_format", None) != "nvfp4": + continue + + def _hook(_module, args, _name=name): + try: + x = args[0] if args else None + if isinstance(x, torch.Tensor): + a = float(x.detach().abs().amax()) + if a > store["amax"].get(_name, 0.0): + store["amax"][_name] = a + store["calls"] += 1 + if store["calls"] % 4096 == 0: + _flush() + except Exception: + pass + + m.register_forward_pre_hook(_hook) + hooked += 1 + except Exception as e: # noqa: BLE001 + logging.warning(f"[FunPack NVFP4] calibration capture install failed: {e}") + return 0 + + if hooked: + # First flush after one full model call so a single generation is enough. + def _first_flush_hook(_module, _args): + if store["writes"] == 0 and len(store["amax"]) >= hooked: + _flush() + try: + last = [m for _, m in model.model.diffusion_model.named_modules() + if getattr(m, "quant_format", None) == "nvfp4"][-1] + last.register_forward_pre_hook(_first_flush_hook) + except Exception: + pass + return hooked + QUANTIZE_SCOPES = [ "video blocks", # video-branch Linears in middle blocks (default, safest) "video + cross-modal", # + audio_to_video_attn / video_to_audio_attn bridges @@ -330,10 +443,19 @@ class FunPackNVFP4ModelLoader: "tooltip": "First N transformer blocks stay full precision."}), "keep_last_blocks": ("INT", {"default": 2, "min": 0, "max": 16, "tooltip": "Last N transformer blocks stay full precision."}), + "input_scales": (INPUT_SCALE_MODES, { + "default": "dynamic", + "tooltip": "auto-calibrate: first run records each layer's input range " + "(slightly slower) into a per-checkpoint sidecar; every later " + "load injects static input scales, removing the per-call amax " + "cost (~0.2 ms x every quantized layer x every step). Reload " + "the model after the calibration run to activate.", + }), } } - def load(self, unet_name, quantize_scope, keep_first_blocks, keep_last_blocks): + def load(self, unet_name, quantize_scope, keep_first_blocks, keep_last_blocks, + input_scales="dynamic"): import comfy.model_management as mm import comfy.sd import comfy.utils @@ -377,9 +499,27 @@ class FunPackNVFP4ModelLoader: f"(scope='{quantize_scope}', full-precision blocks: {sorted(keep) or 'none'})." ) + calibration = load_calibration(path) if input_scales == "auto-calibrate" else None + if calibration: + injected = inject_input_scales(sd, targets, calibration) + logging.info( + f"[FunPack NVFP4] static input scales ACTIVE on {injected}/{len(targets)} layers " + f"(calibrated {len(calibration)}, margin {CALIBRATION_MARGIN}x) — per-call input " + "amax eliminated. Delete the sidecar under refinements/nvfp4_scales/ to recalibrate." + ) + model = comfy.sd.load_diffusion_model_state_dict(sd, metadata=metadata) if model is None: raise RuntimeError("ComfyUI could not detect the model type after NVFP4 patching.") + + if input_scales == "auto-calibrate" and not calibration: + hooked = install_calibration_capture(model, path) + logging.info( + f"[FunPack NVFP4] CALIBRATING input scales on {hooked} layers this session " + "(slightly slower). Run one full generation, then reload the model (change any " + "widget or restart) — static scales activate on the next load." + ) + try: log_nvfp4_diagnostics(model, device) except Exception as e: # noqa: BLE001 diff --git a/tests/test_nvfp4_patch.py b/tests/test_nvfp4_patch.py index 68fb5c0b..c0aa0f69 100644 --- a/tests/test_nvfp4_patch.py +++ b/tests/test_nvfp4_patch.py @@ -187,3 +187,55 @@ def test_diagnostics_reports_layer_census(capsys): assert "nvfp4 layers in loaded model: 3" in text assert "EMULATION" in text assert "GEMM probe skipped" in text or "GEMM probe" in text + + +def test_input_scale_injection_and_core_load(tmp_path, monkeypatch): + """Calibration sidecar -> inject_input_scales matches by suffix across checkpoint + prefixes, and core's MixedPrecisionOps Linear registers the injected input_scale + and still forwards correctly.""" + from nvfp4_patch import inject_input_scales, _scales_path + import nvfp4_patch as np_mod + + ckpt = str(tmp_path / "model.safetensors") + sidecar = _scales_path(ckpt) + assert sidecar.endswith("model.safetensors.json") + + cal = {"transformer_blocks.1.attn1.to_q": 12.5} + w = _w() + sd = {"model.diffusion_model.transformer_blocks.1.attn1.to_q.weight": w, + "model.diffusion_model.transformer_blocks.1.attn1.to_q.bias": torch.zeros(D, dtype=torch.bfloat16)} + out, targets, _ = quantize_state_dict_nvfp4(sd, "video blocks", 0, 0) + n = inject_input_scales(out, targets, cal) + assert n == 1 + key = "model.diffusion_model.transformer_blocks.1.attn1.to_q.input_scale" + assert key in out and out[key].dtype == torch.float32 + assert float(out[key]) == pytest.approx(12.5 * np_mod.CALIBRATION_MARGIN / (448.0 * 6.0)) + + # core load consumes it: input_scale becomes a registered parameter + import comfy.ops + ops = comfy.ops.mixed_precision_ops({}, torch.bfloat16, disabled=[]) + lin = ops.Linear(D, D, bias=True, device="cpu", dtype=torch.bfloat16) + layer_sd = {k.replace("model.diffusion_model.transformer_blocks.1.attn1.to_q.", ""): v + for k, v in out.items()} + missing, unexpected = lin.load_state_dict(layer_sd, strict=False) + assert not missing and not unexpected + assert isinstance(getattr(lin, "input_scale", None), torch.nn.Parameter) + x = torch.randn(2, 4, D, dtype=torch.bfloat16) * 0.5 + with torch.no_grad(): + y = lin.forward(x) + y_ref = torch.nn.functional.linear(x, w) + rel = (y.float() - y_ref.float()).abs().mean() / (y_ref.float().abs().mean() + 1e-8) + assert y.shape == y_ref.shape and rel < 0.3 + + # no calibration entry for a layer -> nothing injected for it + sd2 = {"transformer_blocks.2.attn1.to_q.weight": _w()} + out2, targets2, _ = quantize_state_dict_nvfp4(sd2, "video blocks", 0, 0) + assert inject_input_scales(out2, targets2, cal) == 0 + + +def test_load_calibration_reads_sidecar(tmp_path, monkeypatch): + import nvfp4_patch as np_mod + monkeypatch.setattr(np_mod, "_scales_path", lambda p: str(tmp_path / "s.json")) + assert np_mod.load_calibration("whatever") is None + (tmp_path / "s.json").write_text('{"amax": {"a.b": 3.0}, "margin": 1.25}') + assert np_mod.load_calibration("whatever") == {"a.b": 3.0}