""" Phase G/H - pretraining loop for Abyss-1. AdamW + cosine LR (with warmup), gradient accumulation, mixed precision, periodic eval on the held-out val split, and checkpointing (so a free GPU that disconnects can resume). Reads tokens via ml.data.dataset.get_batch (memmapped .bin from prepare.py). python -m ml.train # train from scratch # resume: set init_from = "resume" below, then rerun the same command """ import os import math import time import signal import shutil from contextlib import nullcontext import torch from ml.model import AbyssConfig, TheAbyssModel from ml.data.dataset import get_batch # ----------------------------- knobs ----------------------------- out_dir = "ml/checkpoints" init_from = os.environ.get("INIT_FROM", "scratch") # "scratch" or "resume" backup_dir = os.environ.get("BACKUP_DIR", "") # if set, checkpoint is copied here after each save (e.g. Drive on Colab) batch_size = 32 # sequences per micro-step grad_accum = 8 # micro-steps per optimizer step (eff. batch 256 seqs) block_size = 512 # context length; must be <= AbyssConfig.block_size max_iters = 20_000 # optimizer steps warmup_iters = 200 lr_decay_iters = 20_000 # cosine decays over the whole run learning_rate = 6e-4 min_lr = 6e-5 weight_decay = 0.1 beta1, beta2 = 0.9, 0.95 grad_clip = 1.0 eval_interval = 500 eval_iters = 100 log_interval = 10 seed = 1337 device = "cuda" if torch.cuda.is_available() else "cpu" dtype = ("bfloat16" if device == "cuda" and torch.cuda.is_bf16_supported() else "float16" if device == "cuda" else "float32") # ----------------------------------------------------------------- ctx = nullcontext() # set properly in main() def get_lr(it): """Linear warmup, then cosine decay to min_lr.""" if it < warmup_iters: return learning_rate * (it + 1) / (warmup_iters + 1) if it > lr_decay_iters: return min_lr ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters) coeff = 0.5 * (1.0 + math.cos(math.pi * ratio)) # 1 -> 0 return min_lr + coeff * (learning_rate - min_lr) def configure_optimizer(model): """Weight-decay the 2D matrices (matmuls/embeddings), not the 1D params (norm gains, biases).""" decay, no_decay = [], [] for p in model.parameters(): if not p.requires_grad: continue (decay if p.dim() >= 2 else no_decay).append(p) groups = [ {"params": decay, "weight_decay": weight_decay}, {"params": no_decay, "weight_decay": 0.0}, ] return torch.optim.AdamW(groups, lr=learning_rate, betas=(beta1, beta2)) @torch.no_grad() def estimate_loss(model): model.eval() out = {} for split in ("train", "val"): losses = torch.zeros(eval_iters) for k in range(eval_iters): x, y = get_batch(split, batch_size, block_size, device) with ctx: _, loss = model(x, y) losses[k] = loss.item() out[split] = losses.mean().item() model.train() return out def main(): global ctx torch.manual_seed(seed) torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True os.makedirs(out_dir, exist_ok=True) config = AbyssConfig() assert block_size <= config.block_size, "block_size exceeds the model's context" ptdtype = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[dtype] ctx = nullcontext() if device == "cpu" else torch.autocast(device_type="cuda", dtype=ptdtype) try: # modern API (torch >= 2.3) scaler = torch.amp.GradScaler(device, enabled=(dtype == "float16")) except (AttributeError, TypeError): # fall back on older torch scaler = torch.cuda.amp.GradScaler(enabled=(dtype == "float16")) model = TheAbyssModel(config).to(device) opt = configure_optimizer(model) start_iter, best_val = 0, float("inf") if init_from == "resume": ck = torch.load(os.path.join(out_dir, "ckpt.pt"), map_location=device, weights_only=False) model.load_state_dict(ck["model"]) opt.load_state_dict(ck["optimizer"]) start_iter = ck["iter"] + 1 best_val = ck["best_val"] print(f"resumed from iter {ck['iter']} (best val {best_val:.4f})") print(f"device {device} | dtype {dtype} | params {model.num_params():,} | " f"eff batch {batch_size*grad_accum} seqs x {block_size} = {batch_size*grad_accum*block_size:,} tok/step") # SIGINT handler — Ctrl+C saves a checkpoint instead of dying mid-step interrupted = False def _handle_sigint(sig, frame): nonlocal interrupted print("\nSIGINT received — finishing this step then saving checkpoint.") interrupted = True signal.signal(signal.SIGINT, _handle_sigint) tokens_per_iter = batch_size * grad_accum * block_size t0 = time.time() for it in range(start_iter, max_iters + 1): for g in opt.param_groups: g["lr"] = get_lr(it) if it % eval_interval == 0: losses = estimate_loss(model) best_val = min(best_val, losses["val"]) print(f"step {it:6d} | train {losses['train']:.4f} | val {losses['val']:.4f} " f"| best {best_val:.4f} | lr {opt.param_groups[0]['lr']:.2e}") torch.save({"model": model.state_dict(), "optimizer": opt.state_dict(), "config": config, "iter": it, "best_val": best_val}, os.path.join(out_dir, "ckpt.pt")) if backup_dir: os.makedirs(backup_dir, exist_ok=True) shutil.copy(os.path.join(out_dir, "ckpt.pt"), os.path.join(backup_dir, "ckpt.pt")) # one optimizer step = grad_accum micro-steps for _ in range(grad_accum): x, y = get_batch("train", batch_size, block_size, device) with ctx: _, loss = model(x, y) loss = loss / grad_accum scaler.scale(loss).backward() if grad_clip: scaler.unscale_(opt) torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip) scaler.step(opt) scaler.update() opt.zero_grad(set_to_none=True) raw_loss = loss.item() * grad_accum # NaN guard if math.isnan(raw_loss): print(f"NaN loss at iter {it} — saving emergency checkpoint and exiting.") torch.save({"model": model.state_dict(), "optimizer": opt.state_dict(), "config": config, "iter": it, "best_val": best_val}, os.path.join(out_dir, "ckpt_nan.pt")) break # SIGINT check if interrupted: print(f"Saving checkpoint at iter {it} and exiting.") torch.save({"model": model.state_dict(), "optimizer": opt.state_dict(), "config": config, "iter": it, "best_val": best_val}, os.path.join(out_dir, "ckpt.pt")) break if it % log_interval == 0: dt = (time.time() - t0) / log_interval t0 = time.time() tok_per_sec = tokens_per_iter / dt eta_sec = dt * (max_iters - it) eta_str = f"{int(eta_sec // 3600)}h{int((eta_sec % 3600) // 60)}m" print(f" iter {it:6d} | loss {raw_loss:.4f} | {tok_per_sec/1000:.1f}k tok/s | ETA {eta_str}") print("training loop finished") if __name__ == "__main__": main()