train.py

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