File size: 5,352 Bytes
0692312 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | #!/usr/bin/env python
"""What each modality's rows actually pay under each candidate lambda, with the real quantizer.
`diag_lambda_crossmodal.py` answers this from the statistics with a bits-lost proxy. This answers
it with the same NVFP4 fake-quantizer the calibration scored on, on real sampled rows, and with the
rank-32 branch included -- so the number is comparable to the `err_smooth_lowrank` column in the
calibration artifacts rather than merely correlated with it.
Why it is needed at all: the calibration's own error column is measured on rows drawn uniformly
from the packed sequence, and video is 98.6% of those rows. That column therefore ranks
`lambda=video` best -- which it is, for video -- while being structurally unable to show what the
text and audio rows paid for it. Those rows are 1.4% of the sequence and every video row attends to
them.
Errors are relative L2 of the quantized layer output against bf16, per modality, using each
artifact's own lambda and its own low-rank factors.
"""
from __future__ import annotations
import argparse
import json
import sys
import time
from pathlib import Path
import torch
REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO / "scripts"))
sys.path.insert(0, str(REPO / "src"))
MODALITIES = ("video", "text", "audio")
# The whole point of this file is that its numbers sit next to `err_smooth_lowrank` in the
# calibration artifacts, so the quantizer has to be the identical object, not a second
# implementation of the same spec -- the per-tensor global scale under the FP8 group scale is easy
# to leave out of a reimplementation and shifts every number if you do.
from calib_smooth_lowrank import quantize_nvfp4 # noqa: E402
def err(x: torch.Tensor, w: torch.Tensor, lam: torch.Tensor,
l1: torch.Tensor | None, l2: torch.Tensor | None) -> float:
"""Relative L2 of the SVDQuant forward against bf16, on these rows."""
ref = x @ w.T
xs = x / lam
ws = w * lam
if l1 is not None and l1.abs().max() > 0:
ws = ws - l1 @ l2
low = (xs @ l2.T) @ l1.T
else:
low = 0.0
y = quantize_nvfp4(xs) @ quantize_nvfp4(ws).T + low
return float((y - ref).norm() / ref.norm().clamp_min(1e-12))
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--samples", required=True, help="directory of msample_*.pt")
ap.add_argument("--calibs", nargs="+", required=True, help="tag=path.pt")
ap.add_argument("--out", required=True)
ap.add_argument("--model-path", default=None)
ap.add_argument("--shard", default="0/1", help="i/n over LAYERS")
ap.add_argument("--max-rows", type=int, default=4096)
ap.add_argument("--device", default="cuda:0")
args = ap.parse_args()
import bench
from h3opt.svdquant_rules import target_linears
from diffusers import ModularPipeline
device = torch.device(args.device)
cal = {}
for spec in args.calibs:
tag, path = spec.split("=", 1)
cal[tag] = torch.load(path, map_location="cpu", weights_only=False)["layers"]
tags = list(cal)
i, n = (int(v) for v in args.shard.split("/"))
names = sorted(next(iter(cal.values())))[i::n]
pipe = ModularPipeline.from_pretrained(args.model_path or bench.DEFAULT_MODEL)
pipe.load_components(names=["transformer"], dtype=torch.bfloat16)
pipe.transformer.to("cpu") # one layer at a time on the GPU; 62 GB resident OOMs
weights = {k: m.weight.detach().to("cpu")
for k, m in target_linears(pipe.transformer).items() if k in names}
del pipe
torch.cuda.empty_cache()
rows: dict[str, dict[str, list]] = {k: {m: [] for m in MODALITIES} for k in names}
for p in sorted(Path(args.samples).glob("msample_*.pt")):
blob = torch.load(p, map_location="cpu", weights_only=False)
for k in names:
if k not in blob:
continue
for m in MODALITIES:
t = blob[k].get(m)
if t is not None and t.numel():
rows[k][m].append(t)
out: dict[str, dict] = {}
t0 = time.perf_counter()
for j, name in enumerate(names):
w = weights[name].to(device, torch.float32)
rec: dict[str, dict[str, float]] = {}
for m in MODALITIES:
if not rows[name][m]:
continue
x = torch.cat(rows[name][m])[: args.max_rows].to(device, torch.float32)
rec[m] = {"rows": int(x.shape[0])}
for t in tags:
r = cal[t][name]
rec[m][t] = err(x, w, r["lambda"].to(device, torch.float32),
r["l1"].to(device, torch.float32),
r["l2"].to(device, torch.float32))
del x
out[name] = rec
del w
torch.cuda.empty_cache()
if (j + 1) % 10 == 0 or j == len(names) - 1:
print(f"[{j+1}/{len(names)}] {name} {rec.get('video', {})}", flush=True)
p = Path(args.out)
p.parent.mkdir(parents=True, exist_ok=True)
p.write_text(json.dumps({"tags": tags, "shard": args.shard,
"wall_s": round(time.perf_counter() - t0, 1),
"layers": out}, indent=1))
print(f"wrote {p}", flush=True)
return 0
if __name__ == "__main__":
raise SystemExit(main())
|