code/modal_trainmix_salvo.py

"""SN3 cycle-63 salvo: FIRST full-objective-coverage CPT, on a 2x4 (batch x peak_lr) grid.

WHY (read /var/lib/mining/sn-packages/sn3/learning/ first):

  Cycle 62 found the root cause of ~15 failed cycles: the trainable `.npy` shards on
  `mining-datasets` covered only **9% of the eval objective weight** (dendrite 0.04 +
  cosmopedia 0.05). `trainmix-v1` fixes that (10/10 corpora, exact validator weights,
  eval-disjointness verified by byte-range intersection, 871 contaminated rows excised).

  But the cycle-62 ladder log carries a second, sharper signal that the coverage story
  alone does not explain: at `peak_lr=1e-5, accum=16` (B = 32 seqs = 65k tok), after only
  8 updates **every corpus got worse, including `dendrite-synth-run` itself** (mu
  -0.002923) — the very corpus being trained on. Training a corpus cannot hurt that
  corpus if the step is signal-dominated. And damage scaled as lr^2
  (u=8: lr1e-5 -0.00132, lr4e-5 -0.01867 ~ 14x for a 4x lr).

  That is the signature of the noise term of SGD around a minimum:

      E[dLoss] ~ -lr * |g|^2  +  (lr^2 / 2) * tr(H Sigma) / B

  i.e. the king sits where the *gradient-noise* term dominates at B=32. Fitting the
  observed damage gives c/B|_{B=32} ~ 1.65e6 per update per lr^2, so the improvement
  window is lr < 2*B*|g|^2 / tr(H Sigma) — it widens **linearly in B**. Lowering lr alone
  shrinks the signal too (it is only linear in lr) and drives mu -> 0; raising B is the
  only knob that moves the optimum *and* raises its height.

  So the single most informative salvo is not a 1-D lr sweep: it is the 2x4 grid
  {accum 16, 64} x {1e-6, 3e-6, 1e-5, 3e-5}. Both marginals are readable, the predicted
  optima for both batches (~5e-6 and ~2e-5 at s~17) are bracketed, and the cell
  (accum 16, lr 1e-5) is a **direct control against cycle-62's own measurement** —
  identical recipe, only the corpus coverage changed (0.09 -> 1.00).

CONTRACTS
  * base = REIGNING king (reign 179, hf:ce6efba4...). Never a dethroned king.
  * train data = /data/trainmix-v1 ONLY. Never the eval set (/data/truemix-v1).
  * u=0 wiring check vs an INDEPENDENT cached per-sequence reference; fail closed.
  * gate-facing `mu_hat` is the FINAL eval point, not the argmax over the trajectory
    (no peeking optimism); the full trajectory is kept in detail.json for analysis.
  * writes lab-bus result.json so `python -m lab.collect|select` applies the real gate.
"""
from __future__ import annotations

import json
import os
import time

import modal

APP = "mining-trainmix-salvo"
SEQ_LEN = 2048
VOCAB_SIZE = 248320

TRAIN_MIX = "/data/trainmix-v1"          # full 10-corpus objective coverage
EVAL_MIX = "/data/truemix-v1"            # eval sample (NEVER trained on)
KING = "/kings/king-647-cp0"             # reign 179
KING_DIGEST = "hf:ce6efba444047328046a511e21d3e37a84c329c5"
REF_PATH = "/data/truemix-results/rebase-newking-n300/per_sequence.json"
LAB_BUS = "/data/lab-results"

image = (
    modal.Image.from_registry("nvidia/cuda:12.8.1-devel-ubuntu24.04", add_python="3.12")
    .env(
        {
            "MAX_JOBS": "8",
            "TORCH_CUDA_ARCH_LIST": "9.0",
            "CAUSAL_CONV1D_FORCE_BUILD": "TRUE",
            "HF_HUB_ENABLE_HF_TRANSFER": "1",
            "CUDA_HOME": "/usr/local/cuda",
        }
    )
    .apt_install("git", "build-essential")
    .pip_install("torch==2.8.0", index_url="https://download.pytorch.org/whl/cu128")
    .pip_install(
        "transformers==5.8.0", "safetensors", "accelerate", "einops", "numpy",
        "huggingface_hub", "hf_transfer", "ninja", "packaging", "setuptools", "wheel",
    )
    .pip_install("flash-linear-attention")
    .pip_install(
        "https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.6.2.post1/causal_conv1d-1.6.2.post1%2Bcu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl"
    )
)

app = modal.App(APP)
kings = modal.Volume.from_name("mining-kings")
ckpts = modal.Volume.from_name("mining-checkpoints")
datasets = modal.Volume.from_name("mining-datasets")
VOLS = {"/kings": kings, "/ckpts": ckpts, "/data": datasets}


def _log(m):
    print("[%s] %s" % (time.strftime("%H:%M:%S"), m), flush=True)


@app.function(image=image, gpu="B200:1", volumes=VOLS, timeout=60 * 180)
def variant(
    tag: str,
    variant_id: str,
    peak_lr: float,
    accum: int,
    max_updates: int = 40,
    micro_bs: int = 2,
    warmup: int = 4,
    min_lr_frac: float = 1.0,
    train_top_k: int = 8,
    train_lm_head: bool = True,
    eval_every: int = 10,
    eval_n: int = 300,
    seed: int = 777,
) -> dict:
    import math

    import numpy as np
    import torch
    from transformers import AutoModelForCausalLM

    t0 = time.time()
    outdir = f"{LAB_BUS}/{tag}/variants/{variant_id}"
    os.makedirs(outdir, exist_ok=True)

    # ---------- eval mix (truemix-v1) + INDEPENDENT cached king reference ----------
    eidx = json.load(open(f"{EVAL_MIX}/index.json"))
    corpora = [(c["corpus"], float(c["weight"]), c["file"]) for c in eidx["corpora"]]
    eval_tok = {}
    for name, _w, fn in corpora:
        arr = np.load(f"{EVAL_MIX}/{fn}", mmap_mode="r")
        if arr.shape[0] < eval_n:
            raise RuntimeError("eval corpus %s has %d < %d sequences" % (name, arr.shape[0], eval_n))
        eval_tok[name] = np.ascontiguousarray(arr[:eval_n])

    blob = json.load(open(REF_PATH))
    if KING not in blob:
        raise RuntimeError("reference %s has no key %s (keys=%s)" % (REF_PATH, KING, list(blob)))
    ref = blob[KING]
    king_ref = {}
    for n_, _w, _f in corpora:
        v = np.array(ref[n_][:eval_n], dtype=np.float64)
        if v.shape[0] != eval_tok[n_].shape[0]:
            raise RuntimeError("ref/eval length mismatch %s: %d vs %d"
                               % (n_, v.shape[0], eval_tok[n_].shape[0]))
        king_ref[n_] = v
    _log("king ref %s key=%s (n=%d/corpus)" % (REF_PATH, KING, eval_n))

    # ---------- FULL-COVERAGE training pool (trainmix-v1) ----------
    tidx = json.load(open(f"{TRAIN_MIX}/index.json"))
    if abs(float(tidx["weight_total"]) - 1.0) > 1e-9:
        raise RuntimeError("trainmix weight_total != 1.0: %r" % tidx["weight_total"])
    if "eval_disjointness" not in tidx:
        raise RuntimeError("trainmix index has no eval_disjointness block - refusing to train")
    pool = []
    for c in tidx["corpora"]:
        p = f"{TRAIN_MIX}/{c['file']}"
        a = np.load(p, mmap_mode="r")
        if a.ndim != 2 or a.shape[1] != SEQ_LEN:
            raise RuntimeError("bad train shard shape %s %s" % (p, a.shape))
        pool.append((c["corpus"], float(c["weight"]), a))
        _log("  train %-45s w=%.2f seqs=%d" % (c["corpus"], float(c["weight"]), a.shape[0]))
    train_names = {n_ for n_, _w, _a in pool}
    eval_names = {n_ for n_, _w, _f in corpora}
    if train_names != eval_names:
        raise RuntimeError("train/eval corpus set mismatch: %s" % (train_names ^ eval_names))
    w = np.array([x[1] for x in pool], dtype=np.float64)
    w = w / w.sum()
    cum = np.cumsum(w)
    n_train_seqs = int(sum(int(a.shape[0]) for _n, _w2, a in pool))
    _log("train pool: %d corpora, %d sequences, coverage=%.2f" % (len(pool), n_train_seqs, w.sum()))

    # up-front OOB audit: prove we are not silently training on zeroed sequences
    rs_audit = np.random.default_rng(12345)
    bad = 0
    for _ in range(2000):
        ci = int(np.searchsorted(cum, rs_audit.random()))
        arr = pool[min(ci, len(pool) - 1)][2]
        si = int(rs_audit.integers(0, arr.shape[0]))
        seq = np.asarray(arr[si])
        if seq.shape[0] != SEQ_LEN or (seq >= VOCAB_SIZE).any():
            bad += 1
    _log("OOB audit: %d/2000 sampled trainmix sequences unusable" % bad)
    if bad > 20:
        raise RuntimeError("trainmix OOB rate %d/2000 too high - fail closed" % bad)

    oob = {"n": 0, "drawn": 0}
    drawn_by_corpus = {n_: 0 for n_, _w2, _a in pool}

    def make_draw(rs):
        def draw(k):
            out = np.zeros((k, SEQ_LEN), dtype=np.int64)
            for j in range(k):
                ci = min(int(np.searchsorted(cum, rs.random())), len(pool) - 1)
                cname, _w2, arr = pool[ci]
                si = int(rs.integers(0, arr.shape[0]))
                seq = np.asarray(arr[si], dtype=np.int64)
                oob["drawn"] += 1
                drawn_by_corpus[cname] += 1
                if seq.shape[0] != SEQ_LEN or (seq >= VOCAB_SIZE).any():
                    oob["n"] += 1
                    seq = np.zeros(SEQ_LEN, dtype=np.int64)
                out[j] = seq
            return out
        return draw

    # ---------- model: pristine REIGNING king ----------
    import sys as _sys

    if KING not in _sys.path:
        _sys.path.insert(0, KING)
    _log("loading king %s" % KING)
    model = AutoModelForCausalLM.from_pretrained(
        KING, trust_remote_code=True, dtype=torch.bfloat16, low_cpu_mem_usage=True
    ).cuda()
    model.config.use_cache = False
    try:
        model.gradient_checkpointing_enable()
        _log("gradient checkpointing on")
    except Exception as e:  # noqa: BLE001
        _log("no gradient checkpointing: %s" % e)

    nl = model.config.num_hidden_layers
    keep_from = nl - train_top_k
    trainable = []
    for n_, p in model.named_parameters():
        want = False
        if ".layers." in n_:
            try:
                li = int(n_.split(".layers.")[1].split(".")[0])
            except Exception:  # noqa: BLE001
                li = -1
            want = li >= keep_from
        elif "lm_head" in n_:
            want = train_lm_head
        elif n_.endswith("model.norm.weight") or n_.endswith("final_layernorm.weight"):
            want = True
        p.requires_grad_(want)
        if want:
            trainable.append((n_, p))
    ntp = sum(p.numel() for _n, p in trainable)
    _log("trainable %d tensors / %.3fB params (layers >= %d of %d, lm_head=%s)"
         % (len(trainable), ntp / 1e9, keep_from, nl, train_lm_head))

    def chunked_ce(logits, tgt, step=256):
        b, tm1, _ = logits.shape
        tot = None
        for s in range(0, tm1, step):
            sl = logits[:, s : s + step, :].float()
            st = tgt[:, s : s + step]
            ce = torch.nn.functional.cross_entropy(
                sl.reshape(-1, sl.shape[-1]), st.reshape(-1), reduction="sum"
            )
            tot = ce if tot is None else tot + ce
        return tot / (b * tm1)

    @torch.no_grad()
    def eval_objective():
        model.eval()
        per = {}
        for name, _w2, _fn in corpora:
            toks = eval_tok[name]
            out = np.zeros(toks.shape[0], dtype=np.float64)
            for i in range(0, toks.shape[0], 2):
                ids = torch.from_numpy(toks[i : i + 2].astype(np.int64)).cuda()
                lg = model(input_ids=ids, use_cache=False).logits[:, :-1, :]
                tg = ids[:, 1:]
                b, tm1, _ = lg.shape
                acc = torch.zeros(b, dtype=torch.float64, device=lg.device)
                for s in range(0, tm1, 256):
                    sl = lg[:, s : s + 256, :].float()
                    st = tg[:, s : s + 256]
                    ce = torch.nn.functional.cross_entropy(
                        sl.reshape(-1, sl.shape[-1]), st.reshape(-1), reduction="none"
                    ).view(b, -1)
                    acc += ce.sum(dim=1).double()
                out[i : i + ids.shape[0]] = (acc / tm1).cpu().numpy()
                del lg
            per[name] = out
        model.train()
        mu, contrib, var = {}, {}, 0.0
        for name, wt, _fn in corpora:
            d = king_ref[name] - per[name]
            mu[name] = float(d.mean())
            contrib[name] = float(wt * mu[name])
            var += (wt ** 2) * float(d.var(ddof=1)) / len(d)
        se = math.sqrt(var)
        m = float(sum(contrib.values()))
        return {
            "mu_hat": m, "se": se, "lcb999": m - 3.09 * se,
            "mu_per_corpus": mu, "weighted_contribution": contrib,
            "cand_loss": {k: float(v.mean()) for k, v in per.items()},
        }

    recipe = {
        "peak_lr": peak_lr, "accum": accum, "micro_bs": micro_bs,
        "effective_batch_seqs": micro_bs * accum,
        "effective_batch_tokens": micro_bs * accum * SEQ_LEN,
        "max_updates": max_updates, "warmup": warmup, "min_lr_frac": min_lr_frac,
        "train_top_k": train_top_k, "train_lm_head": train_lm_head,
        "seed": seed, "data_order_seed": seed, "eval_n": eval_n,
        "train_data": TRAIN_MIX, "objective_coverage": 1.0,
        "king": KING, "king_digest": KING_DIGEST, "precision": "bf16",
    }

    detail = {
        "tag": tag, "variant_id": variant_id, "recipe": recipe,
        "train_index": {k: tidx.get(k) for k in
                        ("name", "seed", "weight_total", "tokenizer", "eval_disjointness")},
        "n_train_sequences": n_train_seqs,
        "eval_ref": REF_PATH, "eval_mix": EVAL_MIX,
    }

    def persist(status, gate_eval, base, hist, note=""):
        used_real = bool(n_train_seqs > 0 and oob["n"] == 0 and bad <= 20)
        promotable = bool(used_real and status == "ok")
        bus = {
            "variant_id": variant_id,
            "tag": tag,
            "mu_hat": gate_eval["mu_hat"] if gate_eval else 0.0,
            "se": gate_eval["se"] if gate_eval else 1.0,
            "lcb999": (gate_eval["mu_hat"] - 3.09 * gate_eval["se"]) if gate_eval else -1.0,
            "king_digest": KING_DIGEST,
            "used_real_data": used_real,
            "real_data": used_real,
            "promotable": promotable,
            "oob_count": oob["n"],
            "oob_drawn": oob["drawn"],
            "effective_accumulation": accum,
            "effective_peak_lr": peak_lr,
            "recipe": recipe,
            "baseline_u0": base,
            "status": status,
            "gate_point": "final_eval_point (no argmax peeking)",
            "note": note,
        }
        with open(f"{outdir}/result.json", "w") as fh:
            json.dump(bus, fh, indent=2)
        with open(f"{outdir}/detail.json", "w") as fh:
            json.dump({**detail, "baseline_u0": base, "history": hist,
                       "oob": dict(oob), "drawn_by_corpus": dict(drawn_by_corpus),
                       "status": status, "elapsed_s": round(time.time() - t0, 1)}, fh, indent=2)
        datasets.commit()

    # ---------- u=0 WIRING CHECK (fail closed) ----------
    _log("=== u=0 baseline (must reproduce the pinned king vs independent reference) ===")
    base = eval_objective()
    _log("u=0 mu_hat=%+.6f se=%.6f" % (base["mu_hat"], base["se"]))
    tol = max(5e-5, 4.0 * base["se"])
    if abs(base["mu_hat"]) > tol:
        persist("wiring_failed", None, base, [], note="u=0 mu_hat %.6g exceeds tol %.6g" % (base["mu_hat"], tol))
        raise RuntimeError("baseline mu_hat=%.6f exceeds tol %.6f -> wrong king/reference pairing"
                           % (base["mu_hat"], tol))
    _log("baseline wiring OK (tol %.6g)" % tol)

    hist = [{"update": 0, "tokens": 0, **base}]
    persist("running", base, base, hist, note="u=0 only")

    # ---------- train ----------
    masters = [p.detach().float().clone() for _n, p in trainable]
    for m in masters:
        m.requires_grad_(False)
    opt = torch.optim.AdamW(masters, lr=peak_lr, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.0)
    rs = np.random.default_rng(seed)
    draw = make_draw(rs)

    def lr_at(u):
        if u < warmup:
            return peak_lr * (u + 1) / warmup
        prog = (u - warmup) / max(1, max_updates - warmup)
        return peak_lr * (min_lr_frac + (1 - min_lr_frac) * 0.5 * (1 + math.cos(math.pi * min(1.0, prog))))

    model.train()
    last = base
    for u in range(max_updates):
        lr = lr_at(u)
        for g in opt.param_groups:
            g["lr"] = lr
        opt.zero_grad(set_to_none=True)
        for _n, p in trainable:
            p.grad = None
        tl, tu = 0.0, time.time()
        for _a in range(accum):
            ids = torch.from_numpy(draw(micro_bs)).cuda()
            logits = model(input_ids=ids, use_cache=False).logits[:, :-1, :]
            loss = chunked_ce(logits, ids[:, 1:]) / accum
            loss.backward()
            tl += float(loss.detach()) * accum
            del logits, loss
        gn = 0.0
        for (nm, p), m in zip(trainable, masters):
            if p.grad is None:
                continue
            m.grad = p.grad.float()
            gn += float(m.grad.pow(2).sum())
        gn = gn ** 0.5
        torch.nn.utils.clip_grad_norm_(masters, 1.0)
        opt.step()
        with torch.no_grad():
            for (nm, p), m in zip(trainable, masters):
                p.data.copy_(m.to(p.dtype))
        for m in masters:
            m.grad = None
        if u % 5 == 0 or u == max_updates - 1:
            _log("  u=%d loss=%.5f lr=%.3e gnorm=%.3f %.1fs/upd mem=%.1fGB"
                 % (u + 1, tl / accum, lr, gn, time.time() - tu,
                    torch.cuda.max_memory_allocated() / 1e9))
        if (u + 1) % eval_every == 0 or (u + 1) == max_updates:
            e = eval_objective()
            last = e
            _log("  *** %s u=%d mu_hat=%+.6f se=%.6f lcb999=%+.6f" %
                 (variant_id, u + 1, e["mu_hat"], e["se"], e["mu_hat"] - 3.09 * e["se"]))
            hist.append({"update": u + 1, "tokens": (u + 1) * micro_bs * accum * SEQ_LEN, **e})
            persist("running", e, base, hist, note="u=%d" % (u + 1))

    # gate-facing = FINAL eval point (no argmax peeking)
    persist("ok", last, base, hist, note="final eval point at u=%d" % max_updates)
    best = max(hist, key=lambda h: h["mu_hat"])
    summary = {
        "tag": tag, "variant_id": variant_id, "peak_lr": peak_lr, "accum": accum,
        "u0_mu_hat": base["mu_hat"],
        "final_mu_hat": last["mu_hat"], "final_se": last["se"],
        "final_lcb999": last["mu_hat"] - 3.09 * last["se"],
        "best_mu_hat": best["mu_hat"], "best_update": best["update"],
        "oob": dict(oob), "drawn_by_corpus": dict(drawn_by_corpus),
        "elapsed_s": round(time.time() - t0, 1),
    }
    _log("done " + json.dumps(summary))
    return summary


@app.local_entrypoint()
def main(
    tag: str = "tmix63",
    variant_id: str = "b16-lr1e-06",
    peak_lr: float = 1e-6,
    accum: int = 16,
    max_updates: int = 40,
    eval_every: int = 10,
    eval_n: int = 300,
    micro_bs: int = 2,
    warmup: int = 4,
    train_top_k: int = 8,
    seed: int = 777,
):
    r = variant.remote(
        tag=tag, variant_id=variant_id, peak_lr=peak_lr, accum=accum,
        max_updates=max_updates, micro_bs=micro_bs, warmup=warmup,
        train_top_k=train_top_k, eval_every=eval_every, eval_n=eval_n, seed=seed,
    )
    print(json.dumps(r, indent=2)[:6000])