File size: 4,696 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 | #!/usr/bin/env python
"""What each modality pays when lambda was calibrated on a different one.
The end-to-end runs put a number on the *outcome* -- lambda=video lands in a different sample than
BF16 does while lambda=all tracks it -- but not on the mechanism. This does, from the per-modality
statistics alone, with no GPU.
NVFP4 groups 16 consecutive input channels under one FP8 scale set by that group's absmax. A
channel whose own absmax is far below its group's absmax spends its mantissa on range it never
uses, and the loss is `log2(group_absmax / channel_absmax)` bits. That quantity is what smoothing
exists to reduce: `X/lambda` is the activation the kernel actually sees, so lambda reshapes exactly
this profile.
The point is that `W * lambda` is one weight and every modality's rows pass through it. Pick lambda
from video's absmax and video's profile flattens -- but text and audio are divided by a vector that
has nothing to do with their own profile, and an uncorrelated divisor makes a profile *sharper*,
not flatter. This prints the bits lost per modality under each lambda, so "video's gain" and
"text's loss" are on the same axis.
The `lambda = 1` column is the no-smoothing floor, i.e. what every modality pays if nobody is
favoured.
"""
from __future__ import annotations
import argparse
from pathlib import Path
import torch
MODALITIES = ("video", "text", "audio")
BLOCK = 16
def bits_lost(absmax: torch.Tensor, lam: torch.Tensor) -> float:
"""Mean bits a channel loses to its 16-wide group's scale, after dividing by lambda."""
a = (absmax.double() / lam.double()).clamp_min(1e-12)
n = (a.numel() // BLOCK) * BLOCK
g = a[:n].reshape(-1, BLOCK)
return float(torch.log2(g.amax(dim=1, keepdim=True) / g).mean())
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--stats", required=True, help="stats_modal_768p.pt")
ap.add_argument("--calibs", nargs="+", required=True, help="tag=path.pt, e.g. all=...pt")
ap.add_argument("--top", type=int, default=8)
args = ap.parse_args()
st = torch.load(args.stats, map_location="cpu", weights_only=False)["stats"]
cal = {}
for spec in args.calibs:
tag, path = spec.split("=", 1)
cal[tag] = torch.load(path, map_location="cpu", weights_only=False)["layers"]
tags = ["1"] + list(cal)
rows = {m: {t: [] for t in tags} for m in MODALITIES}
per_layer = []
for name, s in st.items():
if name not in next(iter(cal.values())):
continue
lams = {"1": torch.ones_like(s["video"]["absmax"])}
lams.update({t: cal[t][name]["lambda"].float() for t in cal})
rec = {"layer": name}
for m in MODALITIES:
am = s[m]["absmax"].float()
if float(am.max()) == 0.0: # audio is absent from these clips for some layers
continue
for t in tags:
v = bits_lost(am, lams[t])
rows[m][t].append(v)
rec[f"{m}/{t}"] = v
per_layer.append(rec)
print(f"bits lost to the block-16 scale, mean over channels, median over layers "
f"({len(per_layer)} layers)\n")
head = "modality".ljust(10) + "".join(f"lam={t}".rjust(12) for t in tags)
print(head)
print("-" * len(head))
for m in MODALITIES:
if not rows[m]["1"]:
print(f"{m:<10}" + " (no rows in calibration data)")
continue
cells = []
for t in tags:
v = sorted(rows[m][t])
cells.append(f"{v[len(v)//2]:.3f}".rjust(12))
print(f"{m:<10}" + "".join(cells))
print("\nchange vs no smoothing (negative = smoothing helps that modality)")
print(head)
print("-" * len(head))
for m in MODALITIES:
if not rows[m]["1"]:
continue
base = sorted(rows[m]["1"]); base = base[len(base) // 2]
cells = []
for t in tags:
v = sorted(rows[m][t]); v = v[len(v) // 2]
cells.append(("--" if t == "1" else f"{v - base:+.3f}").rjust(12))
print(f"{m:<10}" + "".join(cells))
if "video" in cal:
worst = sorted(per_layer, key=lambda r: -(r.get("text/video", 0) - r.get("text/all", 0)))
print(f"\nlayers where lambda=video costs text the most, vs lambda=all:")
for r in worst[: args.top]:
print(f" {r['layer']:<48} text {r.get('text/all', float('nan')):.3f} -> "
f"{r.get('text/video', float('nan')):.3f} "
f"video {r.get('video/all', float('nan')):.3f} -> "
f"{r.get('video/video', float('nan')):.3f}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|