code/experiments/exp_accounting.py

"""Claims 2 & 3 — independent bytes-per-parameter accounting.

Recomputes, from dtype widths alone (no numbers copied from measurements):
  (a) Table 1: AdamW 16 -> 7 bytes/param (5 with gradient release), SGD 12 -> 6 (4),
      including the fp16 group-scale overhead Table 1 omits (2/32 bytes per moment).
  (b) Tables 4, 6, 8: Params & Optim GiB rows for Llama-3.1-8B, ResNet-50, GPT-2 124M.
      Param count N is inferred from each table's Reference Params row (N = GiB*2^30/4)
      and cross-checked against the known architecture param count.
  (c) S3.4 checkpoint-size claim: Adam 12 -> 5 bytes/param; 7B: 84 -> 35 GiB.

Paper values are transcribed in PAPER_* constants and NEVER mixed with predictions.
Exit code 0 iff every prediction matches the paper value after rounding to the
paper's printed precision.
"""

import json
import os
import sys

GIB = 2**30
SCALE = 2.0 / 32.0  # fp16 scale per group of 32 elements, per quantized moment

# ---------- (a) Table 1 accounting ----------------------------------------
def table1():
    rows = {
        "SGD":       {"master": 4, "corr": 0, "grad": 4, "m": 4, "v": 0},
        "FlashSGD":  {"master": 2, "corr": 1, "grad": 2, "m": 1, "v": 0},
        "Adam":      {"master": 4, "corr": 0, "grad": 4, "m": 4, "v": 4},
        "FlashAdam": {"master": 2, "corr": 1, "grad": 2, "m": 1, "v": 1},
    }
    paper_totals = {"SGD": 12, "FlashSGD": 6, "Adam": 16, "FlashAdam": 7}
    paper_gr = {"FlashSGD": 4, "FlashAdam": 5}  # gradient release totals
    out, ok = {}, True
    for k, r in rows.items():
        total = sum(r.values())
        n_moments = (r["m"] == 1) + (r["v"] == 1)
        exact = total + n_moments * SCALE if k.startswith("Flash") else total
        match = round(total) == paper_totals[k]
        ok &= match
        out[k] = {"paper_total": paper_totals[k], "pred_total_table1": total,
                  "pred_total_with_scales": round(exact, 4), "match": match}
        if k in paper_gr:
            gr_exact = exact - r["grad"]
            match_gr = round(gr_exact - n_moments * SCALE) == paper_gr[k]
            ok &= match_gr
            out[k]["paper_total_grad_release"] = paper_gr[k]
            out[k]["pred_total_grad_release_with_scales"] = round(gr_exact, 4)
            out[k]["match_grad_release"] = match_gr
    return out, ok


# ---------- (b) Tables 4/6/8 Params & Optim rows ---------------------------
# bytes/param for each variant's Params row (master weights held by model) and
# Optim row (optimizer-held state: moments (+scales) + ECC correction term).
VARIANTS = {
    # (params_bytes, optim_bytes(n_moments), has_ecc)
    "Reference":    lambda nm: (4, 4 * nm),
    "FlashOptim":   lambda nm: (2, (1 + SCALE) * nm + 1),
    "Weight Split": lambda nm: (2, 4 * nm + 1),
    "Opt. Quant.":  lambda nm: (4, (1 + SCALE) * nm),
}

PAPER_TABLES = {
    # table: (n_moments, arch_param_count, {variant: (params_gib, optim_gib)})
    "Table 4 Llama-3.1-8B AdamW": (2, 8_030_261_248, {
        "Reference": (29.9, 59.8), "FlashOptim": (15.0, 23.4),
        "Weight Split": (15.0, 67.3), "Opt. Quant.": (29.9, 15.9)}),
    "Table 6 ResNet-50 SGD": (1, 25_557_032, {
        "Reference": (0.10, 0.10), "FlashOptim": (0.05, 0.05),
        "Weight Split": (0.05, 0.12), "Opt. Quant.": (0.10, 0.03)}),
    "Table 6 ResNet-50 AdamW": (2, 25_557_032, {
        "Reference": (0.10, 0.19), "FlashOptim": (0.05, 0.08),
        "Weight Split": (0.05, 0.21), "Opt. Quant.": (0.10, 0.05)}),
    "Table 8 GPT-2 124M AdamW": (2, 124_439_808, {
        "Reference": (0.46, 0.93), "FlashOptim": (0.23, 0.36),
        "Weight Split": (0.23, 1.04), "Opt. Quant.": (0.46, 0.25)}),
    "Table 8 GPT-2 124M Lion": (1, 124_439_808, {
        "Reference": (0.46, 0.46), "FlashOptim": (0.23, 0.24),
        "Weight Split": (0.23, 0.58), "Opt. Quant.": (0.46, 0.12)}),
}


def tables_468():
    out, ok = {}, True
    for tname, (nm, arch_n, rows) in PAPER_TABLES.items():
        # infer N from Reference Params row; cross-check vs architecture count
        n_inferred = rows["Reference"][0] * GIB / 4
        res = {"N_inferred_from_paper": int(n_inferred), "N_architecture": arch_n,
               "N_rel_dev": round(abs(n_inferred - arch_n) / arch_n, 4), "rows": {}}
        for v, (p_params, p_optim) in rows.items():
            pb, ob = VARIANTS[v](nm)
            pred_p, pred_o = arch_n * pb / GIB, arch_n * ob / GIB

            # match = prediction rounds to the paper value at the paper's own
            # printed decimal count (e.g. "1.04" -> 2 decimals, "29.9" -> 1)
            def decimals(x):
                s = repr(x)
                return len(s.split(".")[1]) if "." in s else 0
            m = (round(pred_p, decimals(p_params)) == p_params and
                 round(pred_o, decimals(p_optim)) == p_optim)
            ok &= m
            dev = max(abs(pred_p - p_params) / p_params, abs(pred_o - p_optim) / p_optim)
            res["rows"][v] = {"paper": [p_params, p_optim],
                              "pred": [round(pred_p, 3), round(pred_o, 3)],
                              "max_rel_dev": round(dev, 4), "match": m}
        out[tname] = res
    return out, ok


# ---------- (c) checkpoint-size claim --------------------------------------
def checkpoints():
    ref = 4 + 4 + 4            # weights + m + v (S3.4: "12 bytes per parameter")
    flash = 2 + 1 + 1 + 1      # bf16 weights + ecc + m + v ("5 bytes")
    ok = (ref == 12) and (flash == 5)
    n7 = 84 * GIB / 12         # param count implied by "84 GiB" at 12 B/param
    pred35 = n7 * flash / GIB
    ok &= round(pred35) == 35
    halved = flash / ref < 0.5  # "cut checkpoint sizes by more than half"
    ok &= halved
    return {"ref_bytes": ref, "flash_bytes": flash, "flash_with_scales": flash + 2 * SCALE,
            "implied_7B_params": int(n7), "pred_7B_flash_gib": round(pred35, 2),
            "paper_7B_flash_gib": 35, "reduction_pct": round(100 * (1 - flash / ref), 1)}, ok


if __name__ == "__main__":
    t1, ok1 = table1()
    t468, ok2 = tables_468()
    ck, ok3 = checkpoints()
    n_rows = sum(len(r["rows"]) for r in t468.values())
    n_match = sum(int(x["match"]) for r in t468.values() for x in r["rows"].values())
    deviations = {f"{t} / {v}": x for t, r in t468.items()
                  for v, x in r["rows"].items() if not x["match"]}
    result = {"table1": t1, "tables_4_6_8": t468, "checkpoint_claim": ck,
              "table1_all_match": ok1, "checkpoint_all_match": ok3,
              "tables_rows_matching": f"{n_match}/{n_rows}",
              "deviating_rows": deviations}
    os.makedirs("repro_flashoptim/outputs", exist_ok=True)
    with open("repro_flashoptim/outputs/accounting.json", "w") as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))
    print(f"\nSUMMARY: Table1 match={ok1}; checkpoint claim match={ok3}; "
          f"Tables 4/6/8 rows matching at printed precision: {n_match}/{n_rows}")
    for k, v in deviations.items():
        print(f"  DEVIATION: {k}: paper {v['paper']} vs pred {v['pred']} "
              f"(max rel dev {100*v['max_rel_dev']:.1f}%)")
    # exit 0 = computation completed; deviations are findings, not failures
    sys.exit(0)