scripts/ltx_fake_quantize.py

"""Apply a precision plan to the encoder by quantize-dequantize, keeping BF16 on disk.

The point is to measure the *damage* a rung does before paying to pack it. Each
weight is quantized with the real codec's arithmetic and immediately expanded
back to BF16, so the output loads in stock `transformers` with no custom kernel,
no ComfyUI, and no new format - while carrying exactly the error the packed
artifact would carry.

What this does not measure: the size. The output is BF16 and therefore 26 GB no
matter which rung it implements. Sizes come from the plan's own byte accounting,
which counts the packed width. Keeping those two apart is deliberate - this
project has already published one figure that mixed an information-theoretic
bit-width with an on-disk size, and the fix was to report both separately.

Codecs, matched to what ComfyUI's `MixedPrecisionOps` dispatches on:

* `nvfp4` - FP4 E2M1 with an FP8 E4M3 scale per 16 elements and one FP32 scale
  per tensor, which is what keeps the group scales inside E4M3's range
* `int8_tensorwise` - symmetric int8 with an FP32 scale per output channel. The
  vendor's own encoder ships exactly this: an I8 `weight` beside an F32
  `weight_scale` of shape [out, 1].

The vendor also applies `convrot`, an orthogonal rotation baked into the stored
weights, which reduces quantization error by spreading outliers. This script
does not implement it - the inverse lives in an external package, and a stored
weight dequantizes to cosine 0.063 against BF16 without it. So an `r8` result
here is a *lower bound* on the vendor build's quality: same width, same scale
granularity, minus the trick that only helps.
"""

from __future__ import annotations

import argparse
import json
import struct
import sys
from pathlib import Path

import torch
from safetensors import safe_open

#: FP4 E2M1 magnitudes and the midpoints that round to them.
E2M1_LEVELS = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0)
NVFP4_GROUP = 16
#: E4M3's largest finite value; the per-tensor scale is chosen so that the
#: per-group scales divided by it land inside this range.
E4M3_MAX = 448.0


def quantize_nvfp4(w: torch.Tensor) -> torch.Tensor:
    out_features, in_features = w.shape
    if in_features % NVFP4_GROUP:
        raise SystemExit(f"nvfp4 needs a multiple of {NVFP4_GROUP}, got {in_features}")

    amax = w.abs().amax()
    if amax == 0:
        return w
    global_scale = amax / (E2M1_LEVELS[-1] * E4M3_MAX)

    groups = w.reshape(out_features, in_features // NVFP4_GROUP, NVFP4_GROUP)
    group_amax = groups.abs().amax(dim=-1, keepdim=True)
    # the group scale is itself stored in FP8, so quantize it before using it
    scale = (group_amax / E2M1_LEVELS[-1] / global_scale).to(torch.float8_e4m3fn)
    effective = scale.float() * global_scale
    effective = torch.where(effective > 0, effective, torch.ones_like(effective))

    levels = torch.tensor(E2M1_LEVELS, device=w.device, dtype=w.dtype)
    midpoints = torch.tensor(E2M1_MIDPOINTS, device=w.device, dtype=w.dtype)
    normalized = groups / effective
    codes = torch.bucketize(normalized.abs(), midpoints, out_int32=True)
    return (torch.sign(normalized) * levels[codes] * effective).reshape(out_features,
                                                                       in_features)


def quantize_int8(w: torch.Tensor) -> torch.Tensor:
    """Symmetric, one FP32 scale per output channel - the vendor's layout."""
    scale = w.abs().amax(dim=-1, keepdim=True) / 127.0
    scale = torch.where(scale > 0, scale, torch.ones_like(scale))
    return (w / scale).round().clamp_(-127, 127) * scale


CODECS = {4.5: quantize_nvfp4, 8.01: quantize_int8}

#: Exponents searched for the AWQ smoothing scale. 0.0 is the identity, so the
#: search can never choose something worse than plain round-to-nearest under its
#: own objective - the baseline is inside the grid rather than outside it.
AWQ_ALPHAS = (0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0)

#: Quantized tensors that legitimately have no activation statistics, so their
#: absence is not a sign the calibration pass missed something:
#:
#: * `embed_tokens` is a lookup. Its "input" is a one-hot row index, so there is
#:   no per-input-channel magnitude to smooth against - the axis the scale would
#:   act on is the vocabulary.
#: * the two aggregate tables live outside the `Gemma4Unified` module tree, so
#:   no forward hook sees them. They are only ever quantized at 8 bits in this
#:   ladder, where round-to-nearest error is 0.87% and smoothing has little to
#:   recover.
SMOOTHING_EXEMPT = frozenset({
    "model.embed_tokens.weight",
    "text_embedding_projection.video_aggregate_embed.weight",
    "text_embedding_projection.audio_aggregate_embed.weight",
})


#: Smallest value the balanced scale may take. This is the knob behind the
#: near-dead blow-up: the smoothed weight is divided by the scale on the way
#: out, so a channel sitting on the floor has its quantization residual
#: multiplied by `1 / floor`. At the historical 1e-5 that is up to 100000x, and
#: `||Q||/||W||` reaches 303 on the real checkpoint. The balancing makes
#: `max * min == 1` before the clamp, so a floor of `f` bounds the
#: amplification at `1 / f`.
AWQ_SCALE_FLOOR = 1e-5


def awq_scale(activation: torch.Tensor, alpha: float,
              floor: float = AWQ_SCALE_FLOOR) -> torch.Tensor:
    """Per-input-channel scale, balanced so it neither inflates nor shrinks overall.

    `floor` bounds how far a near-dead channel can be smoothed, and therefore
    how much the closing division can amplify. It stays in the identity either
    way - the same clamped scale multiplies the weight and divides it back - so
    raising it is a change of smoothing strength, not a correctness risk.
    """
    s = activation.clamp_min(1e-5).pow(alpha)
    return (s / (s.max() * s.min()).sqrt()).clamp_min(floor)


def quantize_smoothed(w: torch.Tensor, activation: torch.Tensor, codec) -> tuple:
    """AWQ-style smoothing, folded entirely into the weight.

    The runtime identity that makes this possible:

        x @ diag(1/s) @ Q(W @ diag(s))^T  ==  x @ [Q(W @ diag(s)) @ diag(1/s)]^T

    ComfyUI spends the left-hand form - it stores `pre_quant_scale` and divides
    the activations - but the right-hand form is an ordinary dense matrix, so a
    simulation can carry the identical error while still loading in stock
    `transformers`. Which also means a packed artifact and this simulation are
    the same arithmetic, not an approximation of each other.

    The scale is chosen per tensor by searching `AWQ_ALPHAS` against an
    activation-weighted error, since a weight column only matters in proportion
    to what multiplies it. The objective is a surrogate - it uses mean |x| per
    channel rather than the real calibration activations - so it ranks
    candidates rather than predicting the end-to-end drift.
    """
    activation = activation.to(w.device, w.dtype)
    best_loss, best_w, best_alpha = None, None, None
    for alpha in AWQ_ALPHAS:
        s = awq_scale(activation, alpha)
        candidate = codec(w * s) / s
        loss = (((candidate - w) * activation) ** 2).sum().item()
        if best_loss is None or loss < best_loss:
            best_loss, best_w, best_alpha = loss, candidate, alpha
    return best_w, best_alpha


#: A tensor whose float32 form exceeds this is quantized on the CPU.
#:
#: `embed_tokens` is [262144, 3840]: 4.03 GB once promoted to float32, before
#: the codec's own temporaries. On a GPU shared with other work that single
#: allocation is the difference between running and not, and these tensors are
#: a handful of one-off conversions where the CPU's slower arithmetic costs
#: seconds rather than minutes.
GPU_TENSOR_LIMIT = 1_000_000_000


def quantize_tensor(tensor: torch.Tensor, codec, device: str) -> torch.Tensor:
    """Run `codec` on `tensor`, choosing a device that will not run out."""
    original = tensor.dtype
    where = "cpu" if tensor.numel() * 4 > GPU_TENSOR_LIMIT else device
    return codec(tensor.to(where, torch.float32)).to("cpu", original)


def read_header(path: Path) -> tuple[dict, dict, int]:
    with path.open("rb") as f:
        n = struct.unpack("<Q", f.read(8))[0]
        header = json.loads(f.read(n))
    metadata = header.pop("__metadata__", {})
    return header, metadata, n + 8


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source", type=Path, required=True)
    parser.add_argument("--plan", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--activation-stats", type=Path,
                        help="from ltx_activation_stats.py; enables AWQ smoothing")
    args = parser.parse_args()

    header, metadata, _ = read_header(args.source)
    plan = json.loads(args.plan.read_text())
    bits = plan["bits"]

    stats = torch.load(args.activation_stats) if args.activation_stats else {}
    if stats:
        quantizable = {n for n, b in bits.items() if b in CODECS}
        uncovered = sorted(quantizable - set(stats) - SMOOTHING_EXEMPT)
        if uncovered:
            raise SystemExit(
                f"{len(uncovered)} quantized tensors have no activation statistics, "
                f"e.g. {uncovered[:3]}. Smoothing them would silently fall back to "
                "round-to-nearest and the rung would not be what it claims.")
        print(f"AWQ smoothing enabled: {len(stats)} modules, "
              f"alphas {AWQ_ALPHAS[0]}..{AWQ_ALPHAS[-1]}")

    missing = sorted(set(header) - set(bits))
    if missing:
        raise SystemExit(f"{len(missing)} tensors have no entry in the plan: {missing[:5]}")

    names = sorted(header)
    offset = 0
    out_header: dict[str, dict] = {}
    for name in names:
        spec = header[name]
        nbytes = spec["data_offsets"][1] - spec["data_offsets"][0]
        out_header[name] = {"dtype": spec["dtype"], "shape": spec["shape"],
                            "data_offsets": [offset, offset + nbytes]}
        offset += nbytes

    blob = json.dumps({**out_header, "__metadata__": metadata},
                      separators=(",", ":")).encode()
    args.output.parent.mkdir(parents=True, exist_ok=True)

    tally: dict[float, int] = {}
    alphas: dict[float, int] = {}
    with safe_open(str(args.source), framework="pt") as src, args.output.open("wb") as out:
        out.write(struct.pack("<Q", len(blob)))
        out.write(blob)
        for i, name in enumerate(names):
            tensor = src.get_tensor(name)
            width = bits[name]
            codec = CODECS.get(width)
            if codec is not None and tensor.dim() == 2:
                if name in stats:
                    original = tensor.dtype
                    work = tensor.to(args.device, torch.float32)
                    work, alpha = quantize_smoothed(work, stats[name], codec)
                    alphas[alpha] = alphas.get(alpha, 0) + 1
                    tensor = work.to("cpu", original)
                else:
                    tensor = quantize_tensor(tensor, codec, args.device)
            elif codec is not None:
                raise SystemExit(f"{name} is {tensor.dim()}-D but the plan asks for {width}")
            tally[width] = tally.get(width, 0) + 1
            out.write(tensor.contiguous().view(torch.uint8).numpy().tobytes())
            if i % 100 == 0:
                print(f"  {i}/{len(names)}", flush=True)

    print(f"wrote {args.output} ({args.output.stat().st_size/1e9:.2f} GB, BF16 on disk)")
    for width, count in sorted(tally.items()):
        label = {4.5: "nvfp4", 8.01: "int8_tensorwise"}.get(width, "untouched")
        print(f"  {count:>4} tensors @ {width:>5} bit  {label}")
    if alphas:
        chosen = ", ".join(f"{a}:{n}" for a, n in sorted(alphas.items()))
        print(f"  smoothing alpha chosen -> {chosen}")
    print(f"packed size for this rung: {plan['predicted_bytes']/1e9:.2f} GB")
    return 0


if __name__ == "__main__":
    sys.exit(main())