scripts/calib_stats_modal.py

#!/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())