"""Train a SINGLE-token-per-step latent-CoT organism (one z_t carries the whole K-cell row c_t). Loss = answer CE + w_emb * per-cell feedback CE, with teacher-forcing prob annealed 1 -> 0 (scheduled sampling). Trains over a MIX of chain lengths (train_lengths) so the single per-step recurrence generalises to longer chains than seen. See latent_threads/single.py. python -m latent_threads.train_single --config latent_threads/configs/single_k3m6.json --batch-id sg1 """ import os try: import dotenv dotenv.load_dotenv() except Exception: pass _USER = os.environ.get("USER", "jbauer") os.environ.setdefault("HF_HOME", f"/workspace-vast/{_USER}/hf") os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1") os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import argparse import json import random import time import torch import wandb from peft import LoraConfig, get_peft_model from transformers import AutoModelForCausalLM, AutoTokenizer from latent_threads import tasks as LT from latent_threads.single import single_forward, single_readout_acc, single_state_acc def main(): ap = argparse.ArgumentParser() ap.add_argument("--config", required=True) ap.add_argument("--batch-id", default="sg1") ap.add_argument("--wandb-mode", default=None, choices=["online", "offline", "disabled"]) ap.add_argument("--smoke", action="store_true") args = ap.parse_args() if args.wandb_mode: os.environ["WANDB_MODE"] = args.wandb_mode cfg = json.load(open(args.config)) device = "cuda" K = cfg["task_kwargs"]["k"] base_task = LT.make_task(cfg["task"], **cfg["task_kwargs"]) # eval / mastery length train_lengths = cfg.get("train_lengths", [base_task.m]) eval_lengths = cfg.get("eval_lengths", sorted(set(train_lengths + [base_task.m]))) gen_lengths = cfg.get("gen_lengths", [base_task.m + 2, base_task.m + 4]) # longer-than-trained rng = random.Random(cfg["seed"]); erng = random.Random(cfg["seed"] + 777) grng = torch.Generator(device="cpu").manual_seed(cfg["seed"] + 99) torch.manual_seed(cfg["seed"]) tok = AutoTokenizer.from_pretrained(cfg["model_name"]) tok.pad_token_id = tok.pad_token_id or tok.eos_token_id model = AutoModelForCausalLM.from_pretrained(cfg["model_name"], torch_dtype=torch.bfloat16, attn_implementation="eager", device_map="cuda") model = get_peft_model(model, LoraConfig(r=cfg["lora_r"], lora_alpha=cfg["lora_alpha"], lora_dropout=cfg["lora_dropout"], bias="none", task_type="CAUSAL_LM", target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"])) model.config.use_cache = False model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) model.train() d = model.config.hidden_size # head: residual -> K*10 per-cell digit logits ; codebook: the soft-token embedding matrix [K,10,d] head = torch.nn.Sequential(torch.nn.LayerNorm(d), torch.nn.Linear(d, d // 4), torch.nn.GELU(), torch.nn.Linear(d // 4, K * 10)).to(device, torch.bfloat16) dig_E = model.get_input_embeddings().weight[torch.tensor(LT.digit_ids(tok), device=device)].detach() # SEPARABLE in-distribution codebook: K*10 DISTINCT single-token embeddings so the K cells occupy # distinguishable directions from step 1. (Initialising every cell near the same digit embeddings # makes the fed soft token a near-symmetric digit-SUM, which is positionally unrecoverable -> the # per-cell state stays at chance; cf. the stalled sg1 run.) import string _pool, _syms = list(string.digits + string.ascii_letters + string.punctuation), [] for s in _pool: t = tok(s, add_special_tokens=False)["input_ids"] if len(t) == 1: _syms.append(t[0]) if len(_syms) == K * 10: break assert len(_syms) == K * 10, f"need {K*10} single-token symbols, got {len(_syms)}" codebook = torch.nn.Parameter(model.get_input_embeddings().weight[torch.tensor(_syms, device=device)] .detach().view(K, 10, d).clone().to(torch.bfloat16)) q_emb = torch.nn.Parameter(dig_E[0].detach().clone().to(torch.bfloat16)) im_end = tok.convert_tokens_to_ids("<|im_end|>") save_root = os.path.join(cfg["save_dir"], args.batch_id, base_task.name) os.makedirs(save_root, exist_ok=True) if not args.smoke: wandb.init(project=cfg["wandb_project"], entity=cfg.get("wandb_entity") or None, name=f"{args.batch_id}_{base_task.name}_single", group=args.batch_id, config={**cfg, "single": True, "train_lengths": train_lengths}) wandb.define_metric("train/examples_seen") wandb.define_metric("*", step_metric="train/examples_seen") def save_ckpt(tag): dd = os.path.join(save_root, tag); os.makedirs(dd, exist_ok=True) model.save_pretrained(dd); tok.save_pretrained(dd) torch.save({"head": head.state_dict(), "codebook": codebook.detach().cpu(), "q_emb": q_emb.detach().cpu()}, os.path.join(dd, "single_extra.pt")) json.dump({**cfg, "tag": tag, "batch_id": args.batch_id, "single": True, "train_lengths": train_lengths}, open(os.path.join(dd, "lt_cfg.json"), "w"), indent=2) print(f" [ckpt] {dd}", flush=True) params = ([p for p in model.parameters() if p.requires_grad] + list(head.parameters()) + [codebook, q_emb]) opt = torch.optim.AdamW(params, lr=cfg["lr"]) B = cfg["batch"]; w_emb = cfg.get("emb_supervision_weight", 1.0); tf_anneal = cfg.get("tf_anneal_steps", 2500) def ev(task, n, tf_prob=0.0): return single_readout_acc(model, head, codebook, q_emb, tok, task, [task.sample(erng) for _ in range(n)], im_end, device, tf_prob=tf_prob, gt_rng=grng if tf_prob > 0 else None) print(f"[sg] single-token organism K={K} curriculum={cfg.get('curriculum', False)} " f"start={cfg.get('curriculum_start', base_task.m)} target_m={base_task.m} aux_w={cfg.get('emb_supervision_weight',1.0)} " f"chance={base_task.chance:.2f} PRE readout={ev(base_task, 64):.3f}", flush=True) use_curric = cfg.get("curriculum", False) cur_max = cfg.get("curriculum_start", base_task.m) if use_curric else base_task.m grow_thresh = cfg.get("grow_thresh", 0.85) best, streak, seen = 0.0, 0, 0 for gstep in range(1, cfg["max_steps"] + 1): t0 = time.time() tf_p = max(0.0, 1.0 - gstep / tf_anneal) lengths_now = list(range(2, cur_max + 1)) if use_curric else train_lengths m = rng.choice(lengths_now) btask = LT.make_task(cfg["task"], k=K, m=m) probs = [btask.sample(rng) for _ in range(B)] _, _, (aux_logits, aux_gt) = single_forward( model, head, codebook, q_emb, tok, btask, probs, im_end, device, with_answer=False, tf_prob=tf_p, gt_rng=grng) # loss = per-cell state CE only; the answer is read from this same head at z_T (no LM-head # bypass). aux covers ALL cells at ALL steps, including the queried cell at the final step. aux = torch.nn.functional.cross_entropy(aux_logits.reshape(-1, 10).float(), aux_gt.reshape(-1)) loss = w_emb * aux loss.backward() torch.nn.utils.clip_grad_norm_(params, 1.0) opt.step(); opt.zero_grad(); seen += B if gstep % 10 == 0 or gstep == 1: print(f" step {gstep} m={m} aux={aux.item():.4f} tf={tf_p:.2f} t={time.time()-t0:.2f}s", flush=True) if args.smoke and gstep >= 3: print(f"[sg] SMOKE OK: readout={ev(base_task, 32):.3f} state={single_state_acc(model, head, codebook, q_emb, tok, base_task, [base_task.sample(erng) for _ in range(32)], im_end, device):.3f}", flush=True) return if gstep % cfg["eval_every"] == 0: ctask = LT.make_task(cfg["task"], k=K, m=cur_max) acc_cur = ev(ctask, cfg["mastery_eval_n"]) # free-running @ current curriculum length if use_curric and cur_max < base_task.m and acc_cur >= grow_thresh: cur_max += 1; streak = 0 print(f" [sg] GROW curriculum -> cur_max={cur_max} (step {gstep})", flush=True) acc = ev(base_task, cfg["mastery_eval_n"]) # free-running @ target length gen = {f"gen_m{gm}": ev(LT.make_task(cfg["task"], k=K, m=gm), 64) for gm in gen_lengths} state = single_state_acc(model, head, codebook, q_emb, tok, ctask, [ctask.sample(erng) for _ in range(64)], im_end, device) print(f" [eval] step {gstep} cur_max={cur_max} readout@cur={acc_cur:.3f} readout@{base_task.m}={acc:.3f} " f"state@cur={state:.3f} gen={ {k: round(v,3) for k,v in gen.items()} } (tf_p={tf_p:.2f}, best {best:.3f})", flush=True) if not args.smoke: wandb.log({"train/examples_seen": seen, "eval/readout_cur": acc_cur, "eval/cur_max": cur_max, "eval/readout_acc": acc, "eval/state_acc": state, "train/tf_prob": tf_p, "train/aux": float(aux.item()), **{f"eval/{k}": v for k, v in gen.items()}}) if acc > best: best = acc; save_ckpt("best") streak = streak + 1 if (cur_max == base_task.m and acc >= cfg["mastery_threshold"] and tf_p == 0.0) else 0 if streak >= 2: save_ckpt("final"); print(f"[sg] DONE: acc>={cfg['mastery_threshold']} twice at tf=0, cur_max={cur_max} (step {gstep}).", flush=True); break else: print(f"[sg] max_steps reached; best={best:.3f}.", flush=True); save_ckpt("last") if not args.smoke: wandb.finish() if __name__ == "__main__": main()