File size: 4,665 Bytes
7b3a667 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """Test-time loop scaling (TRM host). Train depth is L_cycles=n_backwards_L=6;
this evaluates exact-puzzle accuracy when the loop depth is scaled up at test time
(6, 9, 12, ...). The thesis test: does loop-attn keep using extra loop depth
(accuracy rises / holds) while the carry-last baseline saturates or degrades?
Loads EMA weights. Scales BOTH L_cycles (warmup cycles) and n_backwards_L (the
grad-cycle length, also used at eval) on the inner config; H_cycles and
halt_max_steps are held fixed so only the per-cycle loop depth changes.
Usage:
DISABLE_COMPILE=1 python3 eval_loopscale.py checkpoints/trm-lar-s0 \
--depths 6,9,12,18,24 --max_batches 20
"""
import os, sys, argparse, glob
os.environ.setdefault("DISABLE_COMPILE", "1")
os.environ.setdefault("WANDB_MODE", "offline")
import torch
from omegaconf import OmegaConf
from pretrain_config import PretrainConfig
from create_model import create_model
from pretrain import create_dataloader, autocast_ctx
from models.losses import IGNORE_LABEL_ID
def main():
ap = argparse.ArgumentParser()
ap.add_argument("ckpt_dir")
ap.add_argument("--depths", default="6,9,12,18,24")
ap.add_argument("--step", default="")
ap.add_argument("--max_batches", type=int, default=20)
args = ap.parse_args()
cfg_d = OmegaConf.to_container(OmegaConf.load(os.path.join(args.ckpt_dir, "all_config.yaml")), resolve=True)
cfg_d["load_checkpoint"] = None
cfg_d["resume_from"] = None
cfg_d["metrics_out"] = None
base = PretrainConfig(**cfg_d)
eval_loader, eval_metadata = create_dataloader(
base, "test", test_set_mode=True, epochs_per_iter=1,
global_batch_size=base.global_batch_size, rank=0, world_size=1)
model, _, _ = create_model(base, eval_metadata, rank=0, world_size=1, strict_load=False)
strip = lambda sd: {k.replace("_orig_mod.", ""): v for k, v in sd.items()}
bundles = sorted(glob.glob(os.path.join(args.ckpt_dir, "step_*_train_state.pt")),
key=lambda p: int(p.split("step_")[1].split("_")[0]))
bundle = (os.path.join(args.ckpt_dir, f"step_{args.step}_train_state.pt") if args.step
else (bundles[-1] if bundles else None))
if bundle and os.path.exists(bundle):
# EMA-eval bundle: EMA tracks only the dense params (16 keys); puzzle_emb/H_init/
# L_init live in d["model"] only. Load the full live state first, then overlay
# the EMA shadows (= the eval-time weights training uses).
d = torch.load(bundle, map_location="cpu", weights_only=False)
model.load_state_dict(strip(d["model"]), strict=False)
model.load_state_dict(strip(d["ema"]), strict=False)
src = os.path.basename(bundle) + " (EMA)"
else:
# Raw intermediate checkpoint (checkpoint_every_n_steps) — a flat model
# state_dict, no EMA shadow. Load it directly.
raw = os.path.join(args.ckpt_dir, f"step_{args.step}")
d = torch.load(raw, map_location="cpu", weights_only=False)
model.load_state_dict(strip(d), strict=False)
src = os.path.basename(raw) + " (raw)"
model.eval()
icfg = model.model.inner.config
has_la = model.model.inner.loop_attn is not None
print(f"{args.ckpt_dir}: loop_attn={has_la} ckpt={src} "
f"train depth L_cycles={icfg.L_cycles} n_backwards_L={icfg.n_backwards_L}")
depths = [int(x) for x in args.depths.split(",") if x]
print(f"\n depth | exact-puzzle acc (max_batches={args.max_batches})")
for D in depths:
icfg.L_cycles = D
icfg.n_backwards_L = D
exact_n = exact_tot = 0
with torch.no_grad():
for bi, (set_name, batch, _) in enumerate(eval_loader):
if bi >= args.max_batches:
break
batch = {k: v.cuda() for k, v in batch.items()}
with torch.device("cuda"):
carry = model.initial_carry(batch)
while True:
with autocast_ctx(base):
carry, _, _, _, preds, all_finish = model(carry=carry, batch=batch, return_keys=["preds"])
if all_finish:
break
p, lab = preds["preds"], batch["labels"]
valid = lab != IGNORE_LABEL_ID
cor = valid & (p == lab)
vs, cs = valid.sum(1), cor.sum(1)
sv = vs > 0
exact_n += int((sv & (cs == vs)).sum().item())
exact_tot += int(sv.sum().item())
acc = exact_n / max(exact_tot, 1)
print(f" {D:5d} | {acc:.4f} (n={exact_tot})")
if __name__ == "__main__":
main()
|