colab_train.py

"""
============================================================================
 clankerDiffusion — Colab A100-80GB training script
============================================================================

HOW TO RUN IN COLAB (A100 80GB):

  [Cell 1 — setup, run once]
      !pip install -q torch==2.11.0+cu124 -f https://download.pytorch.org/whl/cu124
      !pip install -q transformers tokenizers datasets accelerate safetensors huggingface_hub numpy
      import os
      os.environ["HF_TOKEN"]   = "hf_xxx"          # your token (also set in Secrets)
      os.environ["CODE_REPO"]  = "clankerDiffusion/base"        # from upload_artifacts.py
      os.environ["CKPT_REPO"] = "clankerDiffusion/checkpoints"

  [Cell 2 — launch in BACKGROUND, then disconnect safely]
      !nohup python colab_train.py > colab_train.log 2>&1 &
      # check later with:  !tail -n 30 colab_train.log
      # it checkpoints + uploads to HF every 250 steps until the runtime dies

The script trains the SAME from-scratch hybrid model (AR + masked
diffusion) on FineWeb-edu, resuming if a checkpoint exists, and pushes
a checkpoint to HuggingFace Hub every 250 steps in a background thread.
============================================================================
"""
import os, sys, json, time, threading, argparse
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import snapshot_download, HfApi

CODE_REPO = os.environ.get("CODE_REPO", "clankerDiffusion/base")
CKPT_REPO = os.environ.get("CKPT_REPO", "clankerDiffusion/checkpoints")
HF_TOKEN = os.environ.get("HF_TOKEN")
api = HfApi(token=HF_TOKEN)

# ---- pull our model code + tokenizer from HF -------------------------------
print(f"[colab] downloading code from {CODE_REPO} ...")
local = snapshot_download(CODE_REPO, repo_type="model")
sys.path.insert(0, local)
from model import YKDiff
from tokenizer import YKTokenizer

# ---- big architecture for A100 80GB -------------------------------------
CFG = dict(
    d_model=2048, n_layers=24, n_heads=16, d_ff=5504,
    max_len=2048, vocab_size=32768,
)

tok = YKTokenizer.load(os.path.join(local, "tokenizer.json"))
CFG["vocab_size"] = tok.vocab_size
V = CFG["vocab_size"]
mask_id, pad_id = tok.mask_id, tok.pad_id
print(f"[colab] vocab={V}")

# ---- streaming fineweb into a rolling token buffer -------------------------
from datasets import load_dataset
ds = load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT",
                  streaming=True, split="train")
BUF_CAP = 60_000_000
buf = []
_buf_lock = threading.Lock()


def _refill():
    for ex in ds:
        ids = tok.encode(ex["text"])
        with _buf_lock:
            buf.extend(ids)
            if len(buf) > BUF_CAP:
                del buf[: len(buf) - BUF_CAP]


threading.Thread(target=_refill, daemon=True).start()


def sample_batch(batch, seq_len):
    with _buf_lock:
        if len(buf) < seq_len + 1:
            return None
        N = len(buf)
        starts = np.random.randint(0, N - seq_len, size=batch)
        return torch.tensor(
            [buf[s:s + seq_len] for s in starts], dtype=torch.long)


# ---- model -----------------------------------------------------------------
model = YKDiff(CFG).cuda()
n_params = sum(p.numel() for p in model.parameters())
print(f"[colab] params = {n_params/1e9:.2f}B")
optim = torch.optim.AdamW(model.parameters(), lr=1e-4, betas=(0.9, 0.95),
                          weight_decay=0.1)

# ---- resume ----------------------------------------------------------------
CKPT_LOCAL = "/content/clanker_ckpts"
os.makedirs(CKPT_LOCAL, exist_ok=True)
step0 = 0
existing = sorted(f for f in os.listdir(CKPT_LOCAL) if f.endswith(".pt"))
if existing:
    sd = torch.load(os.path.join(CKPT_LOCAL, existing[-1]), map_location="cuda")
    model.load_state_dict(sd["model"]); optim.load_state_dict(sd["optim"])
    step0 = sd["step"]
    print(f"[colab] resumed step={step0}")
else:
    # try pulling latest from HF
    try:
        api.create_repo(CKPT_REPO, repo_type="model", exist_ok=True)
    except Exception:
        pass


def _upload(path):
    def _u():
        try:
            api.upload_file(repo_id=CKPT_REPO,
                           path_in_repo=os.path.basename(path),
                           path_or_fileobj=path, repo_type="model")
            print(f"[colab] uploaded {os.path.basename(path)} -> {CKPT_REPO}",
                  flush=True)
        except Exception as e:
            print(f"[colab] upload failed: {e}", flush=True)
    threading.Thread(target=_u, daemon=True).start()


# ---- training loop (hybrid) ----------------------------------------------
@torch.no_grad()
def _cosine_lr(step, warmup, total, base, minlr):
    if step < warmup:
        return base * step / warmup
    p = (step - warmup) / max(total - warmup, 1)
    return minlr + 0.5 * (base - minlr) * (1 + np.cos(np.pi * min(p, 1.0)))


BATCH, SEQ, GRAD_ACCUM = 16, 2048, 4
WARMUP, TOTAL_STEPS = 500, 200_000
BASE_LR, MIN_LR = 1e-4, 1e-5
CKPT_EVERY = 250
amp = torch.cuda.amp.autocast(dtype=torch.bfloat16)

step = step0
model.train()
t0 = time.time()
print("[colab] training started.", flush=True)

while True:                                   # run as long as possible
    try:
        optim.zero_grad(set_to_none=True)
        for micro in range(GRAD_ACCUM):
            idx = None
            while idx is None:
                idx = sample_batch(BATCH, SEQ)
                time.sleep(0.02)
            idx = idx.cuda()
            mode_ar = (torch.rand(1).item() < 0.5)
            with amp:
                if mode_ar:
                    m = torch.zeros(BATCH, dtype=torch.long, device="cuda")
                    logits = model(idx, m, t=None)
                    loss = F.cross_entropy(
                        logits[:, :-1].reshape(-1, V),
                        idx[:, 1:].reshape(-1), ignore_index=pad_id)
                else:
                    m = torch.ones(BATCH, dtype=torch.long, device="cuda")
                    r = torch.rand(BATCH, device="cuda")
                    is_mask = torch.rand(BATCH, SEQ, device="cuda") < 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
            (loss / GRAD_ACCUM).backward()

        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        lr = _cosine_lr(step, WARMUP, TOTAL_STEPS, BASE_LR, MIN_LR)
        for g in optim.param_groups:
            g["lr"] = lr
        optim.step()
        step += 1

        if step % 25 == 0:
            print(f"[colab] step {step} loss~{loss.item():.3f} "
                  f"lr={lr:.2e} t={(time.time()-t0)/60:.1f}m", flush=True)

        if step % CKPT_EVERY == 0:
            # full local ckpt (for resume)
            full = os.path.join(CKPT_LOCAL, f"clanker_{step:07d}.pt")
            torch.save({"model": model.state_dict(),
                        "optim": optim.state_dict(),
                        "step": step, "cfg": CFG, "vocab": V}, full)
            # light bf16 model-only for HF upload
            lite = os.path.join(CKPT_LOCAL, f"clanker_{step:07d}_lite.pt")
            torch.save({"model": {k: v.to(torch.bfloat16)
                                   for k, v in model.state_dict().items()},
                        "cfg": CFG, "vocab": V, "step": step}, lite)
            print(f"[colab] checkpoint {step}", flush=True)
            _upload(lite)
            # keep only last 2 local full ckpts to save disk
            for old in sorted(f for f in os.listdir(CKPT_LOCAL)
                           if f.endswith(".pt") and "lite" not in f)[:-2]:
                os.remove(os.path.join(CKPT_LOCAL, old))

    except torch.cuda.OutOfMemoryError:
        print("[colab] OOM — skipping step", flush=True)
        optim.zero_grad(set_to_none=True)
        torch.cuda.empty_cache()
    except Exception as e:
        print(f"[colab] step error (continuing): {e}", flush=True)
        torch.cuda.empty_cache()

print("[colab] loop ended.")