""" clankerDiffusion — training loop (hybrid AR / masked-diffusion). Hybrid objective (per step, mode chosen at random, p(AR)=0.5): AR (mode 0): causal LM cross-entropy over the whole window. DIFF (mode 1): MDLM absorbing-state masked diffusion -- mask each token independently with ratio r~U(0,1); reconstruct masked tokens with bidirectional attention, conditioned on r via time embed. Runs in bf16, AdamW + cosine LR, grad-clip, checkpoints locally and (optionally) pushes each checkpoint to a HuggingFace repo via the `hf` CLI. Used both locally and on Modal L4 (override --data-dir/--ckpt-dir/--hf-repo and the model dimensions for a bigger model). """ import os, json, time, argparse, subprocess, threading import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from model import YKDiff from tokenizer import YKTokenizer HERE = os.path.dirname(os.path.abspath(__file__)) DATADIR = os.path.join(HERE, "data") CKPTDIR = os.path.join(HERE, "checkpoints") os.makedirs(CKPTDIR, exist_ok=True) # default architecture tuned for 16 GB (RTX 5060 Ti) DEFAULT_CFG = dict( d_model=768, n_layers=12, n_heads=12, d_ff=2048, max_len=8192, vocab_size=32768, rope_scale=1.0, ) def build_cfg(args): cfg = dict(DEFAULT_CFG) for k in ("d_model", "n_layers", "n_heads", "d_ff", "vocab_size", "max_len"): v = getattr(args, k, None) if v is not None: cfg[k] = v if getattr(args, "rope_scale", None) is not None: cfg["rope_scale"] = args.rope_scale return cfg def _push_hf(path, repo): """Upload a single checkpoint file to HF (background thread).""" if not repo: return try: from huggingface_hub import HfApi token = os.environ.get("HF_TOKEN") if not token: for p in (os.path.join(HERE, ".env"), os.path.join(os.path.dirname(HERE), ".env"), os.path.join(os.path.expanduser("~"), ".env")): if os.path.exists(p): for line in open(p, encoding="utf-8"): if line.strip().startswith("HF_TOKEN"): token = line.split("=", 1)[1].strip().strip('"').strip("'") api = HfApi(token=token) api.upload_file(path_or_fileobj=path, path_in_repo=os.path.basename(path), repo_id=repo, repo_type="model") print(f"[hf] pushed {os.path.basename(path)} -> {repo}", flush=True) except Exception as e: print(f"[hf] push failed for {path}: {e}", flush=True) def load_data(data_dir): meta = json.load(open(os.path.join(data_dir, "meta.json"))) arr = np.memmap(os.path.join(data_dir, "train.bin"), dtype=np.uint16, mode="r") return arr, meta["seq_len"], meta["vocab_size"], meta["n_tokens"] def sample_batch(arr, seq_len, batch): N = len(arr) starts = np.random.randint(0, N - seq_len, size=batch) out = np.stack([arr[s:s + seq_len].astype(np.int64) for s in starts]) return torch.from_numpy(out).long() def train(args): data_dir = args.data_dir or DATADIR ckpt_dir = args.ckpt_dir or CKPTDIR os.makedirs(ckpt_dir, exist_ok=True) cfg = build_cfg(args) # ---- device resolution (cuda / xla / cpu) ---- if getattr(args, "device", "cuda") == "xla": import torch_xla.core.xla_model as xm device = xm.xla_device() print(f"[train] device = TPU:XLA ({device})") elif getattr(args, "device", "cuda") == "cpu": device = torch.device("cpu") print("[train] device = CPU") else: device = torch.device("cuda") print(f"[train] device = {device}") tok = YKTokenizer.load(os.path.join(data_dir, "tokenizer.json")) arr, seq_len, vocab, n_tokens = load_data(data_dir) cfg["vocab_size"] = vocab cfg["max_len"] = seq_len print(f"[train] data n_tokens={n_tokens:,} seq_len={seq_len} vocab={vocab}") print(f"[train] model params = {sum(p.numel() for p in YKDiff(cfg).parameters())/1e6:.1f}M") model = YKDiff(cfg).to(device) # Train in bf16 weights so a large model fits a 22 GB L4: fp32 weights would # need ~21.6 GB just for params+grads+Adam states (OOM). bf16 halves that. if device.type in ("cuda", "xla"): model = model.to(torch.bfloat16) print("[train] using bfloat16 weights") n_params = sum(p.numel() for p in model.parameters()) print(f"[train] allocated params = {n_params/1e6:.1f}M") optim = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.1) V = cfg["vocab_size"] pad_id = tok.pad_id mask_id = tok.mask_id # resume step0 = 0 ckpts = sorted([f for f in os.listdir(ckpt_dir) if f.endswith(".pt")]) if ckpts and not args.fresh: path = os.path.join(ckpt_dir, ckpts[-1]) sd = torch.load(path, map_location=device) model.load_state_dict(sd["model"]); optim.load_state_dict(sd["optim"]) step0 = sd["step"] print(f"[train] resumed from {path} step={step0}") model.train() amp = torch.amp.autocast(device_type=device.type, dtype=torch.bfloat16) t0 = time.time() limit = args.hours * 3600.0 step = step0 running = 0.0 while True: if time.time() - t0 > limit: print(f"[train] wall-clock limit {args.hours}h reached at step {step}") break optim.zero_grad(set_to_none=True) mode_ar = (torch.rand(1).item() < 0.5) idx = sample_batch(arr, seq_len, args.batch).to(device) with amp: if mode_ar: m = torch.zeros(args.batch, dtype=torch.long, device=device) logits = model(idx, m, t=None) # causal loss = F.cross_entropy( logits[:, :-1].reshape(-1, V), idx[:, 1:].reshape(-1), ignore_index=pad_id) mname = "AR" else: m = torch.ones(args.batch, dtype=torch.long, device=device) r = torch.rand(args.batch, device=device) # per-sample ratio is_mask = torch.rand(args.batch, seq_len, device=device) < r[:, None] not_pad = idx != pad_id masked = idx.clone(); masked[is_mask] = mask_id logits = model(masked, m, t=r) ce = F.cross_entropy(logits.reshape(-1, V), idx.reshape(-1), reduction="none", ignore_index=-100) ce = ce * is_mask.reshape(-1) * not_pad.reshape(-1) denom = (is_mask & not_pad).reshape(-1).sum().clamp(min=1) loss = ce.sum() / denom mname = "DIFF" loss.backward() nn.utils.clip_grad_norm_(model.parameters(), 1.0) optim.step() if device.type == "xla": xm.mark_step() running = running * 0.9 + float(loss.item()) * 0.1 step += 1 if step % args.log_every == 0: print(f"[train] step {step} [{mname}] loss={running:.3f} " f"t={(time.time()-t0)/60:.1f}m", flush=True) if step % args.ckpt_every == 0: path = os.path.join(ckpt_dir, f"clanker_{step:07d}.pt") torch.save({"model": model.state_dict(), "optim": optim.state_dict(), "step": step, "cfg": cfg, "vocab": V}, path) print(f"[train] checkpoint -> {path}", flush=True) if args.hf_repo: threading.Thread(target=_push_hf, args=(path, args.hf_repo), daemon=True).start() # final save path = os.path.join(ckpt_dir, f"clanker_{step:07d}_final.pt") torch.save({"model": model.state_dict(), "optim": optim.state_dict(), "step": step, "cfg": cfg, "vocab": V}, path) json.dump(cfg, open(os.path.join(ckpt_dir, "config.json"), "w")) print(f"[train] DONE final={path} steps={step}") if __name__ == "__main__": ap = argparse.ArgumentParser() ap.add_argument("--hours", type=float, default=5.0) ap.add_argument("--batch", type=int, default=32) ap.add_argument("--lr", type=float, default=3e-4) ap.add_argument("--log-every", type=int, default=25) ap.add_argument("--ckpt-every", type=int, default=500) ap.add_argument("--fresh", action="store_true") ap.add_argument("--device", default="cuda", choices=["cuda", "xla", "cpu"], help="training device (cuda default; xla for TPU; cpu for tests)") ap.add_argument("--data-dir", default=None) ap.add_argument("--ckpt-dir", default=None) ap.add_argument("--hf-repo", default=None, help="HuggingFace repo id to push checkpoints to (via `hf` CLI)") # model overrides (for a bigger Modal model) ap.add_argument("--d-model", type=int, default=None) ap.add_argument("--n-layers", type=int, default=None) ap.add_argument("--n-heads", type=int, default=None) ap.add_argument("--d-ff", type=int, default=None) ap.add_argument("--rope-scale", type=float, default=None) ap.add_argument("--vocab-size", type=int, default=None) ap.add_argument("--seq-len", type=int, default=None) train(ap.parse_args())