diff options
Diffstat (limited to 'research/flossing/step6_prefloss.py')
| -rw-r--r-- | research/flossing/step6_prefloss.py | 310 |
1 files changed, 310 insertions, 0 deletions
diff --git a/research/flossing/step6_prefloss.py b/research/flossing/step6_prefloss.py new file mode 100644 index 0000000..af84de0 --- /dev/null +++ b/research/flossing/step6_prefloss.py @@ -0,0 +1,310 @@ +"""Step 6: Preflossing experiment (Engelken-style separate phases). + +Phase 1 — Pure flossing: only optimize L_floss, no task loss. + engelken: L = (1/k) Σ λ_i² (push all top-k toward 0, two-sided) + cf: L = (1/k) Σ max(0, λ_i)² (only push positive λ toward 0) + +Phase 2 — Pure task training: standard HRM ACT loss, no flossing. + +Baseline mode (--prefloss-steps 0): skip phase 1, go straight to task training. + +Key fix vs step3: lyap_act_steps defaults to halt_max_steps (16 for HRM). +""" +from __future__ import annotations +import sys, os, yaml, json, math, time, argparse +from pathlib import Path +import numpy as np +import torch +import torch.nn.functional as F + +HRM_DIR = Path("/home/yurenh2/rrm/hrm") +sys.path.insert(0, str(HRM_DIR)) + +from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1 +from models.losses import ACTLossHead +from adam_atan2 import AdamATan2 + + +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_path"]) / "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"], causal=False) + base = HierarchicalReasoningModel_ACTV1(arch_cfg) + head = ACTLossHead(base, loss_type=arch_cfg["loss"]["loss_type"]) + sd = torch.load(ckpt_root / ckpt_name, map_location="cpu", weights_only=True) + stripped = {k.replace("_orig_mod.", ""): v for k, v in sd.items()} + missing, unexpected = head.load_state_dict(stripped, strict=False) + print(f"[load {ckpt_name}] missing={len(missing)} unexpected={len(unexpected)}") + head.to(device) + return head, base, cfg, train_meta + + +def jvp_train(f, x, v): + return torch.autograd.functional.jvp(f, x, v=v, create_graph=True, strict=False) + + +def compute_joint_lyap_spec(base, batch, k_lyap, lyap_act_steps, device, seed): + inner = base.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"], batch["puzzle_identifiers"]) + + 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_steps = 0 + + n_act = min(lyap_act_steps, cfg.halt_max_steps) + for _act in range(n_act): + for _h in range(cfg.H_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, z_H + input_embeddings, **seq_info) + for i in range(k_lyap): + v_i = v_comb[:, :, i].reshape(B, seq_full, hidden).to(inner.forward_dtype) + z_L_new, Dv = jvp_train(f_L, z_L, 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) + z_L = z_L_new + 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_steps += 1 + 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.H_level(z, z_L, **seq_info) + for i in range(k_lyap): + v_i = v_comb[:, :, i].reshape(B, seq_full, hidden).to(inner.forward_dtype) + z_H_new, Dv = jvp_train(f_H, z_H, 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) + z_H = z_H_new + 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_steps += 1 + + return log_R_sum / max(n_steps, 1) + + +def load_train_batches(data_path: Path, batch_size: int, n_iters: int, seed: int = 0): + rng = np.random.default_rng(seed) + inputs = np.load(data_path / "train" / "all__inputs.npy") + labels = np.load(data_path / "train" / "all__labels.npy") + pid = np.load(data_path / "train" / "all__puzzle_identifiers.npy") + N = len(inputs) + for _ in range(n_iters): + idx = rng.choice(N, size=batch_size, replace=False) + yield { + "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)), + } + + +def evaluate(head, base, data_path, n_samples, batch_size, device, seed=42): + 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") + idx_all = rng.choice(len(inputs), size=n_samples, replace=False) + head.eval() + correct = 0; token_correct = 0; token_total = 0 + for s in range(0, n_samples, batch_size): + e = min(s + batch_size, n_samples) + idx = idx_all[s:e] + batch = { + "inputs": torch.from_numpy(inputs[idx].astype(np.int32)).to(device), + "labels": torch.from_numpy(labels[idx].astype(np.int32)).to(device), + "puzzle_identifiers": torch.from_numpy(pid[idx].astype(np.int32)).to(device), + } + with torch.no_grad(): + with torch.device(device): + carry = base.initial_carry(batch) + for _ in range(base.config.halt_max_steps): + carry, outputs = base(carry=carry, batch=batch) + preds = outputs["logits"].argmax(dim=-1) + mask = batch["labels"] > 0 + exact = ((preds == batch["labels"]) | ~mask).all(dim=-1).float() + correct += exact.sum().item() + token_correct += ((preds == batch["labels"]) & mask).sum().item() + token_total += mask.sum().item() + return correct / n_samples, token_correct / max(token_total, 1) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--ckpt-root", required=True) + ap.add_argument("--ckpt-name", default="step_13020") + ap.add_argument("--prefloss-steps", type=int, default=500, + help="Phase 1: pure flossing steps. 0 = skip (baseline).") + ap.add_argument("--train-steps", type=int, default=3000, + help="Phase 2: pure task training steps.") + ap.add_argument("--floss-mode", choices=["engelken", "cf", "volume_cf"], default="engelken", + help="engelken: Σλ_i² (two-sided). cf: Σmax(0,λ_i)² (one-sided hinge). " + "volume_cf: max(0, mean_i λ_i)² over the measured top-k spectrum.") + ap.add_argument("--batch-size", type=int, default=8) + ap.add_argument("--floss-lr", type=float, default=1e-4, + help="LR for flossing phase (Engelken uses higher LR)") + ap.add_argument("--train-lr", type=float, default=1e-5, + help="LR for task training phase") + ap.add_argument("--k-lyap", type=int, default=2) + ap.add_argument("--lyap-act-steps", type=int, default=16, + help="ACT steps for Lyapunov computation (default=halt_max_steps)") + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--eval-every", type=int, default=100) + ap.add_argument("--eval-n", type=int, default=512) + ap.add_argument("--eval-batch-size", type=int, default=32) + ap.add_argument("--out", default="step6_log.json") + args = ap.parse_args() + + device = "cuda" + head, base, cfg, train_meta = load_model(Path(args.ckpt_root), args.ckpt_name, device) + data_path = Path(cfg["data_path"]) + + print(f"\n=== Initial eval (loaded {args.ckpt_name}) ===") + acc0, tacc0 = evaluate(head, base, data_path, args.eval_n, args.eval_batch_size, device) + print(f" initial: exact_acc = {acc0:.4f} token_acc = {tacc0:.4f}") + + log = {"args": vars(args), "initial_acc": acc0, "initial_tok_acc": tacc0, + "phase1_steps": [], "phase1_evals": [], "phase2_steps": [], "phase2_evals": []} + + global_step = 0 + + # ========== PHASE 1: Pure flossing ========== + if args.prefloss_steps > 0: + print(f"\n=== Phase 1: Pure {args.floss_mode} flossing ({args.prefloss_steps} steps, lr={args.floss_lr}) ===") + floss_optim = AdamATan2(head.parameters(), lr=args.floss_lr, betas=(0.9, 0.95), weight_decay=0.0) + floss_iter = load_train_batches(data_path, args.batch_size, args.prefloss_steps, seed=args.seed) + t0 = time.time() + + for step, batch in enumerate(floss_iter): + batch = {k: v.to(device) for k, v in batch.items()} + head.train() + base.inner.puzzle_emb.eval() + for p in base.inner.puzzle_emb.parameters(): + p.requires_grad_(False) + + lyap_spec = compute_joint_lyap_spec( + base, batch, k_lyap=args.k_lyap, + lyap_act_steps=args.lyap_act_steps, device=device, + seed=args.seed + step, + ) + + if args.floss_mode == "engelken": + floss_loss = (lyap_spec ** 2).mean() + elif args.floss_mode == "volume_cf": + floss_loss = (lyap_spec.mean(dim=1).clamp_min(0.0) ** 2).mean() + else: + floss_loss = (lyap_spec.clamp_min(0.0) ** 2).mean() + + floss_optim.zero_grad(set_to_none=True) + floss_loss.backward() + torch.nn.utils.clip_grad_norm_([p for p in head.parameters() if p.requires_grad], 1.0) + floss_optim.step() + + lyap1 = lyap_spec[:, 0].detach() + rec = { + "step": step, "floss_loss": float(floss_loss.item()), + "lyap1_mean": float(lyap1.mean().item()), + "lyap1_max": float(lyap1.max().item()), + "lyap_all_mean": float(lyap_spec.detach().mean().item()), + "lyap_volume_mean": float(lyap_spec.detach().mean(dim=1).mean().item()), + "lyap_volume_max": float(lyap_spec.detach().mean(dim=1).max().item()), + } + log["phase1_steps"].append(rec) + if step % 10 == 0 or step == args.prefloss_steps - 1: + print(f" P1[{step:>4}/{args.prefloss_steps}] dt={time.time()-t0:.1f}s " + f"floss={rec['floss_loss']:.6f} " + f"λ1={rec['lyap1_mean']:+.4f} max={rec['lyap1_max']:+.4f} " + f"λ_all={rec['lyap_all_mean']:+.4f}", flush=True) + + if (step + 1) % args.eval_every == 0: + acc, tacc = evaluate(head, base, data_path, args.eval_n, args.eval_batch_size, device) + print(f" >> P1 EVAL @ step {step+1}: exact_acc={acc:.4f} (Δ={acc-acc0:+.4f})", flush=True) + log["phase1_evals"].append({"step": step + 1, "acc": acc, "tok_acc": tacc}) + + acc_p1, tacc_p1 = evaluate(head, base, data_path, args.eval_n, args.eval_batch_size, device) + print(f" Phase 1 final: exact_acc={acc_p1:.4f} (Δ from init: {acc_p1-acc0:+.4f})") + log["phase1_final_acc"] = acc_p1 + log["phase1_evals"].append({"step": args.prefloss_steps, "acc": acc_p1, "tok_acc": tacc_p1}) + global_step = args.prefloss_steps + else: + print("\n=== Phase 1 skipped (baseline mode) ===") + log["phase1_final_acc"] = acc0 + + # ========== PHASE 2: Pure task training ========== + print(f"\n=== Phase 2: Pure task training ({args.train_steps} steps, lr={args.train_lr}) ===") + train_optim = AdamATan2(head.parameters(), lr=args.train_lr, betas=(0.9, 0.95), + weight_decay=cfg["weight_decay"]) + train_iter = load_train_batches(data_path, args.batch_size, args.train_steps, + seed=args.seed + 10000) + t0 = time.time() + acc_ref = log.get("phase1_final_acc", acc0) + + for step, batch in enumerate(train_iter): + batch = {k: v.to(device) for k, v in batch.items()} + head.train() + base.inner.puzzle_emb.eval() + for p in base.inner.puzzle_emb.parameters(): + p.requires_grad_(False) + + with torch.device(device): + carry = base.initial_carry(batch) + sup_loss_sum = 0.0 + n_loss = 0 + for _ in range(base.config.halt_max_steps): + carry, l, metrics, _, all_finish = head(return_keys=[], carry=carry, batch=batch) + sup_loss_sum = sup_loss_sum + l + n_loss += 1 + if all_finish: + break + sup_loss = sup_loss_sum / max(n_loss, 1) / args.batch_size + + train_optim.zero_grad(set_to_none=True) + sup_loss.backward() + torch.nn.utils.clip_grad_norm_([p for p in head.parameters() if p.requires_grad], 1.0) + train_optim.step() + + rec = {"step": step, "sup_loss": float(sup_loss.item())} + log["phase2_steps"].append(rec) + if step % 50 == 0 or step == args.train_steps - 1: + print(f" P2[{step:>4}/{args.train_steps}] dt={time.time()-t0:.1f}s " + f"sup={rec['sup_loss']:.4f}", flush=True) + + if (step + 1) % args.eval_every == 0: + acc, tacc = evaluate(head, base, data_path, args.eval_n, args.eval_batch_size, device) + print(f" >> P2 EVAL @ step {step+1}: exact_acc={acc:.4f} " + f"(Δ from init: {acc-acc0:+.4f}, Δ from P1: {acc-acc_ref:+.4f})", flush=True) + log["phase2_evals"].append({"step": global_step + step + 1, "acc": acc, "tok_acc": tacc}) + + acc_f, tacc_f = evaluate(head, base, data_path, args.eval_n, args.eval_batch_size, device) + print(f"\n=== Final eval ===") + print(f" initial: {acc0:.4f} phase1_end: {log.get('phase1_final_acc', acc0):.4f} " + f"final: {acc_f:.4f} (total Δ: {acc_f-acc0:+.4f})") + log["final_acc"] = acc_f + log["final_tok_acc"] = tacc_f + log["phase2_evals"].append({"step": global_step + args.train_steps, "acc": acc_f, "tok_acc": tacc_f}) + + Path(args.out).write_text(json.dumps(log, indent=2)) + print(f"log → {args.out}") + + +if __name__ == "__main__": + main() |