code/gpu/equivalence.py

"""Claim 1 & 6 — elementwise equivalence: authors' Triton kernels vs the
independent eager reimplementation (flashsim, written from the paper only).

Runs on CUDA. Compares:
  A. quantize/dequantize (momentum spec: signed+softsign; variance spec:
     unsigned+sqrt) on a battery of tensors -> int-code agreement rate,
     dequant NMSE (theirs vs ours), scale agreement.
  B. compute_ecc_bits / reconstruct_fp32_param vs flashsim compress/decompress
     -> rho agreement, reconstruction agreement, and the AUTHORS' bitwise-exact
     fraction + mean rel err on a stratified exponent sweep (resolves the
     99.92% question from the local exhaustive run).
"""

import json

import torch


def battery(device):
    g = torch.Generator(device="cpu").manual_seed(7)
    ts = {
        "normal": torch.randn(1 << 20, generator=g),
        "heavy_tail": torch.randn(1 << 20, generator=g) * torch.exp(
            torch.randn(1 << 20, generator=g) * 3),
        "tiny": torch.randn(1 << 16, generator=g) * 1e-20,
        "adam_v_like": (torch.randn(1 << 20, generator=g) * 1e-3) ** 2,
        "nonmult32": torch.randn((1 << 16) + 17, generator=g),
        "const_groups": torch.ones(1 << 12) * 0.37,
        "with_zeros": torch.where(torch.rand(1 << 16, generator=g) < 0.3,
                                  torch.zeros(1), torch.randn(1 << 16, generator=g)),
    }
    return {k: v.to(device) for k, v in ts.items()}


def eq_quantization(results):
    import flashoptim.optimizers as fo
    from flashsim import quantization as Q
    specs = {"momentum": dict(signed=True, sqrt=False, softsign=True),
             "variance": dict(signed=False, sqrt=True, softsign=False)}
    out = {}
    for tname, t in battery("cuda").items():
        for sname, sp in specs.items():
            x = t.abs() if sname == "variance" else t
            q_a, s_a = fo.quantize(x, **sp)                     # authors (triton)
            q_m, s_m = Q.quantize(x.cpu(), **sp)                # ours (eager cpu)
            agree = float((q_a.cpu().reshape(-1) == q_m.reshape(-1)).float().mean())
            off1 = float(((q_a.cpu().reshape(-1).to(torch.int16) -
                           q_m.reshape(-1).to(torch.int16)).abs() <= 1).float().mean())
            d_a = fo.dequantize(q_a, s_a, **sp).cpu().reshape(-1)[:x.numel()]
            d_m = Q.dequantize(q_m, s_m, **sp)[:x.numel()]
            xf = x.cpu().float().reshape(-1)
            nmse_a, nmse_m = Q.nmse(xf, d_a), Q.nmse(xf, d_m)
            out[f"{tname}/{sname}"] = dict(
                int_agree=agree, int_within_1=off1,
                nmse_authors=nmse_a, nmse_ours=nmse_m,
                nmse_ratio=(nmse_a / nmse_m) if nmse_m else 1.0)
    results["quantization_equivalence"] = out


def eq_ecc(results):
    import flashoptim as fo
    from flashsim import splitting as S
    torch.manual_seed(11)
    # stratified sweep: 2^16 random mantissas per fp32 exponent bucket
    per_exp = 1 << 16
    exps = list(range(1, 255, 4))
    stats = {"n": 0, "authors_exact": 0, "ours_exact": 0, "rho_agree": 0,
             "recon_agree": 0, "authors_err_sum": 0.0, "ours_err_sum": 0.0}
    for e in exps:
        mant = torch.randint(0, 1 << 23, (per_exp,), dtype=torch.int32)
        sign = torch.randint(0, 2, (per_exp,), dtype=torch.int32) << 31
        bits = sign | (e << 23) | mant
        x = bits.view(torch.float32).cuda()
        xb = x.to(torch.bfloat16)
        rho_a = fo.compute_ecc_bits(x, xb, master_bytewidth=3)
        xa = fo.reconstruct_fp32_param(xb, rho_a)
        xb_m, rho_m = S.compress(x.cpu(), torch.bfloat16, torch.int8)
        xm = S.decompress(xb_m, rho_m)
        stats["n"] += per_exp
        stats["authors_exact"] += int((xa == x).sum())
        stats["ours_exact"] += int((xm == x.cpu()).sum())
        stats["rho_agree"] += int((rho_a.cpu() == rho_m).sum())
        stats["recon_agree"] += int((xa.cpu() == xm).sum())
        stats["authors_err_sum"] += float(((xa - x).abs() / x.abs()).sum())
        stats["ours_err_sum"] += float(((xm - x.cpu()).abs() / x.cpu().abs()).sum())
    # int16 (32-bit master) exact-fraction check on the same protocol
    st16 = {"n": 0, "authors_exact": 0, "ours_exact": 0}
    for e in exps:
        mant = torch.randint(0, 1 << 23, (per_exp,), dtype=torch.int32)
        bits = (e << 23) | mant
        x = bits.view(torch.float32).cuda()
        xb = x.to(torch.bfloat16)
        rho_a = fo.compute_ecc_bits(x, xb, master_bytewidth=4)
        xa = fo.reconstruct_fp32_param(xb, rho_a)
        xb_m, rho_m = S.compress(x.cpu(), torch.bfloat16, torch.int16)
        xm = S.decompress(xb_m, rho_m)
        st16["n"] += per_exp
        st16["authors_exact"] += int((xa == x).sum())
        st16["ours_exact"] += int((xm == x.cpu()).sum())
    n = stats["n"]
    results["ecc_equivalence"] = dict(
        n=n,
        int8=dict(authors_exact_frac=stats["authors_exact"] / n,
                  ours_exact_frac=stats["ours_exact"] / n,
                  rho_agree_frac=stats["rho_agree"] / n,
                  recon_agree_frac=stats["recon_agree"] / n,
                  authors_mean_rel_err=stats["authors_err_sum"] / n,
                  ours_mean_rel_err=stats["ours_err_sum"] / n),
        int16=dict(authors_exact_frac=st16["authors_exact"] / st16["n"],
                   ours_exact_frac=st16["ours_exact"] / st16["n"]))


def run_all():
    results = {}
    eq_quantization(results)
    eq_ecc(results)
    return results


if __name__ == "__main__":
    print(json.dumps(run_all(), indent=2))