code/experiments/exp_charlm_nmse_divergence.py

"""Claims 4 & 6 (local CPU evidence, mechanism-level) — char-LM on real text.

Trains a small character-level transformer on Tiny Shakespeare (real text,
~1.1 MB) with four arms sharing IDENTICAL data order:

  ref_s0    : fp32 AdamW (torch.optim.AdamW), init seed 0     (reference)
  flash_s0  : FlashAdamWSim (independent Alg. 4 impl): companded 8-bit states
              + simulated 24-bit master weights, init seed 0
  linear_s0 : same but LINEAR state quantization (no softsign, no sqrt) —
              the paper's Fig. 5 ablation
  ref_s1    : fp32 AdamW, init seed 1 — seed-noise yardstick for trajectories

Claim 4 (local): |loss(flash_s0) - loss(ref_s0)| compared against the seed
yardstick |loss(ref_s1) - loss(ref_s0)| step-by-step.
Claim 6b (Fig. 4): NMSE of companded vs linear quantization roundtrips of the
REFERENCE run's exp_avg / exp_avg_sq buffers harvested every step.
Claim 6c (Fig. 5, toy scale): does linear_s0 diverge while flash_s0 tracks?

Scale honesty: ~2.4M params, 500 steps, CPU — mechanism test, NOT the paper's
GPT-2 124M / 20k steps. The GPU job repeats the divergence test with the
authors' actual Triton kernels at larger scale.
"""

import json
import math
import os
import sys
import urllib.request

import torch
import torch.nn as nn
import torch.nn.functional as F

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from flashsim import quantization as Q  # noqa: E402
from flashsim.flash_adamw_sim import FlashAdamWSim  # noqa: E402

DATA_URL = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
DATA_PATH = "repro_flashoptim/data/tinyshakespeare.txt"
OUT = "repro_flashoptim/outputs"
CTX, BATCH, STEPS, LR = 256, 32, 500, 1e-3
EVERY = 10  # harvest/  log cadence


class Block(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, h, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x, mask):
        a, _ = self.attn(self.ln1(x), self.ln1(x), self.ln1(x), attn_mask=mask,
                         need_weights=False)
        x = x + a
        return x + self.mlp(self.ln2(x))


class CharLM(nn.Module):
    def __init__(self, vocab, d=192, h=6, layers=4):
        super().__init__()
        self.emb = nn.Embedding(vocab, d)
        self.pos = nn.Embedding(CTX, d)
        self.blocks = nn.ModuleList(Block(d, h) for _ in range(layers))
        self.lnf = nn.LayerNorm(d)
        self.head = nn.Linear(d, vocab, bias=False)
        mask = torch.triu(torch.full((CTX, CTX), float("-inf")), diagonal=1)
        self.register_buffer("mask", mask)

    def forward(self, idx):
        x = self.emb(idx) + self.pos.weight[: idx.shape[1]]
        for b in self.blocks:
            x = b(x, self.mask[: idx.shape[1], : idx.shape[1]])
        return self.head(self.lnf(x))


def get_data():
    os.makedirs(os.path.dirname(DATA_PATH), exist_ok=True)
    if not os.path.exists(DATA_PATH):
        urllib.request.urlretrieve(DATA_URL, DATA_PATH)
    text = open(DATA_PATH).read()
    chars = sorted(set(text))
    stoi = {c: i for i, c in enumerate(chars)}
    data = torch.tensor([stoi[c] for c in text], dtype=torch.long)
    return data, len(chars)


def make_batches(data, n_steps, seed=1234):
    """Fixed data order shared by every arm."""
    g = torch.Generator().manual_seed(seed)
    return [torch.randint(0, len(data) - CTX - 1, (BATCH,), generator=g)
            for _ in range(n_steps)]


def run_arm(name, data, vocab, batches, init_seed, make_opt, harvest=False):
    torch.manual_seed(init_seed)
    model = CharLM(vocab)
    opt = make_opt(model)
    losses, nmse_rows = [], []
    n_params = sum(p.numel() for p in model.parameters())
    for step, ix in enumerate(batches):
        xb = torch.stack([data[i:i + CTX] for i in ix])
        yb = torch.stack([data[i + 1:i + CTX + 1] for i in ix])
        logits = model(xb)
        loss = F.cross_entropy(logits.reshape(-1, vocab), yb.reshape(-1))
        opt.zero_grad(set_to_none=True)
        loss.backward()
        opt.step()
        losses.append(loss.item())
        if harvest and step % EVERY == 0 and step > 0:
            for p in model.parameters():
                st = opt.state.get(p, {})
                if "exp_avg" in st and p.numel() >= 4096:
                    m, v = st["exp_avg"], st["exp_avg_sq"]
                    nmse_rows.append(dict(
                        step=step, numel=p.numel(),
                        m_companded=Q.nmse(m, Q.momentum_companded(m)),
                        m_linear=Q.nmse(m, Q.momentum_linear(m)),
                        v_companded=Q.nmse(v, Q.variance_companded(v)),
                        v_linear=Q.nmse(v, Q.variance_linear(v))))
        if not math.isfinite(loss.item()):
            print(f"[{name}] NONFINITE loss at step {step}; stopping arm")
            losses.extend([float("nan")] * (len(batches) - len(losses)))
            break
    print(f"[{name}] params={n_params} final_loss={losses[-1]:.4f} "
          f"min_loss={min(x for x in losses if math.isfinite(x)):.4f}")
    return losses, nmse_rows


def main():
    os.makedirs(OUT, exist_ok=True)
    data, vocab = get_data()
    batches = make_batches(data, STEPS)
    arms = {}
    arms["ref_s0"], nmse = run_arm("ref_s0", data, vocab, batches, 0,
                                   lambda m: torch.optim.AdamW(m.parameters(), lr=LR),
                                   harvest=True)
    arms["ref_s1"], _ = run_arm("ref_s1", data, vocab, batches, 1,
                                lambda m: torch.optim.AdamW(m.parameters(), lr=LR))
    arms["flash_s0"], _ = run_arm("flash_s0", data, vocab, batches, 0,
                                  lambda m: FlashAdamWSim(m.parameters(), lr=LR,
                                                          companded=True, simulate_split_on_fp32=True))
    arms["linear_s0"], _ = run_arm("linear_s0", data, vocab, batches, 0,
                                   lambda m: FlashAdamWSim(m.parameters(), lr=LR,
                                                           companded=False, simulate_split_on_fp32=True))

    # trajectory deltas (claim 4, local): flash vs ref against seed yardstick
    import csv
    with open(f"{OUT}/charlm_losses.csv", "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["step"] + list(arms))
        for i in range(STEPS):
            w.writerow([i] + [arms[a][i] for a in arms])
    with open(f"{OUT}/charlm_nmse.csv", "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(nmse[0].keys()))
        w.writeheader(); w.writerows(nmse)

    tail = slice(STEPS // 2, STEPS)  # steady-state comparison window

    def dbar(a, b):
        d = [abs(x - y) for x, y in zip(arms[a][tail.start:tail.stop], arms[b][tail.start:tail.stop])
             if math.isfinite(x) and math.isfinite(y)]
        return sum(d) / len(d) if d else float("nan")
    med = lambda k: sorted(r[k] for r in nmse)[len(nmse) // 2]
    summary = {
        "final_loss": {k: arms[k][-1] for k in arms},
        "mean_abs_delta_2nd_half": {
            "flash_vs_ref (method effect)": dbar("flash_s0", "ref_s0"),
            "seed_vs_ref (seed noise yardstick)": dbar("ref_s1", "ref_s0"),
        },
        "linear_diverged": (not math.isfinite(arms["linear_s0"][-1])) or
                            arms["linear_s0"][-1] > 2 * arms["ref_s0"][-1],
        "nmse_median": {k: med(k) for k in ["m_companded", "m_linear", "v_companded", "v_linear"]},
        "nmse_improvement_x": {"momentum": med("m_linear") / med("m_companded"),
                                "variance": med("v_linear") / med("v_companded")},
        "config": dict(ctx=CTX, batch=BATCH, steps=STEPS, lr=LR, vocab=vocab),
    }
    with open(f"{OUT}/charlm_summary.json", "w") as f:
        json.dump(summary, f, indent=2)
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()