code/experiments/make_poster_figs.py

"""Poster figures (SVG, matplotlib) from logged experiment outputs.

fig_bytes.svg : Table 1 accounting as stacked bytes/param bars (exact analytic)
fig_fig3.svg  : exhaustive FP32 reconstruction error, bf16 + fp16 panels
fig_traj.svg  : char-LM trajectories + divergence   (needs charlm_losses.csv)
fig_nmse.svg  : companded vs linear NMSE quantiles  (needs charlm_nmse.csv)
"""

import csv
import json
import os
import sys

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

INK, MUTED, GRID = "#0b0b0b", "#898781", "#e1e0d9"
C = {"blue": "#2a78d6", "green": "#008300", "magenta": "#e87ba4",
     "yellow": "#eda100", "aqua": "#1baf7a", "orange": "#eb6834",
     "red": "#e34948", "violet": "#4a3aa7"}
OUT = "repro_flashoptim/outputs/poster"
os.makedirs(OUT, exist_ok=True)
plt.rcParams.update({"font.family": "sans-serif", "font.size": 11,
                     "axes.edgecolor": MUTED, "axes.labelcolor": INK,
                     "xtick.color": MUTED, "ytick.color": MUTED,
                     "axes.grid": True, "grid.color": GRID, "grid.linewidth": 0.6,
                     "svg.fonttype": "none"})


def fig_bytes():
    # exact analytic accounting incl. fp16 group scales
    comp_order = ["Master wts", "Correction", "Gradients", "Momentum", "Variance", "Scales"]
    colors = [C["blue"], C["violet"], C["yellow"], C["green"], C["aqua"], C["magenta"]]
    bars = {
        "AdamW\n(reference)": [4, 0, 4, 4, 4, 0],
        "FlashAdamW": [2, 1, 2, 1, 1, 0.125],
        "FlashAdamW\n+grad release": [2, 1, 0, 1, 1, 0.125],
        "SGD\n(reference)": [4, 0, 4, 4, 0, 0],
        "FlashSGD": [2, 1, 2, 1, 0, 0.0625],
    }
    fig, ax = plt.subplots(figsize=(7.2, 4.4))
    x = range(len(bars))
    bottoms = [0.0] * len(bars)
    for ci, comp in enumerate(comp_order):
        vals = [v[ci] for v in bars.values()]
        ax.bar(x, vals, bottom=bottoms, width=0.62, color=colors[ci],
               label=comp, edgecolor="white", linewidth=1.2)
        bottoms = [b + v for b, v in zip(bottoms, vals)]
    for xi, (k, v) in zip(x, bars.items()):
        t = sum(v)
        ax.text(xi, t + 0.25, f"{t:.3g}", ha="center", fontsize=12,
                fontweight="bold", color=INK)
    ax.set_xticks(list(x), list(bars.keys()), fontsize=10, color=INK)
    ax.set_ylabel("bytes per parameter")
    ax.set_ylim(0, 18)
    ax.legend(ncol=3, frameon=False, fontsize=9, loc="upper right")
    ax.set_axisbelow(True)
    ax.grid(axis="x", visible=False)
    fig.tight_layout()
    fig.savefig(f"{OUT}/fig_bytes.svg")
    plt.close(fig)


def fig_fig3():
    rows = list(csv.DictReader(open("repro_flashoptim/outputs/fig3_reconstruction.csv")))
    names = {"none": ("No correction", C["blue"]),
             "float": ("+ float error term (prior work)", C["green"]),
             "ulp_int8": ("+ ULP INT8 — ours, 24-bit", C["magenta"]),
             "ulp_int16": ("+ ULP INT16 — ours, 32-bit", C["yellow"])}
    fig, axes = plt.subplots(2, 1, figsize=(7.2, 5.6), sharex=False)
    for ax, target, title in zip(axes, ["bf16", "fp16"],
                                 ["Target BF16", "Target FP16"]):
        for m, (label, color) in names.items():
            pts = sorted((int(r["exp_unbiased"]), float(r["mean_rel_err"]))
                         for r in rows
                         if r["target"] == target and r["method"] == m
                         and int(r["exp_biased"]) > 0 and float(r["mean_rel_err"]) > 0)
            ax.plot([p[0] for p in pts], [p[1] for p in pts], color=color,
                    lw=1.8, label=label if target == "bf16" else None)
        ax.set_yscale("log")
        ax.set_title(title, fontsize=11, color=INK, loc="left")
        ax.set_ylabel("mean rel. error")
        for xb in ([-126] if target == "bf16" else [-14, 15]):
            ax.axvline(xb, ls=":", color=MUTED, lw=1)
    axes[0].legend(frameon=False, fontsize=9, loc="lower left", ncol=2)
    axes[1].set_xlabel("fp32 exponent (unbiased)")
    fig.suptitle("FP32 reconstruction error — exhaustive over all 4.28B finite values",
                 fontsize=12, color=INK, x=0.02, ha="left")
    fig.tight_layout(rect=(0, 0, 1, 0.96))
    fig.savefig(f"{OUT}/fig_fig3.svg")
    plt.close(fig)


def fig_traj():
    if not os.path.exists("repro_flashoptim/outputs/charlm_losses.csv"):
        print("charlm_losses.csv missing; skip fig_traj")
        return
    rows = list(csv.DictReader(open("repro_flashoptim/outputs/charlm_losses.csv")))
    steps = [int(r["step"]) for r in rows]

    def ema(key, a=0.05):
        out, m = [], None
        for r in rows:
            x = float(r[key])
            if x != x:  # NaN
                out.append(float("nan")); continue
            m = x if m is None else (1 - a) * m + a * x
            out.append(m)
        return out
    fig, ax = plt.subplots(figsize=(7.2, 4.0))
    for key, label, color, lw in [
            ("ref_s0", "AdamW fp32 (ref, seed 0)", C["blue"], 2.0),
            ("ref_s1", "AdamW fp32 (ref, seed 1)", C["aqua"], 1.4),
            ("flash_s0", "FlashAdamW sim (companded, seed 0)", C["green"], 2.0),
            ("linear_s0", "linear 8-bit states (no companding)", C["red"], 2.0)]:
        ax.plot(steps, ema(key), color=color, lw=lw, label=label)
    ax.set_xlabel("step"); ax.set_ylabel("train loss (EMA)")
    ax.legend(frameon=False, fontsize=9)
    ax.set_title("Char-LM on Tiny Shakespeare (2.4M params, CPU, identical data order)",
                 fontsize=11, color=INK, loc="left")
    fig.tight_layout(); fig.savefig(f"{OUT}/fig_traj.svg"); plt.close(fig)


def fig_nmse():
    if not os.path.exists("repro_flashoptim/outputs/charlm_nmse.csv"):
        print("charlm_nmse.csv missing; skip fig_nmse")
        return
    rows = list(csv.DictReader(open("repro_flashoptim/outputs/charlm_nmse.csv")))
    import statistics
    keys = [("m_linear", "momentum, linear", C["red"]),
            ("m_companded", "momentum, companded", C["green"]),
            ("v_linear", "variance, linear", C["orange"]),
            ("v_companded", "variance, companded", C["aqua"])]
    data = [[float(r[k]) for r in rows] for k, _, _ in keys]
    fig, ax = plt.subplots(figsize=(7.2, 3.6))
    bp = ax.boxplot(data, vert=False, tick_labels=[l for _, l, _ in keys],
                    patch_artist=True, showfliers=False, widths=0.55)
    for patch, (_, _, color) in zip(bp["boxes"], keys):
        patch.set_facecolor(color); patch.set_alpha(0.75); patch.set_edgecolor(INK)
    for med in bp["medians"]:
        med.set_color(INK)
    ax.set_xscale("log")
    ax.set_xlabel("NMSE of quantize→dequantize roundtrip (log)")
    ax.set_title("Optimizer-state quantization error on real AdamW states",
                 fontsize=11, color=INK, loc="left")
    fig.tight_layout(); fig.savefig(f"{OUT}/fig_nmse.svg"); plt.close(fig)


if __name__ == "__main__":
    fig_bytes(); fig_fig3(); fig_traj(); fig_nmse()
    print("wrote", os.listdir(OUT))