scripts/ablate_variable_rate.py

# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "torch",
#   "numpy",
#   "datasets>=2.19",
#   "huggingface_hub>=0.24",
#   "safetensors>=0.4",
# ]
# ///
"""KG-04 variable-rate proof: matched-cost quality-vs-compute frontier.

Runner for the protocol in trident.md Section 0.5.5.

Modes (env ``MODE``):
  * ``grid``       train the whole frontier in ONE job and push a scorecard
                   (recommended: one GPU allocation, atomic verdict).
  * ``train``      train a single point (env-driven) and push its result JSON.
  * ``aggregate``  collect previously pushed point JSONs into a scorecard.

Every point shares an identical architecture (``r_max`` and all dims fixed), so
parameter count is matched by construction; only the *commit policy* differs:

    fixed     R=k              -> exactly k write factors / patch
    adaptive  lambda_rate=lam  -> learned gates, cost pressure = lam

Cost axis = ``mean_factors`` (mean effective write factors = analytic write-FLOP
proxy; see Section 0.5.5). It is NOT wall-clock. Quality = held-out BPB
(TASK=text) or distance-stratified next-value accuracy (TASK=recall).

KG-04 passes only if the adaptive frontier lies strictly below the fixed-rate
frontier at matched mean cost with disjoint 95% CIs. The test can fail.

Nothing here touches Jobs or repos it did not create. Training runs on the
accelerator; the artefact is a CPU-only inference model after parity.
"""
import gc
import json
import math
import os
import sys
import time
from pathlib import Path

import torch


def env(k, d=None):
    v = os.environ.get(k)
    return v if v is not None and v != "" else d


def env_i(k, d):
    return int(env(k, d))


def env_f(k, d):
    return float(env(k, d))


def log(msg):
    print(f"[ablate] {msg}", flush=True)


def import_trident():
    try:
        import trident  # noqa: F401
        return
    except Exception:
        pass
    from huggingface_hub import snapshot_download
    code_repo = env("CODE_REPO", "farguney/trident")
    code_rev = env("CODE_REVISION", "main")
    workdir = Path("/tmp/trident_code")
    log(f"downloading {code_repo}@{code_rev} (src/**)")
    snapshot_download(
        repo_id=code_repo, revision=code_rev, repo_type="model",
        allow_patterns=["src/**"], local_dir=str(workdir),
    )
    sys.path.insert(0, str(workdir / "src"))


# ----------------------------------------------------------------------------
# synthetic distance-stratified recall control (MQAR-as-bytes)
# ----------------------------------------------------------------------------
# Byte ranges are disjoint so keys, values, filler, and the query marker never
# alias. The model must copy the value bound to a queried key from earlier in
# the stream *through the bounded recurrent state* -- the recurrent core has no
# attention over raw history. Distance = answer_pos - value_pos, so accuracy can
# be reported as a function of how far back the binding sits.
KEY_LO, KEY_HI = 1, 65          # keys in [1,65)
VAL_LO, VAL_HI = 65, 129        # values in [65,129)
FILL_LO, FILL_HI = 129, 254     # filler in [129,254)
ANSWER_MARK = 254               # "emit the recalled value next" cue
Q_MARK = 255                    # query marker byte
# Trident reads its associative state with one-patch latency: the decoder at a
# position conditions on the *previous* patch's committed state (causal design).
# So the query key and the scored answer must sit in DIFFERENT patches, else the
# query-conditioned read is never exposed and recall is at chance regardless of
# difficulty (verified empirically). Q_GAP filler bytes between the query key and
# the answer cue force >=1 patch boundary.
Q_GAP = 8


def make_recall_batch(B, seq_len, n_pairs, rng, device, pack="dense"):
    """Return bytes_in (B,L), answer_pos (B,), answer_val (B,), distance (B,).

    ``pack`` controls whether the R=1-vs-R>1 write bottleneck is exercised:

      * ``dense``  -- the ``n_pairs`` (key,val) bindings are laid *contiguously*
        in one block, so several bindings land inside a single entropy patch
        (~3-5 bytes). A patch commits one rank-1 delta write per factor, so with
        R=1 multiple co-patch bindings must share a single write and interfere;
        R>1 supplies extra independent writes. This is the controlled stimulus
        that can discriminate the commit rate. The block is placed so the query
        is a fixed distance away (retention held roughly constant across R).
      * ``sparse`` -- bindings scattered (mostly one per patch); R=1 usually
        suffices, so this is the negative control.
    """
    x = torch.from_numpy(rng.integers(FILL_LO, FILL_HI, size=(B, seq_len))).long()
    answer_pos = torch.zeros(B, dtype=torch.long)
    answer_val = torch.zeros(B, dtype=torch.long)
    distance = torch.zeros(B, dtype=torch.long)
    # query key at seq_len-Q_GAP, then filler (forces a patch boundary), then the
    # answer cue at seq_len-2; the value is scored as the byte predicted after it.
    q_mark_pos = seq_len - Q_GAP
    qk_pos = q_mark_pos + 1
    ans_pos = seq_len - 2
    for b in range(B):
        keys = rng.choice(range(KEY_LO, KEY_HI), size=n_pairs, replace=False)
        vals = rng.integers(VAL_LO, VAL_HI, size=n_pairs)
        if pack == "dense":
            base = 4 + int(rng.integers(0, max(1, seq_len // 16)))
            slots = [base + 2 * i for i in range(n_pairs)]
        else:
            region = max(4, int(seq_len * 2 // 3) - 2)
            slots = sorted(rng.choice(range(region), size=n_pairs, replace=False))
        val_by_key, val_pos = {}, {}
        for (k, v, s) in zip(keys, vals, slots):
            if s + 1 >= q_mark_pos - 1:  # keep bindings clear of the query block
                continue
            x[b, s] = int(k)
            x[b, s + 1] = int(v)
            val_pos[int(k)] = s + 1
            val_by_key[int(k)] = int(v)
        placed = list(val_by_key.keys())
        qk = int(rng.choice(placed))
        x[b, q_mark_pos] = Q_MARK
        x[b, qk_pos] = qk
        x[b, ans_pos] = ANSWER_MARK  # model predicts x[ans_pos+1]; we score that logit
        answer_pos[b] = ans_pos
        answer_val[b] = val_by_key[qk]
        distance[b] = ans_pos - val_pos[qk]
    return (x.to(device), answer_pos.to(device), answer_val.to(device), distance.to(device))


@torch.no_grad()
def eval_recall(model, device, seq_len, n_pairs, n_batches, B, amp_dtype, pack="dense"):
    import numpy as np
    rng = np.random.default_rng(1234)
    buckets = {}
    mf_sum, aR_sum, nb = 0.0, 0.0, 0
    model.eval()
    for _ in range(n_batches):
        x, apos, aval, dist = make_recall_batch(B, seq_len, n_pairs, rng, device, pack)
        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=device == "cuda"):
            out = model(x, return_logits=True)
        logits = out["logits"].float()
        pred = logits[torch.arange(x.shape[0]), apos].argmax(-1)
        correct = (pred == aval)
        mf_sum += out["mean_factors"].item(); aR_sum += out["active_R"].item(); nb += 1
        for c, d in zip(correct.tolist(), dist.tolist()):
            bkt = "<=64" if d <= 64 else ("65-256" if d <= 256 else ">256")
            slot = buckets.setdefault(bkt, [0, 0])
            slot[0] += int(c); slot[1] += 1
    model.train()
    acc = {k: (v[0] / v[1] if v[1] else 0.0) for k, v in buckets.items()}
    tot = sum(v[0] for v in buckets.values()) / max(1, sum(v[1] for v in buckets.values()))
    return {"recall_acc_overall": tot, "recall_acc_by_distance": acc,
            "mean_factors": mf_sum / max(1, nb), "active_R": aR_sum / max(1, nb)}


# ----------------------------------------------------------------------------
# config
# ----------------------------------------------------------------------------
def make_cfg(rate_mode, fixed_r, lambda_rate, task, extra_over=None):
    from trident import TridentConfig
    profile = env("PROFILE", "micro")
    over = {}
    for key, cast in [
        ("d_model", int), ("n_blocks", int), ("n_heads", int), ("d_k", int),
        ("d_v", int), ("n_dec_layers", int), ("r_max", int), ("b_max", int),
        ("tau_b", float), ("max_patches", int), ("exact_ring", int),
        ("d_code", int), ("chunk_size", int), ("rho_min", float),
        ("gate_bias_init", float), ("factor_emb_init", float),
        ("lambda_fast_head", float),
    ]:
        val = env(key.upper())
        if val is not None:
            over[key] = cast(val)
    ula = env("USE_LOCAL_ATTN")
    if ula is not None:
        over["use_local_attn"] = ula not in ("0", "false", "False", "")
    if rate_mode == "fixed":
        over["fixed_r"] = int(fixed_r)
        over["lambda_rate"] = 0.0
    elif rate_mode == "adaptive":
        over["fixed_r"] = None
        over["lambda_rate"] = float(lambda_rate)
    else:
        raise ValueError(f"rate_mode must be fixed|adaptive, got {rate_mode}")
    if task == "recall":
        over.setdefault("lambda_fast_head", 0.0)
    if extra_over:  # per-point overrides (e.g. operator pool config)
        over.update(extra_over)
    if profile in ("micro", "scout"):
        return getattr(TridentConfig, profile)(**over)
    over.setdefault("profile", profile)
    return TridentConfig(**over)


def read_hp():
    return dict(
        max_steps=env_i("MAX_STEPS", 3000), warmup=env_i("WARMUP", 200),
        lr=env_f("LR", 6e-4), min_lr=env_f("MIN_LR", 6e-5),
        wd=env_f("WEIGHT_DECAY", 0.1), grad_clip=env_f("GRAD_CLIP", 1.0),
        batch=env_i("BATCH", 16), grad_accum=env_i("GRAD_ACCUM", 1),
        seq_len=env_i("SEQ_LEN", 2048), n_pairs=env_i("N_PAIRS", 8),
        val_windows=env_i("VAL_WINDOWS", 64), val_batch=env_i("VAL_BATCH", 8),
        val_batch_n=env_i("VAL_BATCH_N", 16), log_every=env_i("LOG_EVERY", 200),
        eval_every=env_i("EVAL_EVERY", 0), pack=env("RECALL_PACK", "dense"),
    )


# ----------------------------------------------------------------------------
# train a single frontier point (pure: returns a result dict, no Hub I/O)
# ----------------------------------------------------------------------------
def train_one(task, rate_mode, fixed_r, lambda_rate, seed, hp, device, amp_dtype,
              extra_over=None, point_id=None):
    import numpy as np
    from trident import Trident

    torch.manual_seed(seed); np.random.seed(seed)
    cfg = make_cfg(rate_mode, fixed_r, lambda_rate, task, extra_over=extra_over)
    model = Trident(cfg).to(device)
    nparams = sum(p.numel() for p in model.parameters())
    # GLR-1 refiner: sample a refinement budget per step (Universal-Transformer
    # depth-by-iteration). None => no refiner (budget arg ignored by the model).
    refine_budgets = ([int(b) for b in cfg.refine_budgets]
                      if getattr(cfg, "use_refiner", False) else None)
    if point_id is None:
        point_id = f"fixed_R{fixed_r}" if rate_mode == "fixed" else f"adaptive_L{lambda_rate}"
    # active payload FLOPs (constant in pool size N) for the operator experiment
    active_flops = 0
    stored_op = 0
    if getattr(cfg, "use_operators", False) and model.decoder.operators is not None:
        active_flops = sum(op.active_payload_flops() for op in model.decoder.operators)
        stored_op = sum(op.stored_operator_params() for op in model.decoder.operators)
    log(f"== point={point_id} seed={seed} task={task} params={nparams/1e6:.3f}M "
        f"stored_op={stored_op/1e6:.3f}M active_op_flops={active_flops} "
        f"identity={cfg.identity_hash()[:12]}")

    decay = [p for _, p in model.named_parameters() if p.ndim >= 2]
    no_decay = [p for _, p in model.named_parameters() if p.ndim < 2]
    opt = torch.optim.AdamW(
        [{"params": decay, "weight_decay": hp["wd"]}, {"params": no_decay, "weight_decay": 0.0}],
        lr=hp["lr"], betas=(0.9, 0.95), eps=1e-8,
    )

    def lr_at(s):
        if s < hp["warmup"]:
            return hp["lr"] * (s + 1) / hp["warmup"]
        t = min(1.0, (s - hp["warmup"]) / max(1, hp["max_steps"] - hp["warmup"]))
        return hp["min_lr"] + 0.5 * (hp["lr"] - hp["min_lr"]) * (1 + math.cos(math.pi * t))

    if task == "text":
        from trident.data import ByteWindowIterable, collate, load_fixed_byte_windows
        train_ds = ByteWindowIterable(
            dataset=env("DATASET", "HuggingFaceFW/fineweb-edu"), split=env("SPLIT", "train"),
            seq_len=hp["seq_len"], text_field=env("TEXT_FIELD", "text"),
            name=env("DATASET_NAME", "sample-10BT"),
            shuffle_buffer=env_i("SHUFFLE_BUFFER", 10000), seed=seed,
        )
        loader = torch.utils.data.DataLoader(
            train_ds, batch_size=hp["batch"], collate_fn=collate,
            num_workers=env_i("NUM_WORKERS", 2), drop_last=True,
        )
        data_iter = iter(loader)
        val_windows = load_fixed_byte_windows(
            dataset=env("VAL_DATASET", "Salesforce/wikitext"), split=env("VAL_SPLIT", "validation"),
            name=env("VAL_NAME", "wikitext-103-raw-v1"), text_field=env("VAL_TEXT_FIELD", "text"),
            seq_len=hp["seq_len"], num_windows=hp["val_windows"], add_eod=False,
        ).to(device)

        @torch.no_grad()
        def evaluate(budget=None):
            model.eval()
            tot_nll, tot_bytes, mf, aR, nb = 0.0, 0, 0.0, 0.0, 0
            for i in range(0, val_windows.shape[0], hp["val_batch"]):
                chunk = val_windows[i:i + hp["val_batch"]]
                with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=device == "cuda"):
                    out = model(chunk, return_logits=True, refine_budget=budget)
                lg = out["logits"].float()
                tgt = chunk.clamp(0, lg.shape[-1] - 1)
                nll = torch.nn.functional.cross_entropy(
                    lg.reshape(-1, lg.shape[-1]), tgt.reshape(-1), reduction="sum")
                tot_nll += nll.item(); tot_bytes += tgt.numel()
                mf += out["mean_factors"].item(); aR += out["active_R"].item(); nb += 1
            model.train()
            return {"val_bpb": tot_nll / tot_bytes / math.log(2),
                    "mean_factors": mf / nb, "active_R": aR / nb}
    else:
        rng = np.random.default_rng(seed)

        def evaluate():
            return eval_recall(model, device, hp["seq_len"], hp["n_pairs"],
                               hp["val_batch_n"], hp["val_batch"], amp_dtype, hp["pack"])

    model.train()
    skipped, t0, last_mf = 0, time.time(), 0.0
    for step in range(hp["max_steps"]):
        for g in opt.param_groups:
            g["lr"] = lr_at(step)
        opt.zero_grad(set_to_none=True)
        loss_val = 0.0
        for _ in range(hp["grad_accum"]):
            rb = int(np.random.choice(refine_budgets)) if refine_budgets else None
            if task == "text":
                try:
                    bd = next(data_iter)
                except StopIteration:
                    data_iter = iter(loader); bd = next(data_iter)
                bytes_in = bd["bytes_in"].to(device, non_blocking=True)
                valid = bd["valid"].to(device, non_blocking=True)
                with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=device == "cuda"):
                    out = model(bytes_in, valid=valid, refine_budget=rb)
                    loss = out["loss"] / hp["grad_accum"]
            else:
                x, apos, aval, _ = make_recall_batch(hp["batch"], hp["seq_len"], hp["n_pairs"], rng, device, hp["pack"])
                with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=device == "cuda"):
                    out = model(x, return_logits=True, refine_budget=rb)
                    lg = out["logits"].float()
                    ans_logits = lg[torch.arange(x.shape[0]), apos]
                    ce = torch.nn.functional.cross_entropy(ans_logits, aval)
                    loss = (ce + cfg.lambda_rate * out["mean_factors"]) / hp["grad_accum"]
            loss.backward()
            loss_val += loss.item()
        last_mf = out["mean_factors"].item()
        gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), hp["grad_clip"])
        if torch.isfinite(gnorm):
            opt.step()
        else:
            opt.zero_grad(set_to_none=True); skipped += 1
        if step % hp["log_every"] == 0:
            log(f"   step {step}/{hp['max_steps']} loss {loss_val:.4f} mean_factors {last_mf:.3f} lr {lr_at(step):.2e}")
        if hp["eval_every"] and (step + 1) % hp["eval_every"] == 0:
            ev = evaluate()
            qk = "val_bpb" if task == "text" else "recall_acc_overall"
            log(f"   [eval] {point_id} step {step + 1} {qk}={ev[qk]:.4f} mf={ev['mean_factors']:.3f}")

    metrics = evaluate()
    # GLR-1 depth test: test-time budget sweep on the SAME trained refiner model
    # (perfectly param-matched). If BPB falls as budget grows, iteration buys
    # quality -- the core "test-time compute buys depth" claim, on real text.
    val_bpb_by_budget = None
    if refine_budgets and task == "text":
        val_bpb_by_budget = {str(b): round(evaluate(b)["val_bpb"], 5)
                             for b in sorted(set(refine_budgets))}
        log(f"   BUDGET_SWEEP {point_id} seed{seed} :: {json.dumps(val_bpb_by_budget)}")
    result = {
        "point_id": point_id, "task": task, "rate_mode": rate_mode,
        "fixed_r": cfg.fixed_r, "lambda_rate": cfg.lambda_rate, "seed": seed,
        "params": nparams, "identity": cfg.identity_hash(),
        "use_refiner": bool(getattr(cfg, "use_refiner", False)),
        "refine_budgets": list(getattr(cfg, "refine_budgets", ())),
        "val_bpb_by_budget": val_bpb_by_budget,
        "use_operators": bool(getattr(cfg, "use_operators", False)),
        "n_operators": int(getattr(cfg, "n_operators", 0)),
        "n_active": int(getattr(cfg, "n_active", 0)),
        "operator_mode": getattr(cfg, "operator_mode", "learned"),
        "active_op_flops": active_flops, "stored_op_params": stored_op,
        "max_steps": hp["max_steps"], "batch": hp["batch"], "grad_accum": hp["grad_accum"],
        "seq_len": hp["seq_len"], "tokens": hp["max_steps"] * hp["batch"] * hp["grad_accum"] * hp["seq_len"],
        "skipped": skipped, "wall_s": round(time.time() - t0, 1), **metrics,
    }
    log(f"   RESULT {json.dumps(result)}")
    del model, opt
    gc.collect()
    if device == "cuda":
        torch.cuda.empty_cache()
    return result


# ----------------------------------------------------------------------------
# aggregate frontier points -> verdict (pure)
# ----------------------------------------------------------------------------
_T95 = {1: 12.706, 2: 4.303, 3: 3.182, 4: 2.776, 5: 2.571, 6: 2.447,
        7: 2.365, 8: 2.306, 9: 2.262, 10: 2.228}


def _mean_ci(xs):
    n = len(xs)
    m = sum(xs) / n
    if n < 2:
        return m, float("nan")
    var = sum((x - m) ** 2 for x in xs) / (n - 1)
    t = _T95.get(n - 1, 1.96)
    return m, t * math.sqrt(var) / math.sqrt(n)


def build_scorecard(pts, task):
    qkey = "val_bpb" if task == "text" else "recall_acc_overall"
    pts = [p for p in pts if p.get("task") == task]
    groups = {}
    for p in pts:
        groups.setdefault(p["point_id"], []).append(p)
    rows = []
    for pid, ps in groups.items():
        q_m, q_ci = _mean_ci([p[qkey] for p in ps])
        c_m, c_ci = _mean_ci([p["mean_factors"] for p in ps])
        rows.append({"point_id": pid, "rate_mode": ps[0]["rate_mode"],
                     "cost_mean_factors": c_m, "cost_ci": c_ci,
                     "quality": q_m, "quality_ci": q_ci, "n_seeds": len(ps),
                     "active_R": sum(p["active_R"] for p in ps) / len(ps)})
    rows.sort(key=lambda r: r["cost_mean_factors"])
    fixed = [r for r in rows if r["rate_mode"] == "fixed"]
    adaptive = [r for r in rows if r["rate_mode"] == "adaptive"]

    def fixed_quality_at(cost):
        if not fixed:
            return None
        below = [r for r in fixed if r["cost_mean_factors"] <= cost]
        above = [r for r in fixed if r["cost_mean_factors"] >= cost]
        if not below:
            return fixed[0]["quality"], fixed[0]["quality_ci"]
        if not above:
            return fixed[-1]["quality"], fixed[-1]["quality_ci"]
        lo = max(below, key=lambda r: r["cost_mean_factors"])
        hi = min(above, key=lambda r: r["cost_mean_factors"])
        if hi["cost_mean_factors"] == lo["cost_mean_factors"]:
            return lo["quality"], lo["quality_ci"]
        w = (cost - lo["cost_mean_factors"]) / (hi["cost_mean_factors"] - lo["cost_mean_factors"])
        q = lo["quality"] + w * (hi["quality"] - lo["quality"])
        return q, max(lo["quality_ci"], hi["quality_ci"])

    wins = []
    for a in adaptive:
        fq = fixed_quality_at(a["cost_mean_factors"])
        if fq is None:
            continue
        fqm, fqci = fq
        if task == "text":  # lower bpb better
            strictly = (a["quality"] + (a["quality_ci"] or 0)) < (fqm - (fqci or 0))
        else:               # higher recall better
            strictly = (a["quality"] - (a["quality_ci"] or 0)) > (fqm + (fqci or 0))
        wins.append({"adaptive_point": a["point_id"], "cost": a["cost_mean_factors"],
                     "adaptive_quality": a["quality"], "fixed_quality_at_cost": fqm,
                     "strict_win": bool(strictly)})

    passed = any(w["strict_win"] for w in wins)
    return {
        "task": task, "quality_metric": qkey, "n_points": len(rows),
        "fixed_frontier": fixed, "adaptive_frontier": adaptive,
        "matched_cost_comparisons": wins, "KG04_pass": passed,
        "verdict": ("PASS: adaptive beats fixed at matched cost (disjoint CIs)" if passed
                    else "NOT PROVEN: adaptive frontier does not strictly beat fixed within CIs"),
    }


# ----------------------------------------------------------------------------
# Hub push helpers
# ----------------------------------------------------------------------------
def _push(local_path, path_in_repo, msg):
    from huggingface_hub import HfApi
    try:
        HfApi().upload_file(
            path_or_fileobj=str(local_path), path_in_repo=path_in_repo,
            repo_id=env("CODE_REPO", "farguney/trident"), repo_type="model",
            commit_message=msg,
        )
        log(f"pushed {path_in_repo}")
    except Exception as e:  # noqa
        log(f"push failed for {path_in_repo} ({e})")


# ----------------------------------------------------------------------------
# entry points
# ----------------------------------------------------------------------------
def device_amp():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cuda":
        log(f"gpu {torch.cuda.get_device_name(0)}  torch {torch.__version__}")
    else:
        log("WARNING: no CUDA device; running on CPU")
    return device, (torch.bfloat16 if device == "cuda" else torch.float32)


def run_grid():
    device, amp = device_amp()
    hp = read_hp()
    task = env("TASK", "text")
    fixed_rs = [int(x) for x in env("FIXED_RS", "1,2,4").split(",") if x != ""]
    lam_env = os.environ.get("LAMBDAS", "0.02,0.1,0.5")
    lambdas = ([] if lam_env.strip().lower() in ("", "none", "skip")
               else [float(x) for x in lam_env.split(",") if x != ""])
    seeds = [int(x) for x in env("SEEDS", "0,1,2").split(",") if x != ""]
    run = env("RUN", "vr")
    results_dir = env("RESULTS_DIR", "ablation/variable_rate")
    log(f"GRID task={task} fixed_rs={fixed_rs} lambdas={lambdas} seeds={seeds} "
        f"points={(len(fixed_rs)+len(lambdas))*len(seeds)} steps/pt={hp['max_steps']}")

    outdir = Path("/tmp/vr_grid"); outdir.mkdir(exist_ok=True)
    results = []

    def flush(tag):
        # push after every point so a timeout/preemption never loses completed
        # (expensive) points; MODE=aggregate can rebuild the scorecard any time.
        (outdir / "results.jsonl").write_text("\n".join(json.dumps(r) for r in results))
        _push(outdir / "results.jsonl", f"{results_dir}/{run}/results_{task}.jsonl",
              f"vr grid results ({task}, {len(results)} pts, {tag})")

    plan = [("fixed", r, 0.0) for r in fixed_rs] + [("adaptive", None, lam) for lam in lambdas]
    for (mode, r, lam) in plan:
        for s in seeds:
            results.append(train_one(task, mode, r, lam, s, hp, device, amp))
            flush(f"{mode}:{r if mode == 'fixed' else lam}:s{s}")

    score = build_scorecard(results, task)
    (outdir / "scorecard.json").write_text(json.dumps(score, indent=2))
    _push(outdir / "scorecard.json", f"{results_dir}/{run}/scorecard_{task}.json",
          f"vr grid scorecard ({task})")
    log(f"VERDICT [{task}] KG04_pass={score['KG04_pass']} :: {score['verdict']}")
    log(f"SCORECARD {json.dumps(score, indent=2)}")


def build_operator_scorecard(pts):
    """KG-05 (operator value) + KG-08 (stored-capacity slope) on text BPB.

    All points share one fixed-R1 backbone and identical active payload FLOPs;
    only the operator pool differs. KG-08 asks whether BPB *falls* as the stored
    pool N grows at constant active compute. KG-05 asks whether the learned pool
    beats the no-operator baseline AND the anti-fooling controls (random-frozen
    payloads, matched-active dense) with disjoint 95% CIs. Both can fail.
    """
    pts = [p for p in pts if p.get("task") == "text"]
    groups = {}
    for p in pts:
        groups.setdefault(p["point_id"], []).append(p)
    rows = {}
    for pid, ps in groups.items():
        bpb_m, bpb_ci = _mean_ci([p["val_bpb"] for p in ps])
        rows[pid] = {
            "point_id": pid, "n_seeds": len(ps), "val_bpb": bpb_m, "val_bpb_ci": bpb_ci,
            "n_operators": ps[0].get("n_operators", 0),
            "operator_mode": ps[0].get("operator_mode", "-"),
            "use_operators": ps[0].get("use_operators", False),
            "active_op_flops": ps[0].get("active_op_flops", 0),
            "stored_op_params": ps[0].get("stored_op_params", 0),
            "params": ps[0].get("params", 0),
        }

    def lb(r):  # upper edge of the better (lower) direction
        return r["val_bpb"] - (r["val_bpb_ci"] or 0.0)

    def ub(r):
        return r["val_bpb"] + (r["val_bpb_ci"] or 0.0)

    learned = sorted([r for r in rows.values()
                      if r["use_operators"] and r["operator_mode"] == "learned"],
                     key=lambda r: r["n_operators"])
    # KG-08: BPB strictly lower at the largest pool vs the smallest, CIs disjoint.
    slope_ok, slope_desc = False, "insufficient learned points for a slope"
    if len(learned) >= 2:
        lo, hi = learned[0], learned[-1]
        gain = lo["val_bpb"] - hi["val_bpb"]
        disjoint = ub(hi) < lb(lo)
        slope_ok = gain > 0 and disjoint
        slope_desc = (f"N {lo['n_operators']}->{hi['n_operators']}: BPB "
                      f"{lo['val_bpb']:.4f}->{hi['val_bpb']:.4f} (gain {gain:+.4f}, "
                      f"CIs {'disjoint' if disjoint else 'overlap'})")

    best = min(learned, key=lambda r: r["val_bpb"]) if learned else None
    comparisons, kg05_ok = {}, best is not None
    ctrl_map = {"baseline": [r for r in rows.values() if not r["use_operators"]],
                "random_frozen": [r for r in rows.values() if r["operator_mode"] == "random_frozen"],
                "dense": [r for r in rows.values() if r["operator_mode"] == "dense"]}
    for name, cs in ctrl_map.items():
        if best and cs:
            c = cs[0]
            win = ub(best) < lb(c)  # best learned strictly better (lower BPB)
            comparisons[name] = {"ctrl_bpb": c["val_bpb"], "ctrl_ci": c["val_bpb_ci"],
                                 "best_bpb": best["val_bpb"], "best_beats_ctrl": bool(win)}
            kg05_ok = kg05_ok and win
        elif name == "baseline":
            kg05_ok = False  # no baseline => cannot assert value
    return {
        "task": "text",
        "rows": sorted(rows.values(), key=lambda r: (not r["use_operators"], r["n_operators"])),
        "best_learned": best["point_id"] if best else None,
        "KG08_slope_ok": slope_ok, "KG08_verdict": slope_desc,
        "KG05_comparisons": comparisons, "KG05_pass": bool(kg05_ok),
        "KG05_verdict": ("PASS: learned operators beat the no-operator baseline and every "
                         "anti-fooling control at matched active compute (disjoint 95% CIs)"
                         if kg05_ok else
                         "NOT PROVEN: learned operators do not strictly beat all controls within CIs"),
    }


def run_operators():
    device, amp = device_amp()
    hp = read_hp()
    task = "text"  # operator value is measured on real-text held-out BPB
    op_ns = [int(x) for x in env("OP_NS", "512,2048,8192").split(",") if x != ""]
    seeds = [int(x) for x in env("SEEDS", "0,1").split(",") if x != ""]
    n_active = env_i("N_ACTIVE", 4)
    d_operator = env_i("D_OPERATOR", 0)
    n_op_layers = env_i("N_OPERATOR_LAYERS", 2)
    include_baseline = env("BASELINE", "1") not in ("0", "false", "")
    include_controls = env("CONTROLS", "1") not in ("0", "false", "")
    run = env("RUN", "ops")
    results_dir = env("RESULTS_DIR", "ablation/operators")

    def op_over(**kw):
        base = dict(use_operators=True, n_active=n_active, n_operator_layers=n_op_layers)
        if d_operator > 0:
            base["d_operator"] = d_operator
        base.update(kw)
        return base

    plan = []
    if include_baseline:
        plan.append(("baseline", {"use_operators": False}))
    for N in op_ns:
        plan.append((f"learned_N{N}", op_over(n_operators=N, operator_mode="learned")))
    if include_controls and op_ns:
        maxN = max(op_ns)
        plan.append((f"randfrozen_N{maxN}",
                     op_over(n_operators=maxN, operator_mode="random_frozen")))
        plan.append(("dense_matched", op_over(n_operators=maxN, operator_mode="dense")))

    log(f"OPERATORS points={len(plan)} seeds={seeds} n_active={n_active} "
        f"n_op_layers={n_op_layers} op_ns={op_ns} steps/pt={hp['max_steps']}")

    outdir = Path("/tmp/op_grid"); outdir.mkdir(exist_ok=True)
    results = []

    def flush(tag):
        (outdir / "results.jsonl").write_text("\n".join(json.dumps(r) for r in results))
        _push(outdir / "results.jsonl", f"{results_dir}/{run}/results.jsonl",
              f"operator results ({len(results)} pts, {tag})")

    for (pid, over) in plan:
        for s in seeds:
            results.append(train_one(task, "fixed", 1, 0.0, s, hp, device, amp,
                                     extra_over=over, point_id=pid))
            flush(f"{pid}:s{s}")

    score = build_operator_scorecard(results)
    (outdir / "scorecard.json").write_text(json.dumps(score, indent=2))
    _push(outdir / "scorecard.json", f"{results_dir}/{run}/scorecard.json",
          "operator scorecard (KG-05/KG-08)")
    log(f"VERDICT KG05_pass={score['KG05_pass']} :: {score['KG05_verdict']}")
    log(f"VERDICT KG08_slope_ok={score['KG08_slope_ok']} :: {score['KG08_verdict']}")
    log(f"SCORECARD {json.dumps(score, indent=2)}")


def build_refiner_scorecard(pts):
    """GLR-1 iterative-refiner value on real-text held-out BPB.

    Two falsifiable questions, both able to fail:
      * refiner_beats_baseline -- Trident+refiner (eval at max budget) has strictly
        lower BPB than the no-refiner baseline at matched steps/data, disjoint 95%
        CIs. (Params differ: the refiner adds one weight-shared block; reported.)
      * budget_monotone -- on the SAME trained refiner model, test-time BPB falls as
        the refinement budget grows (perfectly param-matched depth test). This is
        the core GLR "test-time compute buys depth" claim on real bytes.
    """
    pts = [p for p in pts if p.get("task") == "text"]
    groups = {}
    for p in pts:
        groups.setdefault(p["point_id"], []).append(p)
    rows = {}
    for pid, ps in groups.items():
        m, ci = _mean_ci([p["val_bpb"] for p in ps])
        rows[pid] = {"point_id": pid, "n_seeds": len(ps), "val_bpb": m, "val_bpb_ci": ci,
                     "params": ps[0].get("params", 0),
                     "use_refiner": ps[0].get("use_refiner", False),
                     "refine_budgets": ps[0].get("refine_budgets", []),
                     "val_bpb_by_budget": ps[0].get("val_bpb_by_budget")}

    def lb(r):
        return r["val_bpb"] - (r["val_bpb_ci"] or 0.0)

    def ub(r):
        return r["val_bpb"] + (r["val_bpb_ci"] or 0.0)

    base = next((r for r in rows.values() if not r["use_refiner"]), None)
    ref = next((r for r in rows.values() if r["use_refiner"]), None)
    beats, verdict = False, "no baseline/refiner pair to compare"
    if base and ref:
        beats = ub(ref) < lb(base)
        disj = ub(ref) < lb(base) or ub(base) < lb(ref)
        verdict = (f"{'PASS' if beats else 'NOT PROVEN'}: refiner {ref['val_bpb']:.4f} vs "
                   f"baseline {base['val_bpb']:.4f} (delta {ref['val_bpb']-base['val_bpb']:+.4f}, "
                   f"CIs {'disjoint' if disj else 'overlap'}); params "
                   f"{ref['params']/1e6:.2f}M vs {base['params']/1e6:.2f}M")

    mono, bdesc = False, "no budget sweep recorded"
    if ref and ref.get("val_bpb_by_budget"):
        sweep = ref["val_bpb_by_budget"]
        bs = sorted(sweep, key=lambda k: int(k))
        vals = [sweep[b] for b in bs]
        mono = all(vals[i] >= vals[i + 1] - 1e-4 for i in range(len(vals) - 1)) and (vals[0] - vals[-1] > 0)
        bdesc = (" ".join(f"b{b}={sweep[b]:.4f}" for b in bs)
                 + f" (gain b{bs[0]}->b{bs[-1]}: {vals[0]-vals[-1]:+.4f})")

    return {
        "task": "text",
        "rows": sorted(rows.values(), key=lambda r: r["use_refiner"]),
        "refiner_beats_baseline": bool(beats), "verdict": verdict,
        "budget_monotone": bool(mono), "budget_desc": bdesc,
    }


def run_refiner():
    device, amp = device_amp()
    hp = read_hp()
    task = env("TASK", "text")
    seeds = [int(x) for x in env("SEEDS", "0,1").split(",") if x != ""]
    budgets = tuple(int(x) for x in env("REFINE_BUDGETS", "1,2,4,8").split(",") if x != "")
    ffn_mult = env_i("REFINER_FFN_MULT", 4)
    dsw = env_f("DEEP_SUP_W", 0.1)
    hkw = env_f("HALT_KL_W", 0.01)
    hprior = env_f("HALT_PRIOR", 0.2)
    include_baseline = env("BASELINE", "1") not in ("0", "false", "")
    run = env("RUN", "refiner")
    results_dir = env("RESULTS_DIR", "ablation/refiner")

    ref_over = dict(use_refiner=True, refine_budgets=budgets, refiner_ffn_mult=ffn_mult,
                    deep_supervision_weight=dsw, halt_kl_weight=hkw, halt_prior=hprior)
    plan = []
    if include_baseline:
        plan.append(("baseline", {"use_refiner": False}))
    plan.append((f"refiner_b{max(budgets)}", ref_over))

    log(f"REFINER points={len(plan)} seeds={seeds} budgets={budgets} ffn_mult={ffn_mult} "
        f"dsw={dsw} hkw={hkw} steps/pt={hp['max_steps']}")
    outdir = Path("/tmp/ref_grid"); outdir.mkdir(exist_ok=True)
    results = []

    def flush(tag):
        (outdir / "results.jsonl").write_text("\n".join(json.dumps(r) for r in results))
        _push(outdir / "results.jsonl", f"{results_dir}/{run}/results.jsonl",
              f"refiner results ({len(results)} pts, {tag})")

    for (pid, over) in plan:
        for s in seeds:
            results.append(train_one(task, "fixed", 1, 0.0, s, hp, device, amp,
                                     extra_over=over, point_id=pid))
            flush(f"{pid}:s{s}")

    score = build_refiner_scorecard(results)
    (outdir / "scorecard.json").write_text(json.dumps(score, indent=2))
    _push(outdir / "scorecard.json", f"{results_dir}/{run}/scorecard.json",
          "refiner scorecard (GLR-1 port)")
    log(f"VERDICT refiner_beats_baseline={score['refiner_beats_baseline']} :: {score['verdict']}")
    log(f"VERDICT budget_monotone={score['budget_monotone']} :: {score['budget_desc']}")
    log(f"SCORECARD {json.dumps(score, indent=2)}")


def train_point():
    device, amp = device_amp()
    hp = read_hp()
    task = env("TASK", "text")
    rate_mode = env("RATE_MODE", "fixed")
    fixed_r = env_i("FIXED_R", 1)
    lambda_rate = env_f("LAMBDA_RATE", 0.0)
    seed = env_i("SEED", 0)
    res = train_one(task, rate_mode, fixed_r if rate_mode == "fixed" else None,
                    lambda_rate, seed, hp, device, amp)
    run = env("RUN", "vr")
    results_dir = env("RESULTS_DIR", "ablation/variable_rate")
    fname = f"{task}__{res['point_id']}__seed{seed}.json"
    outdir = Path("/tmp/vr_result"); outdir.mkdir(exist_ok=True)
    (outdir / fname).write_text(json.dumps(res, indent=2))
    _push(outdir / fname, f"{results_dir}/{run}/{fname}", f"vr point {res['point_id']} seed{seed}")


def aggregate():
    from huggingface_hub import HfApi, hf_hub_download
    api = HfApi()
    repo = env("CODE_REPO", "farguney/trident")
    run = env("RUN", "vr")
    results_dir = env("RESULTS_DIR", "ablation/variable_rate")
    task = env("TASK", "text")
    prefix = f"{results_dir}/{run}/"
    files = [f for f in api.list_repo_files(repo, repo_type="model")
             if f.startswith(prefix) and f.endswith(".json") and "scorecard" not in f]
    pts = []
    for f in files:
        pts.append(json.loads(Path(hf_hub_download(repo, f, repo_type="model")).read_text()))
    # also fold any results_*.jsonl
    for f in api.list_repo_files(repo, repo_type="model"):
        if f.startswith(prefix) and f.endswith(".jsonl"):
            for line in Path(hf_hub_download(repo, f, repo_type="model")).read_text().splitlines():
                if line.strip():
                    pts.append(json.loads(line))
    log(f"aggregating {len(pts)} points for task={task}")
    score = build_scorecard(pts, task)
    outdir = Path("/tmp/vr_score"); outdir.mkdir(exist_ok=True)
    (outdir / "scorecard.json").write_text(json.dumps(score, indent=2))
    _push(outdir / "scorecard.json", f"{prefix}scorecard_{task}.json", f"vr scorecard ({task})")
    log(f"SCORECARD {json.dumps(score, indent=2)}")


def main():
    import_trident()
    mode = env("MODE", "grid")
    if mode == "grid":
        run_grid()
    elif mode == "operators":
        run_operators()
    elif mode == "refiner":
        run_refiner()
    elif mode == "train":
        train_point()
    elif mode == "aggregate":
        aggregate()
    else:
        raise ValueError(f"MODE must be grid|operators|refiner|train|aggregate, got {mode}")


if __name__ == "__main__":
    main()