training_code/latent_threads/train_single.py

"""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()