""" ============================================================================ 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.")