File size: 7,160 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 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | #!/usr/bin/env python
"""Per-input-channel statistics, split by modality of the packed sequence.
`calib_stats.py` accumulates over every row of `[text | video | audio]` without a mask. For a
mean that is harmless -- non-video is 1.36% of rows -- but `absmax` is a maximum, and one extreme
row can set a channel's span alone. Measured on `blocks.0.attn.to_q`, **79.8% of channels have
their absmax set by the 102 text rows**, which are 0.27% of the sequence, and the resulting span
is 1.88x the video-only span at the median. The span becomes lambda, and `W * lambda` is
quantized once for every token, so a lambda pulled by the text rows is paid for by the video rows.
This keeps three independent accumulators per layer so the choice of which rows define lambda is
made later, from data, instead of being fixed by an unmasked `reshape(-1, C)`.
Row modality comes from the transformer's own `video_indices` / `audio_indices` / `text_indices`
kwargs rather than from any assumed tag encoding. The token refiner is a special case: it runs on
the text rows only, so its layers see a row count that matches neither the packed sequence nor a
slice of it, and those are recorded as text.
"""
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 = ("text", "video", "audio")
class ModalStats:
"""absmax / sums per input channel, kept separately for each modality."""
__slots__ = ("absmax", "sq_sum", "abs_sum", "rows")
def __init__(self, in_features: int, device: torch.device) -> None:
z = lambda: torch.zeros(in_features, dtype=torch.float64, device=device)
self.absmax = {m: z() for m in MODALITIES}
self.sq_sum = {m: z() for m in MODALITIES}
self.abs_sum = {m: z() for m in MODALITIES}
self.rows = {m: 0 for m in MODALITIES}
@torch.no_grad()
def update(self, x: torch.Tensor, masks: dict[str, torch.Tensor] | None) -> None:
flat = x.reshape(-1, x.shape[-1])
parts = ({m: flat[masks[m]] for m in MODALITIES} if masks is not None
else {"text": flat, "video": flat[:0], "audio": flat[:0]})
for m, part in parts.items():
if part.shape[0] == 0:
continue
p = part.to(torch.float64)
self.absmax[m] = torch.maximum(self.absmax[m], p.abs().amax(dim=0))
self.sq_sum[m] += p.square().sum(dim=0)
self.abs_sum[m] += p.abs().sum(dim=0)
self.rows[m] += p.shape[0]
def result(self) -> dict:
out = {}
for m in MODALITIES:
n = max(self.rows[m], 1)
rms = (self.sq_sum[m] / n).sqrt()
out[m] = {"absmax": self.absmax[m].float().cpu(), "rms": rms.float().cpu(),
"mean_abs": (self.abs_sum[m] / n).float().cpu(), "rows": self.rows[m]}
return out
def main() -> int:
p = argparse.ArgumentParser(description="Per-modality activation statistics for MiniMax-H3")
p.add_argument("--caches", required=True)
p.add_argument("--out", required=True)
p.add_argument("--model-path", default=None)
p.add_argument("--device", default="cuda:0")
p.add_argument("--shard", default=None, help="i/n; streaming stats, merges exactly")
p.add_argument("--attention-backend", default="_flash_3_hub")
args = p.parse_args()
import bench
from h3opt.svdquant_rules import inventory, target_linears
from diffusers import ModularPipeline
device = torch.device(args.device)
t0 = time.perf_counter()
pipe = ModularPipeline.from_pretrained(args.model_path or bench.DEFAULT_MODEL)
pipe.load_components(names=["transformer"], dtype=torch.bfloat16)
transformer = pipe.transformer.to(device).eval()
if args.attention_backend:
transformer.set_attention_backend(args.attention_backend)
print(f"denoiser loaded in {time.perf_counter() - t0:.1f}s", flush=True)
inv = inventory(transformer)
targets = target_linears(transformer)
stats = {n: ModalStats(m.in_features, device) for n, m in targets.items()}
state: dict[str, dict[str, torch.Tensor] | None] = {"masks": None, "seq": 0}
def set_layout(_module, _args, kwargs):
pos = kwargs.get("position_ids")
if pos is None:
return None
seq = int(pos.shape[0])
masks = {}
for m, key in (("text", "text_indices"), ("video", "video_indices"), ("audio", "audio_indices")):
idx = kwargs.get(key)
v = torch.zeros(seq, dtype=torch.bool, device=device)
if idx is not None:
v[idx.to(device)] = True
masks[m] = v
state["masks"], state["seq"] = masks, seq
return None
transformer.register_forward_pre_hook(set_layout, with_kwargs=True)
for name, module in targets.items():
def hook(_m, inputs, _n=name):
x = inputs[0]
rows = x.reshape(-1, x.shape[-1]).shape[0]
# The token refiner runs on the text rows alone, so its row count does not match the
# packed sequence; those rows are text by construction.
stats[_n].update(x, state["masks"] if rows == state["seq"] else None)
module.register_forward_pre_hook(hook)
cache_dir = Path(args.caches)
files = sorted(q for q in cache_dir.glob("*.pt") if ".cond" not in q.name)
if args.shard:
i, n = (int(v) for v in args.shard.split("/"))
files = files[i::n]
print(f"{len(files)} cached steps", flush=True)
cond: dict[str, dict] = {}
t1 = time.perf_counter()
with torch.no_grad():
for k, path in enumerate(files):
rec = torch.load(path, map_location="cpu", weights_only=False)
clip = rec["clip"]
if clip not in cond:
cond.clear()
cond[clip] = torch.load(cache_dir / f"{clip}.cond.pt", map_location="cpu",
weights_only=False)
kw = {kk: v for kk, v in rec.items()
if kk not in ("outputs", "clip", "step") and torch.is_tensor(v)}
kw.update(cond[clip])
kw = {kk: (v.to(device, torch.bfloat16) if v.is_floating_point() else v.to(device))
for kk, v in kw.items()}
transformer(**kw, return_dict=False)
if (k + 1) % 25 == 0:
r = (k + 1) / (time.perf_counter() - t1)
print(f" {k+1}/{len(files)} {r:.2f} steps/s eta {(len(files)-k-1)/r/60:.1f} min",
flush=True)
out = Path(args.out)
out.parent.mkdir(parents=True, exist_ok=True)
torch.save({"inventory": inv, "num_cache_files": len(files), "shard": args.shard,
"caches": str(cache_dir), "wall_s": round(time.perf_counter() - t1, 1),
"stats": {n: s.result() for n, s in stats.items()}}, out)
print(f"wrote {out} ({out.stat().st_size/1e6:.1f} MB)", flush=True)
return 0
if __name__ == "__main__":
raise SystemExit(main())
|