diff options
Diffstat (limited to 'research/flossing/diagnose_trm_joint.py')
| -rw-r--r-- | research/flossing/diagnose_trm_joint.py | 225 |
1 files changed, 225 insertions, 0 deletions
diff --git a/research/flossing/diagnose_trm_joint.py b/research/flossing/diagnose_trm_joint.py new file mode 100644 index 0000000..160a7cb --- /dev/null +++ b/research/flossing/diagnose_trm_joint.py @@ -0,0 +1,225 @@ +"""TRM Sudoku joint Lyapunov diagnostic — TRM version of diagnose_hrm_joint.py. + +Key differences from HRM: +- TRM has ONE shared L_level (H_layers config is "ignored") +- z_L update: z_L = L_level(z_L, z_H + input_embeddings) +- z_H update: z_H = L_level(z_H, z_L) ← same L_level! +- H_cycles=3, L_cycles=6 (vs HRM 2,2) + +Joint tangent block structure: +- L step: v_L_new = J · (v_L + v_H), v_H_new = v_H, J at (z_L + z_H + ie) +- H step: v_H_new = J' · (v_H + v_L), v_L_new = v_L, J' at (z_H + z_L) +J and J' share weights but evaluated at different points. +""" +from __future__ import annotations +import sys, os, yaml, math, argparse, json, time +from pathlib import Path +import numpy as np +import torch + +TRM_DIR = Path("/home/yurenh2/rrm/trm") +sys.path.insert(0, str(TRM_DIR)) + +from models.recursive_reasoning.trm import TinyRecursiveReasoningModel_ACTV1 + + +def load_model(ckpt_root: Path, ckpt_name: str, device: str): + cfg = yaml.safe_load((ckpt_root / "all_config.yaml").read_text()) + arch_cfg = dict(cfg["arch"]) + train_meta = json.loads((Path(cfg["data_paths"][0]) / "train" / "dataset.json").read_text()) + arch_cfg.update(batch_size=cfg["global_batch_size"], seq_len=train_meta["seq_len"], + vocab_size=train_meta["vocab_size"], + num_puzzle_identifiers=train_meta["num_puzzle_identifiers"]) + model = TinyRecursiveReasoningModel_ACTV1(arch_cfg) + sd = torch.load(ckpt_root / ckpt_name, map_location="cpu", weights_only=True) + stripped = {k.replace("_orig_mod.", "").replace("model.", ""): v for k, v in sd.items()} + missing, unexpected = model.load_state_dict(stripped, strict=False) + print(f"[load] missing={len(missing)} unexpected={len(unexpected)}") + if missing[:3]: print(f" sample missing: {missing[:3]}") + if unexpected[:3]: print(f" sample unexpected: {unexpected[:3]}") + model.to(device).eval() + return model, cfg, train_meta + + +def load_test_samples(data_path, n_total, shard_id, num_shards, seed): + rng = np.random.default_rng(seed) + inputs = np.load(data_path / "test" / "all__inputs.npy") + labels = np.load(data_path / "test" / "all__labels.npy") + pid = np.load(data_path / "test" / "all__puzzle_identifiers.npy") + all_idx = rng.choice(len(inputs), size=n_total, replace=False) + shard_size = (n_total + num_shards - 1) // num_shards + s, e = shard_id * shard_size, min((shard_id + 1) * shard_size, n_total) + idx = all_idx[s:e] + return { + "inputs": torch.from_numpy(inputs[idx].astype(np.int32)), + "labels": torch.from_numpy(labels[idx].astype(np.int32)), + "puzzle_identifiers": torch.from_numpy(pid[idx].astype(np.int32)), + "idx": idx, + } + + +def jvp_through(f, x, v): + return torch.autograd.functional.jvp(f, x, v=v, create_graph=False, strict=False) + + +def run_diagnose_batch(model, batch, device, k_lyap, t_ons, seed): + inner = model.inner + cfg = inner.config + B = batch["inputs"].shape[0] + seq_full = cfg.seq_len + inner.puzzle_emb_len + hidden = cfg.hidden_size + D = seq_full * hidden + + z_H = inner.H_init.unsqueeze(0).expand(B, seq_full, hidden).clone().to(inner.forward_dtype) + z_L = inner.L_init.unsqueeze(0).expand(B, seq_full, hidden).clone().to(inner.forward_dtype) + seq_info = dict(cos_sin=inner.rotary_emb() if hasattr(inner, "rotary_emb") else None) + input_embeddings = inner._input_embeddings(batch["inputs"].to(device), + batch["puzzle_identifiers"].to(device)) + + g = torch.Generator(device=device).manual_seed(seed) + Q0 = torch.randn(B, 2*D, k_lyap, device=device, dtype=torch.float32, generator=g) + Q, _ = torch.linalg.qr(Q0) + log_R_sum = torch.zeros(B, k_lyap, device=device, dtype=torch.float32) + n_lyap_steps = 0 + step_counter = 0 + + drift_zH_per_step, drift_zL_per_step = [], [] + halted_at = torch.zeros(B, dtype=torch.long, device=device) + q_halt_hist, q_continue_hist = [], [] + + for act_step in range(cfg.halt_max_steps): + z_H_prev = z_H.detach().clone() + z_L_prev = z_L.detach().clone() + + with torch.enable_grad(): + zH, zL = z_H.detach(), z_L.detach() + for _h in range(cfg.H_cycles): + # L cycles + for _l in range(cfg.L_cycles): + v_H_j = Q[:, :D, :] + v_L_j = Q[:, D:, :] + v_comb = v_H_j + v_L_j + new_v_L_cols = [] + f_L = lambda z: inner.L_level(z, zH + input_embeddings, **seq_info) + for i in range(k_lyap): + v_i = v_comb[:, :, i].reshape(B, seq_full, hidden).to(inner.forward_dtype) + zL_new, Dv = jvp_through(f_L, zL, v_i) + new_v_L_cols.append(Dv.reshape(B, D).to(torch.float32)) + new_v_L = torch.stack(new_v_L_cols, dim=-1) + Q = torch.cat([v_H_j, new_v_L], dim=1) + zL = zL_new + step_counter += 1 + if step_counter % t_ons == 0: + Q, R = torch.linalg.qr(Q) + log_R_sum = log_R_sum + R.diagonal(dim1=-2, dim2=-1).abs().clamp_min(1e-30).log() + n_lyap_steps += 1 + + # H step (uses SAME L_level!) + v_H_j = Q[:, :D, :] + v_L_j = Q[:, D:, :] + v_comb = v_H_j + v_L_j + new_v_H_cols = [] + f_H = lambda z: inner.L_level(z, zL, **seq_info) + for i in range(k_lyap): + v_i = v_comb[:, :, i].reshape(B, seq_full, hidden).to(inner.forward_dtype) + zH_new, Dv = jvp_through(f_H, zH, v_i) + new_v_H_cols.append(Dv.reshape(B, D).to(torch.float32)) + new_v_H = torch.stack(new_v_H_cols, dim=-1) + Q = torch.cat([new_v_H, v_L_j], dim=1) + zH = zH_new + step_counter += 1 + if step_counter % t_ons == 0: + Q, R = torch.linalg.qr(Q) + log_R_sum = log_R_sum + R.diagonal(dim1=-2, dim2=-1).abs().clamp_min(1e-30).log() + n_lyap_steps += 1 + + z_H, z_L = zH, zL + + drift_zH_per_step.append((z_H - z_H_prev).float().flatten(1).norm(dim=1).cpu()) + drift_zL_per_step.append((z_L - z_L_prev).float().flatten(1).norm(dim=1).cpu()) + + with torch.no_grad(): + q_logits = inner.q_head(z_H[:, 0]).float() + q_halt, q_continue = q_logits[..., 0], q_logits[..., 1] + q_halt_hist.append(q_halt.cpu()); q_continue_hist.append(q_continue.cpu()) + new_halt = (q_halt > q_continue) & (halted_at == 0) + halted_at[new_halt] = act_step + 1 + output = inner.lm_head(z_H)[:, inner.puzzle_emb_len:].float() + final_logits = output + + lyap_spec = (log_R_sum / max(n_lyap_steps, 1)).cpu().numpy() + + with torch.no_grad(): + preds = final_logits.argmax(dim=-1) + labels = batch["labels"].to(device) + mask = labels > 0 + exact = ((preds == labels) | ~mask).all(dim=-1).cpu().float() + token_acc = ((preds == labels) & mask).sum(-1).float() / mask.sum(-1).float().clamp_min(1) + token_acc = token_acc.cpu() + + return { + "drift_zH": torch.stack(drift_zH_per_step, dim=1).numpy(), + "drift_zL": torch.stack(drift_zL_per_step, dim=1).numpy(), + "halted_at": halted_at.cpu().numpy(), + "q_halt": torch.stack(q_halt_hist, dim=1).numpy(), + "q_continue": torch.stack(q_continue_hist, dim=1).numpy(), + "lyap_spec": lyap_spec, + "exact_correct": exact.numpy(), + "token_acc": token_acc.numpy(), + } + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--ckpt-root", required=True) + ap.add_argument("--ckpt-name", default="step_13020") + ap.add_argument("--n-samples", type=int, default=512) + ap.add_argument("--shard-id", type=int, default=0) + ap.add_argument("--num-shards", type=int, default=1) + ap.add_argument("--batch-size", type=int, default=16) + ap.add_argument("--k-lyap", type=int, default=8) + ap.add_argument("--t-ons", type=int, default=1) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--out", default="diag_trm.npz") + args = ap.parse_args() + + device = "cuda" + model, cfg, train_meta = load_model(Path(args.ckpt_root), args.ckpt_name, device) + print(f"loaded {args.ckpt_name}: hidden={model.inner.config.hidden_size}, " + f"seq_full={train_meta['seq_len'] + model.inner.puzzle_emb_len}, " + f"halt_max_steps={model.inner.config.halt_max_steps}, " + f"H={model.inner.config.H_cycles} L={model.inner.config.L_cycles}") + + test = load_test_samples(Path(cfg["data_paths"][0]), args.n_samples, args.shard_id, args.num_shards, args.seed) + n = len(test["inputs"]) + print(f"shard {args.shard_id}/{args.num_shards}: {n} samples") + + res = {k: [] for k in ["drift_zH","drift_zL","halted_at","q_halt","q_continue","lyap_spec","exact_correct","token_acc","idx"]} + t0 = time.time() + for s in range(0, n, args.batch_size): + e = min(s + args.batch_size, n) + batch = {k: test[k][s:e].to(device) for k in ["inputs","labels","puzzle_identifiers"]} + out = run_diagnose_batch(model, batch, device, args.k_lyap, args.t_ons, args.seed + s) + for k, v in out.items(): + res[k].append(v) + res["idx"].append(test["idx"][s:e]) + ls = out["lyap_spec"] + print(f" [{e}/{n}] dt={time.time()-t0:.1f}s exact={out['exact_correct'].mean():.3f} " + f"λ_1={ls[:,0].mean():+.4f} λ_{args.k_lyap}={ls[:,-1].mean():+.4f}", flush=True) + + saved = {} + for k, v in res.items(): + if not v: continue + try: saved[k] = np.concatenate(v, 0) + except ValueError: saved[k] = np.stack(v, 0) + np.savez_compressed(args.out, **saved) + succ = saved["exact_correct"] > 0.5 + print(f"\nN={len(succ)} acc={succ.mean():.4f}") + print(f"{'i':>3} {'all':>10} {'succ':>10} {'fail':>10} {'Δ':>9}") + for i in range(saved["lyap_spec"].shape[1]): + li = saved["lyap_spec"][:, i] + print(f"{i+1:>3} {li.mean():+10.4f} {li[succ].mean():+10.4f} {li[~succ].mean():+10.4f} {li[~succ].mean()-li[succ].mean():+9.4f}") + + +if __name__ == "__main__": + main() |