train/dist_bed.py

"""dist_bed.py — the CLIP-distillation training bed that produced every
released checkpoint.

Student: 8.66M-parameter ViT (d=240, depth 12, heads 4, patch 16, 160px,
CLS readout, linear head to the target dimension). Optimizer: Adam lr 3e-4,
weight_decay 0, fp32, TF32 off, batch 256, fixed CRC-derived seeds. The
one exception is the cbert_full replication arm, which reproduces its
source recipe exactly (AdamW wd 0.01 + 1000-step linear warmup + cosine to
1e-6 + grad-clip 1.0) — flagged in its startup banner.

ARMS
  feature_mse     per-element MSE vs the deployment teacher (mimicry)
  infonce         symmetric InfoNCE vs the deployment teacher, temp 0.07
  siglip_pairwise decoupled pairwise sigmoid vs the SigLIP teacher (768-d)
  affinity_kl     similarity-row KL (affinity mimicking)
  blueprint       infonce + 0.3*procrustes_sq + 1e-3 spread force
  consensus_gpa   per-sample MSE vs the 5-teacher GPA consensus mean
  consensus_nce_mse  infonce + per-sample MSE vs the GPA mean (champion)
  x3_full/x3_bce/x3_autograd/nce_mse_cv/cbert_full  battery variants
  zs_floor / teacher_ceiling  controls

EVALS (identical for every arm; deterministic — re-runs reproduce to 4dp)
  CIFAR-10/100 zero-shot via the matched text tower; COCO val retrieval
  R@1/R@5; cosine agreement to the arm target; embedding-geometry readouts
  (spread CV on a fixed 16-d projection; effective rank). JSONL ledgers.

DATA LAYOUT (set DIST_DATA_ROOT; defaults to ./data)
  $DIST_DATA_ROOT/dist/bank/{config}/{split}-*.parquet   COCO teacher bank
  $DIST_DATA_ROOT/dist/coco/{val2017,train2017,annotations}
  $DIST_DATA_ROOT/dist/ckpts/                            outputs
  CC12M split: see cc12m_data.py (tar-offset index + feature shards).

RUN
  python dist_bed.py --smoke [--arm feature_mse]
  python dist_bed.py --gates
  python dist_bed.py --arm infonce --seed 0 --steps 88000 --split cc12m
  python dist_bed.py --eval-only <ckpt.pt> [--extra-evals]
"""
import argparse
import json
import math
import os
import sys
import time
import zlib

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets as tvds
from torchvision import transforms as T
from torchvision.transforms import InterpolationMode


def _root():
    d = os.path.abspath(os.getcwd())
    while True:
        if os.path.exists(os.path.join(d, "TRAINING.md")):
            return d
        p = os.path.dirname(d)
        if p == d:
            return os.getcwd()
        d = p

ROOT = _root()
if os.path.join(ROOT, "tools") not in sys.path:
    sys.path.insert(0, os.path.join(ROOT, "tools"))

from dist_align_gate import (CFG_MODEL, CLIP_MEAN, CLIP_STD, PRIMARY,   # noqa: E402
                             _feat, first_captions, laion_local_dir,
                             load_bank, load_siglip_text_tokenizer,
                             load_vision)
from vitals import pentachoron_cv                                # noqa: E402
from losses import (EmbeddingAutograd, a7_grid_infonce,             # noqa: E402
                        affinity_kl_loss, collinearity_novelty, prim_sq,
                        procrustes_sq, siglip_pairwise_loss,
                        x3_centered_cos_align, x3_cv_loss)

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
DEV = "cuda" if torch.cuda.is_available() else "cpu"

DATA_ROOT = os.environ.get("DIST_DATA_ROOT", "./data")
DATA = os.path.join(DATA_ROOT, "dist")
COCO = os.path.join(DATA, "coco")
RUNS = os.path.join(ROOT, "tools", "dist_runs")
CKPTS = os.path.join(DATA, "ckpts")

BATCH = 256
SIGLIP_CFG = "siglip_b16_384"
CONSENSUS = ("clip_b16_laion2b", "clip_b32_openai", "clip_b32_laion2b",
             "clip_b32_datacomp", "clip_b16_openai")
TRAIN_ARMS = ("feature_mse", "infonce", "siglip_pairwise", "affinity_kl",
              "blueprint", "consensus_gpa", "consensus_nce_mse",
              "x3_full", "x3_bce", "x3_autograd", "nce_mse_cv",
              "cbert_full")
X3_ARMS = ("x3_full", "x3_bce", "x3_autograd", "nce_mse_cv")
ARM_TARGET_GPA_EXTRA = ("cbert_full",)
ARMS = TRAIN_ARMS + ("zs_floor", "teacher_ceiling", "hybrid_best")
ARM_TARGET = {a: PRIMARY for a in ARMS}
ARM_TARGET["siglip_pairwise"] = SIGLIP_CFG
ARM_TARGET["consensus_gpa"] = "gpa"
ARM_TARGET["consensus_nce_mse"] = "gpa"
for _a in X3_ARMS + ARM_TARGET_GPA_EXTRA:
    ARM_TARGET[_a] = "gpa"
CIFAR10_CLASSES = ("airplane", "automobile", "bird", "cat", "deer", "dog",
                   "frog", "horse", "ship", "truck")
BLUEPRINT_CV_WEIGHT = 1e-3
EXTRA_EVALS = False


def seed_for(name):
    return zlib.crc32(name.encode()) & 0x7FFFFFFF


# ================================================================== STUDENT
class Block(nn.Module):
    def __init__(self, d, heads):
        super().__init__()
        self.n1 = nn.LayerNorm(d)
        self.qkv = nn.Linear(d, 3 * d)
        self.proj = nn.Linear(d, d)
        self.n2 = nn.LayerNorm(d)
        self.fc1 = nn.Linear(d, 4 * d)
        self.fc2 = nn.Linear(4 * d, d)
        self.heads = heads

    def forward(self, x):
        B, N, C = x.shape
        q, k, v = (self.qkv(self.n1(x))
                   .reshape(B, N, 3, self.heads, C // self.heads)
                   .permute(2, 0, 3, 1, 4))
        a = F.scaled_dot_product_attention(q, k, v)
        x = x + self.proj(a.transpose(1, 2).reshape(B, N, C))
        return x + self.fc2(F.gelu(self.fc1(self.n2(x))))


class Student(nn.Module):
    """ViT-Ti-class: d=240 depth 12 heads 4 patch 16 img 160. CLS-token
    readout, Linear head to the arm's target dim."""
    def __init__(self, out_dim=512, d=240, depth=12, heads=4, patch=16,
                 img=160):
        super().__init__()
        self.patch = nn.Conv2d(3, d, patch, patch)
        self.cls = nn.Parameter(torch.zeros(1, 1, d))
        self.pos = nn.Parameter(torch.zeros(1, (img // patch) ** 2 + 1, d))
        self.blocks = nn.ModuleList(Block(d, heads) for _ in range(depth))
        self.norm = nn.LayerNorm(d)
        self.head = nn.Linear(d, out_dim)
        self.apply(self._init)
        nn.init.trunc_normal_(self.cls, std=0.02)
        nn.init.trunc_normal_(self.pos, std=0.02)

    @staticmethod
    def _init(m):
        if isinstance(m, (nn.Linear, nn.Conv2d)):
            nn.init.trunc_normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)

    def forward_features(self, x):
        x = self.patch(x).flatten(2).transpose(1, 2)
        x = torch.cat([self.cls.expand(x.shape[0], -1, -1), x], 1) + self.pos
        for b in self.blocks:
            x = b(x)
        return self.norm(x)[:, 0]

    def forward(self, x):
        return self.head(self.forward_features(x))


def build_student(arm, seed, out_dim):
    torch.manual_seed(seed_for(f"dist:{arm}:{seed}"))
    m = Student(out_dim=out_dim).to(DEV)
    tot = sum(p.numel() for p in m.parameters())
    hd = sum(p.numel() for p in m.head.parameters())
    print(f"[{arm} s{seed}] params total {tot:,} trunk {tot - hd:,} "
          f"head({out_dim}) {hd:,}", flush=True)
    return m


# ===================================================================== DATA
STUDENT_TF = T.Compose([
    T.Resize(182, interpolation=InterpolationMode.BICUBIC),
    T.CenterCrop(160), T.ToTensor(), T.Normalize(CLIP_MEAN, CLIP_STD)])


class CocoImages(Dataset):
    """image_id-joined COCO jpgs -> (tensor, row_idx). Targets are indexed
    by row_idx in the main process (keeps DataLoader workers light)."""
    def __init__(self, split, ids):
        root = os.path.join(COCO, f"{split}2017")
        if not os.path.isdir(root):
            raise FileNotFoundError(
                f"COCO {split}2017 images absent at {root} — still "
                f"downloading? The bed cannot run this split yet.")
        self.paths = [os.path.join(root, f"{int(i):012d}.jpg") for i in ids]
        missing = [p for p in self.paths if not os.path.isfile(p)]
        if missing:
            raise FileNotFoundError(
                f"{len(missing)}/{len(ids)} {split}2017 jpgs missing "
                f"(e.g. {missing[0]}) — download still landing?")
        self.tf = STUDENT_TF

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, i):
        return self.tf(Image.open(self.paths[i]).convert("RGB")), i


def build_gpa(split):
    """GPA mean shape over the consensus five (all 512-d), cached once.
    L2-normalize towers -> 5 iters of { orthogonal-Procrustes each tower to
    the mean over a crc32-fixed 10k subsample (SVD fp64), re-mean, renorm }.
    Rows in the PRIMARY parquet order (towers reindexed by image_id)."""
    path = os.path.join(DATA, f"gpa_targets_{split}.pt")
    if os.path.isfile(path):
        return path
    towers, ids0 = [], None
    for cfg in CONSENSUS:
        ids, feats = load_bank(cfg, split)
        z = F.normalize(feats, dim=-1)
        if ids0 is None:
            ids0 = ids
        else:
            pos = {int(v): i for i, v in enumerate(ids)}
            z = z[torch.tensor([pos[int(v)] for v in ids0])]
        towers.append(z)
    X = torch.stack(towers)                             # (5, N, 512)
    N = X.shape[1]
    g = torch.Generator().manual_seed(seed_for(f"dist-gpa-sub:{split}"))
    sub = torch.randperm(N, generator=g)[:min(10000, N)]
    mean = F.normalize(X.mean(0), dim=-1)
    aligned = X.clone()
    print(f"[gpa {split}] towers {X.shape[0]} rows {N} subsample {len(sub)}")
    for it in range(5):
        for t in range(X.shape[0]):
            U, _, Vt = torch.linalg.svd(
                X[t][sub].t().double() @ mean[sub].double())
            aligned[t] = X[t] @ (U @ Vt).float()
        mean = F.normalize(aligned.mean(0), dim=-1)
        cos = (F.normalize(aligned, dim=-1) * mean.unsqueeze(0)).sum(-1).mean(1)
        print(f"[gpa {split}] iter {it + 1} align cos "
              + " ".join(f"{float(c):.4f}" for c in cos)
              + f" | mean {float(cos.mean()):.4f}", flush=True)
    torch.save({"image_ids": torch.from_numpy(np.array(ids0, dtype=np.int64)),
                "targets": mean, "towers": list(CONSENSUS), "iters": 5}, path)
    print(f"[gpa {split}] cached -> {path}")
    return path


def load_targets(arm, split):
    """-> (ids int64 np, targets fp32 L2-normalized, row-aligned to ids)."""
    t = ARM_TARGET[arm]
    if t == "gpa":
        d = torch.load(build_gpa(split), map_location="cpu",
                       weights_only=True)
        return d["image_ids"].numpy(), d["targets"]
    ids, feats = load_bank(t, split)
    return ids, F.normalize(feats, dim=-1)


# ============================================================ TEACHER HEADS
def _text_tower(tower):
    from transformers import AutoTokenizer
    from dist_align_gate import _from_pretrained_fp32
    if tower == "laion":
        d = laion_local_dir()
        return (_from_pretrained_fp32(d).to(DEV).eval(),
                AutoTokenizer.from_pretrained(d), False)
    m = _from_pretrained_fp32(CFG_MODEL[SIGLIP_CFG]).to(DEV).eval()
    return m, load_siglip_text_tokenizer(), True


@torch.no_grad()
def _embed_text(model, tok, texts, siglip, bs=128):
    outs = []
    for i in range(0, len(texts), bs):
        kw = (dict(padding="max_length", truncation=True, max_length=64)
              if siglip else dict(padding=True, truncation=True))
        tk = tok(texts[i:i + bs], return_tensors="pt", **kw)
        args = {"input_ids": tk["input_ids"].to(DEV)}
        if not siglip:          # canonical SigLIP: full attention over pads
            args["attention_mask"] = tk["attention_mask"].to(DEV)
        outs.append(_feat(model.get_text_features(**args)).float().cpu())
    return F.normalize(torch.cat(outs), dim=-1)


def text_val(tower):
    """First-caption text embeddings for the 5k val images, cached."""
    path = os.path.join(DATA, f"text_val_{tower}.pt")
    if os.path.isfile(path):
        return torch.load(path, map_location="cpu", weights_only=True)
    ids, _ = load_bank(PRIMARY if tower == "laion" else SIGLIP_CFG, "val")
    caps = first_captions()
    model, tok, sg = _text_tower(tower)
    emb = _embed_text(model, tok, [caps[int(i)] for i in ids], sg)
    del model
    torch.cuda.empty_cache()
    obj = {"image_ids": torch.from_numpy(np.array(ids, dtype=np.int64)),
           "emb": emb}
    torch.save(obj, path)
    print(f"[text_val {tower}] cached {tuple(emb.shape)} -> {path}")
    return obj


def text_classes(tower, name, classes):
    path = os.path.join(DATA, f"text_{name}_{tower}.pt")
    if os.path.isfile(path):
        return torch.load(path, map_location="cpu", weights_only=True)
    model, tok, sg = _text_tower(tower)
    emb = _embed_text(model, tok,
                      [f"a photo of a {c}" for c in classes], sg)
    del model
    torch.cuda.empty_cache()
    obj = {"classes": list(classes), "emb": emb}
    torch.save(obj, path)
    return obj


# ===================================================================== LOSS
def _diff_cv(z16, cv_target=0.20, weight=BLUEPRINT_CV_WEIGHT, n_sets=64):
    """Differentiable pentachoron-CV band term on PROJECTED student
    embeddings (16-d rows on the unit sphere) — same statistic as
    exp017:154-186 but fp32 and on embeddings, per the campaign spec. The
    weight is capped at 1e-3; subset draw crc32-fixed so the
    subset pattern is deterministic across steps."""
    assert weight <= 1e-3, "CV force above 1e-3 is prohibited"
    x = F.normalize(z16, dim=-1)
    n = x.shape[0]
    g = torch.Generator().manual_seed(seed_for("dist-cv-subsets"))
    idx = torch.stack([torch.randperm(n, generator=g)[:5]
                       for _ in range(n_sets)]).to(x.device)
    pts = x[idx]                                        # (n_sets, 5, 16)
    d2 = torch.cdist(pts, pts).pow(2)
    cm = torch.ones(n_sets, 6, 6, dtype=pts.dtype, device=x.device)
    cm[:, 0, 0] = 0.0
    cm[:, 1:, 1:] = d2
    v = (-torch.linalg.det(cm) / 9216.0).clamp_min(1e-24).sqrt()
    cv = v.std() / v.mean().clamp_min(1e-12)
    return weight * (cv - cv_target).abs()


def _fit_teacher_head(targets, labels, epochs=5, bs=2048, lr=1e-3):
    """Frozen teacher-space classifier for the x3 BCE term — the analog of
    x3's frozen calibrated soup head: a linear 80-class head fit ON THE GPA
    TARGETS (teacher side), then frozen. Deterministic (crc32 seeds)."""
    torch.manual_seed(seed_for("x3-teacher-head"))
    head = nn.Linear(targets.shape[-1], labels.shape[-1]).to(DEV)
    opt = torch.optim.Adam(head.parameters(), lr=lr, weight_decay=0.0)
    Tt, Y = targets.float().to(DEV), labels.float().to(DEV)
    n = Tt.shape[0]
    for ep in range(epochs):
        g = torch.Generator().manual_seed(seed_for(f"x3-head-ep:{ep}"))
        perm = torch.randperm(n, generator=g)
        for i in range(0, n - bs + 1, bs):
            j = perm[i:i + bs]
            opt.zero_grad(set_to_none=True)
            F.binary_cross_entropy_with_logits(head(Tt[j]), Y[j]).backward()
            opt.step()
    head.eval()
    for p in head.parameters():
        p.requires_grad_(False)
    return head


def loss_of(arm, zs, zt, ctx):
    """zs, zt L2-normalized. consensus_gpa uses the feature_mse form against
    the GPA mean (the target changes, not the primitive)."""
    if arm in ("feature_mse", "consensus_gpa"):
        return prim_sq(zs, zt).mean()
    if arm == "consensus_nce_mse":
        # CaptionBert Stage-A composite on the GPA target. Per-SAMPLE MSE
        # (= ||zs-zt||^2 = 2(1-cos)) so the 1.0/1.0 weighting is real —
        # per-element mean is ~1000x smaller than the NCE term at init.
        return (a7_grid_infonce(zs, zt, temp=0.07)
                + prim_sq(zs, zt).sum(-1).mean())
    if arm == "cbert_full":
        # Source recipe VERBATIM (replication arm): per-ELEMENT
        # MSE at 1.0 and CV weight 0.1 — 100x the usual ceiling,
        # sanctioned ONLY as a flagged replication exception.
        return (a7_grid_infonce(zs, zt, temp=0.07)
                + prim_sq(zs, zt).mean()
                + 0.1 * x3_cv_loss(zs, ctx["cv_target"]))
    if arm in X3_ARMS:
        if arm in ("x3_full", "x3_autograd"):
            zs = EmbeddingAutograd.apply(zs, zs, ctx["anchors"], 0.01, 1.0)
        if arm == "x3_full":
            # faithful x3 stack incl. its per-ELEMENT MSE at 0.5 (measured
            # ~inert vs the NCE term — disclosed, kept for fidelity) and
            # its CV at 0.001 = the house ceiling.
            return (a7_grid_infonce(zs, zt, temp=0.07)
                    + 0.5 * prim_sq(zs, zt).mean()
                    + 0.3 * F.binary_cross_entropy_with_logits(
                        ctx["head"](zs), ctx["yb"])
                    + 0.5 * x3_centered_cos_align(zs, zt)
                    + 1e-3 * x3_cv_loss(zs, ctx["cv_target"]))
        base = (a7_grid_infonce(zs, zt, temp=0.07)
                + prim_sq(zs, zt).sum(-1).mean())
        if arm == "x3_bce":
            return base + 0.3 * F.binary_cross_entropy_with_logits(
                ctx["head"](zs), ctx["yb"])
        if arm == "nce_mse_cv":
            return base + 1e-3 * x3_cv_loss(zs, ctx["cv_target"])
        return base                       # x3_autograd: EA-wrapped base
    if arm == "infonce":
        return a7_grid_infonce(zs, zt, temp=0.07)
    if arm == "siglip_pairwise":
        return siglip_pairwise_loss(zs, zt)
    if arm == "affinity_kl":
        return affinity_kl_loss(zs, zt, temp=0.07)
    if arm == "blueprint":
        return (a7_grid_infonce(zs, zt, temp=0.07)
                + 0.3 * procrustes_sq(zs, zt)
                + _diff_cv(zs @ ctx["P16"]))
    raise ValueError(arm)


# ==================================================================== EVALS
def student_embed01(model):
    mean = torch.tensor(CLIP_MEAN).view(1, 3, 1, 1).to(DEV)
    std = torch.tensor(CLIP_STD).view(1, 3, 1, 1).to(DEV)

    @torch.no_grad()
    def f(x01):
        x = F.interpolate(x01, size=(160, 160), mode="bicubic",
                          align_corners=False).clamp(0, 1)
        return F.normalize(model((x - mean) / std), dim=-1)
    return f


def teacher_embed01(model):
    """LAION B/16 vision tower on [0,1] tensors at its native 224 (tensor
    bicubic upscale — CIFAR is not in the bank, so the ceiling embeds it with
    the real tower)."""
    mean = torch.tensor(CLIP_MEAN).view(1, 3, 1, 1).to(DEV)
    std = torch.tensor(CLIP_STD).view(1, 3, 1, 1).to(DEV)

    @torch.no_grad()
    def f(x01):
        x = F.interpolate(x01, size=(224, 224), mode="bicubic",
                          align_corners=False).clamp(0, 1)
        return F.normalize(_feat(model.get_image_features(
            pixel_values=(x - mean) / std)), dim=-1)
    return f


@torch.no_grad()
def zeroshot_classes(embed01, txt, dataset, subset_seed, subset=None, bs=500):
    x, y = dataset
    if subset:
        g = torch.Generator().manual_seed(seed_for(subset_seed))
        keep = torch.randperm(len(y), generator=g)[:subset]
        x, y = x[keep], y[keep]
    txt = txt.to(DEV)
    ok = 0
    for i in range(0, len(y), bs):
        z = embed01(x[i:i + bs].to(DEV))
        ok += int((z @ txt.t()).argmax(-1).cpu().eq(y[i:i + bs]).sum())
    return ok / len(y)


def _cifar10():
    ds = tvds.CIFAR10(root=DATA_ROOT, train=False, download=True)
    x = torch.from_numpy(ds.data).permute(0, 3, 1, 2).contiguous()
    return x.float().div_(255.0), torch.tensor(ds.targets)


@torch.no_grad()
def embed_val_student(model, ids, subset=None, workers=6):
    if subset is not None:
        g = torch.Generator().manual_seed(seed_for("dist-eval-sub"))
        keep = torch.randperm(len(ids), generator=g)[:subset].numpy()
    else:
        keep = np.arange(len(ids))
    ds = CocoImages("val", ids[keep])
    dl = DataLoader(ds, batch_size=BATCH, shuffle=False, num_workers=workers,
                    pin_memory=True)
    emb = torch.empty(len(ds), model.head.out_features)
    for x, idx in dl:
        emb[idx] = F.normalize(model(x.to(DEV, non_blocking=True)),
                               dim=-1).float().cpu()
    return emb, ids[keep]


@torch.no_grad()
def coco_retrieval(emb, row_ids, tower):
    """Image->text R@1/R@5 against the cached first-caption embeddings."""
    tc = text_val(tower)
    pos = {int(v): i for i, v in enumerate(tc["image_ids"].tolist())}
    txt = tc["emb"][torch.tensor([pos[int(v)] for v in row_ids])].to(DEV)
    sim = emb.to(DEV) @ txt.t()
    ranks = (sim > sim.diag().unsqueeze(1)).sum(1)
    return (float((ranks == 0).float().mean()),
            float((ranks < 5).float().mean()))


@torch.no_grad()
def agreement(emb, row_ids, t_ids, targets):
    pos = {int(v): i for i, v in enumerate(t_ids)}
    idx = torch.tensor([pos[int(v)] for v in row_ids])
    return float((emb * targets[idx]).sum(-1).mean())


@torch.no_grad()
def geometry(emb):
    """Readout channel (CV is a readout here — the only CV FORCE in this bed
    is the blueprint arm's gated 1e-3 term): pentachoron CV on a fixed
    random D->16 projection + erank of the val matrix (fp64 gauges)."""
    D = emb.shape[1]
    gp = torch.Generator().manual_seed(seed_for(f"dist-proj16-{D}"))
    P = torch.randn(D, 16, generator=gp)
    cv = pentachoron_cv(emb.cpu() @ P,
                        generator=torch.Generator().manual_seed(
                            seed_for("dist-cv-eval")))
    s = torch.linalg.svdvals(emb.double().cpu())
    p = (s / s.sum()).clamp_min(1e-24)
    return float(cv), float(torch.exp(-(p * p.log()).sum()))


def evaluate_student(arm, model, tower, smoke, workers=6):
    model.eval()
    ids_v, targets_v = load_targets(arm, "val")
    emb, row_ids = embed_val_student(model, ids_v,
                                     subset=500 if smoke else None,
                                     workers=workers)
    txt10 = text_classes(tower, "cifar10", CIFAR10_CLASSES)["emb"]
    zs = zeroshot_classes(student_embed01(model), txt10, _cifar10(),
                          "dist-cifar-sub", subset=1000 if smoke else None)
    r1, r5 = coco_retrieval(emb, row_ids, tower)
    ag = agreement(emb, row_ids, ids_v, targets_v)
    cv, er = geometry(emb)
    rec = {"zs_cifar10": round(zs, 4), "coco_r1": round(r1, 4),
           "coco_r5": round(r5, 4), "agree_cos": round(ag, 4),
           "cv16": round(cv, 4), "erank": round(er, 2),
           "n_eval": int(len(row_ids)), "target": ARM_TARGET[arm]}
    if EXTRA_EVALS:
        try:
            ds = tvds.CIFAR100(root=DATA_ROOT, train=False, download=False)
            x = torch.from_numpy(ds.data).permute(0, 3, 1, 2).contiguous()
            data = (x.float().div_(255.0), torch.tensor(ds.targets))
            cls = [c.replace("_", " ") for c in ds.classes]
            txt100 = text_classes(tower, "cifar100", cls)["emb"]
            rec["zs_cifar100"] = round(zeroshot_classes(
                student_embed01(model), txt100, data, "dist-c100-sub",
                subset=1000 if smoke else None), 4)
        except RuntimeError as e:
            print(f"[extra-evals] CIFAR-100 SKIP: {e}")
    model.train()
    return rec


# ================================================================= TRAINING
def _ledger(rec, smoke):
    os.makedirs(RUNS, exist_ok=True)
    name = "smoke.jsonl" if smoke else f"{rec['arm']}_s{rec['seed']}.jsonl"
    out = os.path.join(RUNS, name)
    with open(out, "a", encoding="utf-8") as f:
        f.write(json.dumps(rec) + "\n")
    return out


def train_arm(arm, seed, steps, split="val", smoke=False, workers=6,
              bench=False):
    assert arm in ARMS, f"unknown arm {arm} (see --list)"
    if arm == "hybrid_best":
        raise SystemExit(
            "hybrid_best is RESERVED: minted only after the Night-0 verdicts "
            "(best vector arm + best grid arm at 0.5/0.5, and only if the "
            "gates show novelty >= 0.3 between them).")
    if DEV != "cuda":
        raise SystemExit("no CUDA device — this bed never CPU-trains "
                         "(smoke mode checks shapes/parse only).")
    torch.cuda.set_per_process_memory_fraction(0.73)
    torch.cuda.reset_peak_memory_stats()
    t0 = time.time()
    if arm == "teacher_ceiling":
        return teacher_ceiling_arm(seed, smoke, t0)

    tower = "siglip" if arm == "siglip_pairwise" else "laion"
    out_dim = 768 if arm == "siglip_pairwise" else 512
    cc_keys = None
    if split == "cc12m":
        from cc12m_data import (load_cc12m_gpa, load_cc12m_targets,
                                load_cc12m_tower)
        assert arm not in ("x3_full", "x3_bce"), \
            (f"{arm} needs the COCO 80-class multilabel BCE channel — "
             f"CC12M has no gold labels (COCO-scoped arm)")
        if ARM_TARGET[arm] == "gpa":
            cc_keys, t16 = load_cc12m_gpa()
        elif arm == "siglip_pairwise":
            cc_keys, t16 = load_cc12m_tower(SIGLIP_CFG)
        else:
            assert arm in ("feature_mse", "infonce", "affinity_kl",
                           "blueprint"), \
                f"{arm} has no CC12M target path yet"
            cc_keys, t16 = load_cc12m_targets()
        # fp16-RESIDENT, normalized per-batch on GPU (row-wise normalize is
        # batch-independent — numerically identical, HALF the cgroup
        # residency; two siglip OOM kills taught this)
        targets = t16
    else:
        ids, targets = load_targets(arm, split)
    model = build_student(arm, seed, out_dim)
    ctx = {}
    if arm == "blueprint":
        gp = torch.Generator().manual_seed(seed_for("dist-proj16-512"))
        ctx["P16"] = torch.randn(512, 16, generator=gp).to(DEV)
    if arm == "cbert_full":
        from bank_utils import _sample_cv
        g = torch.Generator().manual_seed(seed_for("x3-cv-target"))
        ctx["cv_target"] = _sample_cv(F.normalize(targets[:20000].float(), dim=-1), g)
        print(f"[{arm} s{seed}] REPLICATION ARM (prototype-replication "
              f"law): CaptionBert Stage-A verbatim — per-element MSE 1.0, "
              f"CV 0.1 (ceiling EXCEPTION, flagged), AdamW+warmup+cosine+"
              f"clip. cv_target {ctx['cv_target']:.4f}", flush=True)
    if arm in X3_ARMS:
        from bank_utils import _sample_cv, coco_multilabel
        ai = torch.linspace(0, len(targets) - 1, 512).long()
        ctx["anchors"] = F.normalize(targets[ai].float(), dim=-1).to(DEV)
        g = torch.Generator().manual_seed(seed_for("x3-cv-target"))
        ctx["cv_target"] = _sample_cv(F.normalize(targets[:20000].float(), dim=-1), g)
        if arm in ("x3_full", "x3_bce"):     # only these consume BCE
            ctx["labels"] = coco_multilabel(COCO, split, ids)
            ctx["head"] = _fit_teacher_head(targets, ctx["labels"])
        print(f"[{arm} s{seed}] x3-ctx: cv_target "
              f"{ctx['cv_target']:.4f}"
              + (" head frozen (teacher-space)" if "head" in ctx else ""),
              flush=True)

    train_steps = 0 if arm == "zs_floor" else steps
    stats, s_per_step = {}, 0.0
    if train_steps:
        if split == "cc12m":
            from cc12m_data import Cc12mImages
            ds = Cc12mImages(cc_keys, STUDENT_TF)
        else:
            ds = CocoImages(split, ids)
        print(f"[{arm} s{seed}] split {split} rows {len(ds)} steps "
              f"{train_steps} (~{train_steps * BATCH / len(ds):.1f} epochs)",
              flush=True)
        g = torch.Generator().manual_seed(seed_for(f"dist-data:{arm}:{seed}"))
        dl = DataLoader(ds, batch_size=BATCH, shuffle=True, generator=g,
                        num_workers=workers, persistent_workers=workers > 0,
                        pin_memory=True, drop_last=True)
        sched = None
        if arm == "cbert_full":
            # REPLICATION EXCEPTION (prototype-replication law): the
            # source recipe's AdamW(wd .01) + 1000-step linear warmup +
            # cosine to 1e-6 + grad-clip 1.0. Variants use pure Adam.
            opt = torch.optim.AdamW(model.parameters(), lr=3e-4,
                                    weight_decay=0.01)
            warm = torch.optim.lr_scheduler.LinearLR(
                opt, start_factor=0.01, end_factor=1.0, total_iters=1000)
            cos = torch.optim.lr_scheduler.CosineAnnealingLR(
                opt, T_max=max(1, train_steps - 1000), eta_min=1e-6)
            sched = torch.optim.lr_scheduler.SequentialLR(
                opt, [warm, cos], [1000])
        else:
            opt = torch.optim.Adam(model.parameters(), lr=3e-4,
                                   weight_decay=0.0)
        it = iter(dl)
        t_train = time.time()
        for step in range(train_steps):
            try:
                x, idx = next(it)
            except StopIteration:
                it = iter(dl)
                x, idx = next(it)
            zt = targets[idx].to(DEV, non_blocking=True)
            if zt.dtype == torch.half:       # cc12m fp16-resident path
                zt = F.normalize(zt.float(), dim=-1)
            if "labels" in ctx:
                ctx["yb"] = ctx["labels"][idx].to(DEV, non_blocking=True)
            opt.zero_grad(set_to_none=True)
            zs = F.normalize(model(x.to(DEV, non_blocking=True)), dim=-1)
            L = loss_of(arm, zs, zt, ctx)
            L.backward()
            if arm == "cbert_full":
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            if sched is not None:
                sched.step()
            lf = float(L.detach())
            if step == 0:
                stats["loss_step1"] = round(lf, 4)
            if step == 10:
                print(f"[{arm} s{seed}] step10 loss {lf:.4f} peak "
                      f"{torch.cuda.max_memory_allocated() / 2**30:.2f}GB "
                      f"{(time.time() - t_train) / 11:.2f}s/step", flush=True)
            ema = lf if step == 0 else 0.98 * ema + 0.02 * lf
            if step and step % 2000 == 0:
                print(f"[{arm} s{seed}] step{step}/{train_steps} "
                      f"loss {lf:.4f} ema {ema:.4f} "
                      f"{(time.time() - t_train) / (step + 1):.2f}s/step",
                      flush=True)
            if not smoke and step and step % 20000 == 0:
                os.makedirs(CKPTS, exist_ok=True)
                mid = os.path.join(CKPTS, f"{arm}_s{seed}_mid.pt")
                torch.save({"arm": arm, "seed": seed, "steps": step,
                            "split": split, "out_dim": out_dim,
                            "state_dict": model.state_dict()}, mid + ".tmp")
                os.replace(mid + ".tmp", mid)
            stats["loss_final"] = round(lf, 4)
        s_per_step = (time.time() - t_train) / train_steps
        del dl, it

    if bench:
        print(f"[BENCH {arm} s{seed}] {train_steps} steps "
              f"{s_per_step:.3f}s/step peak "
              f"{torch.cuda.max_memory_allocated() / 2**30:.2f}GB")
        return None

    ckpt = None
    if not smoke and train_steps:
        os.makedirs(CKPTS, exist_ok=True)
        ckpt = os.path.join(CKPTS, f"{arm}_s{seed}_t{train_steps}.pt")
        torch.save({"arm": arm, "seed": seed, "steps": train_steps,
                    "split": split, "out_dim": out_dim,
                    "state_dict": model.state_dict()}, ckpt)

    rec = {"arm": arm, "seed": seed, "steps": train_steps, "split": split,
           **stats, **evaluate_student(arm, model, tower, smoke, workers),
           "params": sum(p.numel() for p in model.parameters()),
           "s_per_step": round(s_per_step, 3),
           "wall_s": round(time.time() - t0, 1),
           "peak_gb": round(torch.cuda.max_memory_allocated() / 2**30, 2),
           "ckpt": ckpt, "smoke": smoke}
    out = _ledger(rec, smoke)
    print(f"[DONE {arm} s{seed}] zs {rec['zs_cifar10']:.4f} "
          f"r@1 {rec['coco_r1']:.4f} r@5 {rec['coco_r5']:.4f} "
          f"agree {rec['agree_cos']:.4f} cv16 {rec['cv16']:.3f} "
          f"erank {rec['erank']:.1f} ({rec['wall_s']}s) -> {out}", flush=True)
    return rec


def teacher_ceiling_arm(seed, smoke, t0):
    """Night-0 reference row: stored laion b16 bank val features stand in as
    the student for retrieval/agreement/geometry; the real LAION tower embeds
    CIFAR (the bank holds no CIFAR rows)."""
    ids, targets = load_targets("teacher_ceiling", "val")
    if smoke:
        g = torch.Generator().manual_seed(seed_for("dist-eval-sub"))
        keep = torch.randperm(len(ids), generator=g)[:500].numpy()
    else:
        keep = np.arange(len(ids))
    emb, row_ids = targets[keep], ids[keep]
    model, _ = load_vision(PRIMARY)
    txt10 = text_classes("laion", "cifar10", CIFAR10_CLASSES)["emb"]
    zs = zeroshot_classes(teacher_embed01(model), txt10, _cifar10(),
                          "dist-cifar-sub", subset=1000 if smoke else None)
    del model
    torch.cuda.empty_cache()
    r1, r5 = coco_retrieval(emb, row_ids, "laion")
    cv, er = geometry(emb)
    rec = {"arm": "teacher_ceiling", "seed": seed, "steps": 0, "split": "val",
           "zs_cifar10": round(zs, 4), "coco_r1": round(r1, 4),
           "coco_r5": round(r5, 4), "agree_cos": 1.0, "cv16": round(cv, 4),
           "erank": round(er, 2), "n_eval": int(len(row_ids)),
           "target": PRIMARY, "params": 0, "s_per_step": 0.0,
           "wall_s": round(time.time() - t0, 1),
           "peak_gb": round(torch.cuda.max_memory_allocated() / 2**30, 2),
           "ckpt": None, "smoke": smoke}
    out = _ledger(rec, smoke)
    print(f"[DONE teacher_ceiling] zs {zs:.4f} r@1 {r1:.4f} r@5 {r5:.4f} "
          f"cv16 {cv:.3f} erank {er:.1f} ({rec['wall_s']}s) -> {out}",
          flush=True)
    return rec


# ==================================================================== GATES
def run_gates(workers=0):
    """Collinearity matrix across all six trainable losses on ONE fixed val
    batch at init state (reuses loss_forms.collinearity_novelty; base =
    feature_mse). Shared trunk + both heads; grads zero-filled on the shared
    parameter support."""
    if DEV != "cuda":
        raise SystemExit("gates need the GPU (one forward on a 256 batch)")
    torch.cuda.set_per_process_memory_fraction(0.73)
    ids_l, t_l = load_targets("feature_mse", "val")
    ids_s, t_s = load_targets("siglip_pairwise", "val")
    ids_g, t_g = load_targets("consensus_gpa", "val")
    g = torch.Generator().manual_seed(seed_for("dist-gates-batch"))
    keep = torch.randperm(len(ids_l), generator=g)[:BATCH].numpy()
    ds = CocoImages("val", ids_l[keep])
    x = torch.stack([ds[i][0] for i in range(len(ds))]).to(DEV)

    def rows(t_ids, T_):
        pos = {int(v): i for i, v in enumerate(t_ids)}
        return T_[torch.tensor([pos[int(v)] for v in ids_l[keep]])].to(DEV)

    zt_l, zt_s, zt_g = rows(ids_l, t_l), rows(ids_s, t_s), rows(ids_g, t_g)
    model = build_student("gates", 0, 512)
    torch.manual_seed(seed_for("dist:gates-head768"))
    head768 = nn.Linear(240, 768).to(DEV)
    gp = torch.Generator().manual_seed(seed_for("dist-proj16-512"))
    P16 = torch.randn(512, 16, generator=gp).to(DEV)

    f = model.forward_features(x)
    zs512 = F.normalize(model.head(f), dim=-1)
    zs768 = F.normalize(head768(f), dim=-1)
    L = {"feature_mse": prim_sq(zs512, zt_l).mean(),
         "infonce": a7_grid_infonce(zs512, zt_l, temp=0.07),
         "siglip_pairwise": siglip_pairwise_loss(zs768, zt_s),
         "affinity_kl": affinity_kl_loss(zs512, zt_l, temp=0.07),
         "blueprint": (a7_grid_infonce(zs512, zt_l, temp=0.07)
                       + 0.3 * procrustes_sq(zs512, zt_l)
                       + _diff_cv(zs512 @ P16)),
         "consensus_gpa": prim_sq(zs512, zt_g).mean()}
    params = list(model.parameters()) + list(head768.parameters())
    names = list(L)
    print(f"=== DIST GATES: collinearity novelty 1-|cos(grad_i,grad_j)| "
          f"(init state, fixed val batch n={BATCH}) ===")
    print("losses: " + " ".join(f"{k}={float(L[k]):.4f}" for k in names))
    nov = {}
    for i, a in enumerate(names):
        for b in names[i + 1:]:
            nov[(a, b)] = collinearity_novelty(L[a], L[b], params)
    base = "feature_mse"
    print("novelty vs feature_mse (base): "
          + " | ".join(f"{b} {nov[(base, b)]:.4f}" for b in names[1:]))
    col = {"feature_mse": "featmse", "infonce": "infonce",
           "siglip_pairwise": "siglip", "affinity_kl": "affkl",
           "blueprint": "blprnt", "consensus_gpa": "gpa"}
    print(f"{'':<16}" + "".join(f"{col[b]:>9}" for b in names))
    for i, a in enumerate(names):
        cells = []
        for j, b in enumerate(names):
            if i == j:
                cells.append(f"{'--':>9}")
            else:
                key = (a, b) if (a, b) in nov else (b, a)
                cells.append(f"{nov[key]:>9.4f}")
        print(f"{col[a]:<16}" + "".join(cells))
    return nov


# ==================================================================== SMOKE
def smoke_summary(rows):
    print("\n=== DIST SMOKE TABLE (25 steps on the val2017 join; tiny eval: "
          "CIFAR 1000, retrieval 500) ===")
    hdr = (f"{'arm':<16}{'loss1':>9}{'loss25':>9}{'s/step':>8}{'peakGB':>8}"
           f"{'zs@1k':>8}{'r@1':>8}{'r@5':>8}{'agree':>8}{'cv16':>7}"
           f"{'erank':>8}  verdict")
    print(hdr)
    for r in rows:
        if r is None:
            continue
        l1, l2 = r.get("loss_step1", float("nan")), r.get("loss_final",
                                                          float("nan"))
        if "loss_final" not in r:            # eval-only control arms
            verdict = "PASS" if r["peak_gb"] < 18.0 else "FAIL"
        else:
            ok = (math.isfinite(l2) and r["peak_gb"] < 18.0)
            fell = math.isfinite(l1) and math.isfinite(l2) and l2 < l1
            verdict = "PASS" if ok and fell else ("WARN" if ok else "FAIL")
        print(f"{r['arm']:<16}{l1:>9.4f}{l2:>9.4f}{r['s_per_step']:>8.3f}"
              f"{r['peak_gb']:>8.2f}{r['zs_cifar10']:>8.4f}"
              f"{r['coco_r1']:>8.4f}{r['coco_r5']:>8.4f}"
              f"{r['agree_cos']:>8.4f}{r['cv16']:>7.3f}{r['erank']:>8.1f}"
              f"  {verdict}")


def print_launch_matrix():
    print("\n=== OVERNIGHT LAUNCH MATRIX (after train2017 lands; sequential; "
          "~12000 steps ~ 26 epochs of train2017 at batch 256) ===")
    for arm in TRAIN_ARMS:
        for seed in (0, 1):
            print(f"python tools/dist_bed.py --arm {arm} --seed {seed} "
                  f"--steps 12000 --split train")
    print("python tools/dist_bed.py --arm zs_floor --seed 0")
    print("python tools/dist_bed.py --arm teacher_ceiling --seed 0")


# ===================================================================== MAIN
if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", default=None)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--steps", type=int, default=4000)
    ap.add_argument("--split", default="val",
                    choices=("val", "train", "cc12m"))
    ap.add_argument("--workers", type=int, default=6)
    ap.add_argument("--smoke", action="store_true")
    ap.add_argument("--smoke-all", action="store_true", dest="smoke_all")
    ap.add_argument("--gates", action="store_true")
    ap.add_argument("--bench", action="store_true")
    ap.add_argument("--list", action="store_true")
    ap.add_argument("--eval-only", default=None, dest="eval_only")
    ap.add_argument("--extra-evals", action="store_true", dest="extra")
    a, _ = ap.parse_known_args()
    EXTRA_EVALS = a.extra

    if a.list:
        print("arms:", " ".join(ARMS))
        sys.exit(0)
    if a.gates:
        run_gates()
        sys.exit(0)
    if a.smoke_all:
        rows = [train_arm(arm, a.seed, 25, "val", smoke=True,
                          workers=a.workers) for arm in TRAIN_ARMS]
        smoke_summary(rows)
        print_launch_matrix()
        sys.exit(0)
    if a.smoke:
        rec = train_arm(a.arm or "feature_mse", a.seed, 25, "val",
                        smoke=True, workers=a.workers)
        smoke_summary([rec])
        print_launch_matrix()
        sys.exit(0)
    if a.eval_only:
        ck = torch.load(a.eval_only, map_location="cpu", weights_only=True)
        torch.cuda.set_per_process_memory_fraction(0.73)
        torch.cuda.reset_peak_memory_stats()
        t0 = time.time()
        model = build_student(ck["arm"], ck["seed"], ck["out_dim"])
        model.load_state_dict(ck["state_dict"])
        tower = "siglip" if ck["arm"] == "siglip_pairwise" else "laion"
        rec = {"arm": ck["arm"], "seed": ck["seed"], "steps": ck["steps"],
               "split": ck.get("split", "?"), "eval_only": True,
               **evaluate_student(ck["arm"], model, tower, False, a.workers),
               "wall_s": round(time.time() - t0, 1)}
        out = _ledger(rec, False)
        print(f"[DONE eval-only {ck['arm']} s{ck['seed']}] "
              f"zs {rec['zs_cifar10']:.4f} r@1 {rec['coco_r1']:.4f} -> {out}")
        sys.exit(0)
    if a.bench:
        train_arm(a.arm or "feature_mse", a.seed, 10, a.split, smoke=True,
                  workers=a.workers, bench=True)
        sys.exit(0)
    if not a.arm:
        print("arms:", " ".join(ARMS))
        sys.exit(0)
    train_arm(a.arm, a.seed, a.steps, a.split, workers=a.workers)