Add convrot4 and convrot8 quantizations to the ui. A lot of improvements to the convrot quantization methods.

scripts/test_quantizations.py

@@ -0,0 +1,339 @@
+"""
+Reproducible speed / VRAM / accuracy benchmarks for the toolkit quantization
+backends.
+
+Compares bf16 against the custom OstrisLinear backends (convrot8, convrot4 for
+now; add more qtypes to QTYPES as they land).
+
+Measures, per qtype:
+  - layer inference latency across DiT-representative shapes (vs bf16)
+  - layer training latency (forward + backward through the frozen layer)
+  - VRAM on a transformer-ish block stack: resident weights, peak during a
+    no-grad forward, peak during a train step
+  - accuracy drift vs bf16: output relative error per layer shape and
+    accumulated through the block stack
+  - weight reconstruction error and one-time quantize (conversion) time
+
+Usage:
+    python scripts/test_quantizations.py --gpu 1
+    python scripts/test_quantizations.py --gpu 1 --qtypes bf16 convrot8
+"""
+
+import argparse
+import math
+import os
+import sys
+import time
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import torch  # noqa: E402
+
+# (tokens, in_features, out_features) — FLUX/Wan-class projections
+SPEED_SHAPES = [
+    (4096, 3072, 3072),
+    (4096, 3072, 12288),
+    (4096, 12288, 3072),
+    (1024, 3072, 12288),
+]
+
+# block stack used for the vram/drift tests (mimics a DiT block's linears)
+VRAM_BLOCKS = 8
+VRAM_BLOCK_SHAPES = [(3072, 12288), (12288, 3072), (3072, 3072), (3072, 3072)]
+VRAM_TOKENS = 4096
+
+QTYPES = ["bf16", "qfloat8", "float8", "orbit4", "orbitvq4", "convrot8", "convrot4"]
+
+STACK_KEY = f"{VRAM_BLOCKS}-block stack"
+
+
+def convert(module: torch.nn.Linear, qtype: str) -> torch.nn.Linear:
+    """Quantize a linear with the given qtype. Returns the (possibly replaced)
+    module — quanto swaps the module object, the ostris backends convert in place."""
+    if qtype == "bf16":
+        return module
+    from toolkit.util.ostris_quant import convert_linear_to_ostris, get_ostris_quantizer
+
+    q = get_ostris_quantizer(qtype)
+    if q is not None:
+        assert convert_linear_to_ostris(module, q), f"conversion refused for {qtype}"
+        return module
+
+    # quanto / torchao qtypes go through the shared toolkit quantize flow; use a
+    # holder so quanto's module replacement has a parent to swap into
+    from optimum.quanto import freeze
+    from toolkit.util.quantize import get_qtype, quantize
+
+    holder = torch.nn.Sequential(module)
+    quantize(holder, weights=get_qtype(qtype))
+    freeze(holder)
+    return holder[0]
+
+
+def fp_weight(module: torch.nn.Linear) -> torch.Tensor:
+    """Dequantized weight in float32, whatever the backend."""
+    if hasattr(module, "dequantize_weight"):
+        return module.dequantize_weight().float()
+    w = module.weight
+    if hasattr(w, "dequantize"):
+        return w.dequantize().float()
+    return w.detach().float()
+
+
+def bench(fn, iters: int, device) -> float:
+    for _ in range(max(3, iters // 5)):
+        fn()
+    torch.cuda.synchronize(device)
+    t0 = time.perf_counter()
+    for _ in range(iters):
+        fn()
+    torch.cuda.synchronize(device)
+    return (time.perf_counter() - t0) / iters * 1000  # ms
+
+
+def gb(nbytes: int) -> str:
+    return f"{nbytes / 1e9:6.2f} GB"
+
+
+def make_layer(k: int, n: int, device) -> torch.nn.Linear:
+    lin = torch.nn.Linear(k, n, bias=True, dtype=torch.bfloat16, device=device)
+    with torch.no_grad():
+        lin.weight.mul_(0.02)
+    return lin
+
+
+def make_stack(device) -> torch.nn.ModuleList:
+    # default nn.Linear init (~1/sqrt(in) std) so block branches contribute at a
+    # realistic O(1) scale to the residual stream — scaling weights down further
+    # makes accumulated quantization drift look artificially tiny
+    torch.manual_seed(0)
+    blocks = torch.nn.ModuleList()
+    for _ in range(VRAM_BLOCKS):
+        blocks.append(torch.nn.ModuleList([
+            torch.nn.Linear(k, n, bias=True, dtype=torch.bfloat16, device=device)
+            for k, n in VRAM_BLOCK_SHAPES
+        ]))
+    return blocks
+
+
+def stack_forward(blocks, x):
+    # pre-norm residual blocks like a real transformer, so activations stay at a
+    # sane scale and quantization drift accumulates realistically across depth
+    h = x
+    for b in blocks:
+        r = torch.nn.functional.layer_norm(h, h.shape[-1:])
+        r = b[0](r)          # 3072 -> 12288
+        r = torch.nn.functional.gelu(r)
+        h = h + b[1](r)      # 12288 -> 3072
+        r = torch.nn.functional.layer_norm(h, h.shape[-1:])
+        h = h + b[3](b[2](r))  # 3072 -> 3072 -> 3072
+    return h
+
+
+def run_speed(qtype: str, device, iters: int, results: dict):
+    for m, k, n in SPEED_SHAPES:
+        torch.manual_seed(0)
+        lin = make_layer(k, n, device)
+        lin = convert(lin, qtype)
+        x = torch.randn(m, k, device=device, dtype=torch.bfloat16)
+
+        with torch.no_grad():
+            t_inf = bench(lambda: lin(x), iters, device)
+
+        def train_step():
+            xi = x.detach().requires_grad_(True)
+            lin(xi).sum().backward()
+
+        t_train = bench(train_step, max(10, iters // 3), device)
+        results[(qtype, "inf", (m, k, n))] = t_inf
+        results[(qtype, "train", (m, k, n))] = t_train
+        torch.cuda.empty_cache()
+
+
+def run_vram(qtype: str, device, results: dict):
+    torch.cuda.empty_cache()
+    base = torch.cuda.memory_allocated(device)
+
+    blocks = make_stack(device)
+    for b in blocks:
+        for i in range(len(b)):
+            b[i] = convert(b[i], qtype)
+    torch.cuda.empty_cache()
+    results[(qtype, "vram_weights")] = torch.cuda.memory_allocated(device) - base
+
+    x = torch.randn(VRAM_TOKENS, 3072, device=device, dtype=torch.bfloat16)
+
+    # no-grad forward peak (sampling); warm up first so lazy-init allocations are
+    # not counted as steady-state peak
+    with torch.no_grad():
+        stack_forward(blocks, x)
+    torch.cuda.synchronize(device)
+    torch.cuda.reset_peak_memory_stats(device)
+    with torch.no_grad():
+        stack_forward(blocks, x)
+    torch.cuda.synchronize(device)
+    results[(qtype, "vram_fwd_peak")] = torch.cuda.max_memory_allocated(device) - base
+
+    # train step peak (frozen base; grads flow to the input like lora training)
+    def train_step():
+        xi = x.detach().requires_grad_(True)
+        stack_forward(blocks, xi).float().pow(2).mean().backward()
+
+    train_step()
+    torch.cuda.synchronize(device)
+    torch.cuda.reset_peak_memory_stats(device)
+    train_step()
+    torch.cuda.synchronize(device)
+    results[(qtype, "vram_train_peak")] = torch.cuda.max_memory_allocated(device) - base
+
+    blocks = x = None  # release before the allocator accounting of the next run
+    torch.cuda.empty_cache()
+
+
+def run_drift(qtype: str, device, results: dict):
+    """Output error vs the bf16 reference, per layer shape and through the stack."""
+    for m, k, n in SPEED_SHAPES:
+        torch.manual_seed(0)
+        lin = make_layer(k, n, device)
+        x = torch.randn(m, k, device=device, dtype=torch.bfloat16)
+        with torch.no_grad():
+            y_ref = lin(x).float()
+            lin = convert(lin, qtype)
+            y_q = lin(x).float()
+        results[(qtype, "drift", (m, k, n))] = ((y_q - y_ref).norm() / y_ref.norm()).item()
+        torch.cuda.empty_cache()
+
+    blocks = make_stack(device)
+    x = torch.randn(VRAM_TOKENS, 3072, device=device, dtype=torch.bfloat16)
+    with torch.no_grad():
+        y_ref = stack_forward(blocks, x).float()
+        for b in blocks:
+            for i in range(len(b)):
+                b[i] = convert(b[i], qtype)
+        y_q = stack_forward(blocks, x).float()
+    results[(qtype, "drift", STACK_KEY)] = ((y_q - y_ref).norm() / y_ref.norm()).item()
+    blocks = x = None
+    torch.cuda.empty_cache()
+
+
+def run_quality_and_quantize_time(qtype: str, device, results: dict):
+    torch.manual_seed(0)
+    lin = make_layer(3072, 3072, device)
+    w0 = lin.weight.detach().float().clone()
+    torch.cuda.synchronize(device)
+    t0 = time.perf_counter()
+    lin = convert(lin, qtype)
+    torch.cuda.synchronize(device)
+    results[(qtype, "quantize_ms")] = (time.perf_counter() - t0) * 1000
+    if qtype == "bf16":
+        results[(qtype, "weight_err")] = 0.0
+    else:
+        wq = fp_weight(lin)
+        results[(qtype, "weight_err")] = ((wq - w0).norm() / w0.norm()).item()
+    torch.cuda.empty_cache()
+
+
+def print_speed_table(title: str, kind: str, qts, results):
+    print(f"\n=== {title} (ms; speedup vs bf16) ===")
+    print(f"{'M x K -> N':<22}" + "".join(f"{qt:>18}" for qt in qts))
+    for shape in SPEED_SHAPES:
+        m, k, n = shape
+        row = f"{f'{m} x {k} -> {n}':<22}"
+        ref = results.get(("bf16", kind, shape))
+        for qt in qts:
+            t = results.get((qt, kind, shape))
+            if t is None:
+                row += f"{'-':>18}"
+                continue
+            sp = f" ({ref / t:4.2f}x)" if ref and qt != "bf16" else " " * 8
+            row += f"{t:8.3f}ms{sp}"
+        print(row)
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--gpu", type=int, default=0, help="cuda device id to run on")
+    ap.add_argument("--qtypes", nargs="+", default=QTYPES, help=f"subset of {QTYPES}")
+    ap.add_argument("--iters", type=int, default=50, help="timing iterations per case")
+    args = ap.parse_args()
+
+    device = torch.device(f"cuda:{args.gpu}")
+    torch.cuda.set_device(device)
+    props = torch.cuda.get_device_properties(device)
+    print(f"device: cuda:{args.gpu} ({props.name}, sm_{props.major}{props.minor}, "
+          f"{props.total_memory / 1e9:.0f} GB)")
+    print(f"torch {torch.__version__}\n")
+
+    # warm the toolkit import chain (module imports + custom-op registration) so it
+    # isn't charged to the first qtype's quantize timing
+    if any(qt != "bf16" for qt in args.qtypes):
+        from toolkit.util.ostris_quant import get_ostris_quantizer
+        for qt in args.qtypes:
+            if qt != "bf16":
+                get_ostris_quantizer(qt)
+
+    results = {}
+    for qt in args.qtypes:
+        print(f"benchmarking {qt} ...")
+        run_quality_and_quantize_time(qt, device, results)
+        run_drift(qt, device, results)
+        run_speed(qt, device, args.iters, results)
+        run_vram(qt, device, results)
+
+    qts = args.qtypes
+    print_speed_table("layer latency, inference", "inf", qts, results)
+    print_speed_table("layer latency, train fwd+bwd", "train", qts, results)
+
+    print(f"\n=== vram on the block stack ({VRAM_BLOCKS} blocks, {VRAM_TOKENS} tokens) ===")
+    print(f"{'':<28}" + "".join(f"{qt:>18}" for qt in qts))
+    for key, label in (("vram_weights", "weights resident"),
+                       ("vram_fwd_peak", "peak, no-grad fwd"),
+                       ("vram_train_peak", "peak, train step")):
+        row = f"{label:<28}"
+        for qt in qts:
+            row += f"{gb(results[(qt, key)]):>18}"
+        print(row)
+
+    print("\n=== accuracy drift vs bf16 (output rel err, no-grad) ===")
+    print(f"{'':<28}" + "".join(f"{qt:>18}" for qt in qts))
+    for shape in SPEED_SHAPES + [STACK_KEY]:
+        label = f"{shape[0]} x {shape[1]} -> {shape[2]}" if isinstance(shape, tuple) else shape
+        row = f"{label:<28}"
+        for qt in qts:
+            row += f"{results[(qt, 'drift', shape)]:>18.5f}"
+        print(row)
+
+    print("\n=== quantization ===")
+    print(f"{'':<28}" + "".join(f"{qt:>18}" for qt in qts))
+    row = f"{'weight rel err':<28}"
+    for qt in qts:
+        row += f"{results[(qt, 'weight_err')]:>18.5f}"
+    print(row)
+    row = f"{'quantize time (ms)':<28}"
+    for qt in qts:
+        row += f"{results[(qt, 'quantize_ms')]:>18.1f}"
+    print(row)
+
+    # ---- clean per-qtype breakdown: speed (geomean over shapes) + accuracy ----
+    def geomean_speedup(qt, kind):
+        logs = []
+        for shape in SPEED_SHAPES:
+            ref = results.get(("bf16", kind, shape))
+            t = results.get((qt, kind, shape))
+            if ref and t:
+                logs.append(math.log(ref / t))
+        return math.exp(sum(logs) / len(logs)) if logs else float("nan")
+
+    print("\n=== summary (speed = geomean speedup vs bf16; drift lower is better) ===")
+    print(f"{'':<12}{'inference':>18}{'train':>18}{'accuracy drift':>20}{'max vram':>16}")
+    for qt in qts:
+        max_vram = max(results[(qt, "vram_fwd_peak")], results[(qt, "vram_train_peak")])
+        print(f"{qt:<12}"
+              f"{geomean_speedup(qt, 'inf'):>17.2f}x"
+              f"{geomean_speedup(qt, 'train'):>17.2f}x"
+              f"{results[(qt, 'drift', STACK_KEY)]:>20.5f}"
+              f"{gb(max_vram):>16}")
+
+
+if __name__ == "__main__":
+    main()