Spaces:
Runtime error
Runtime error
| """ | |
| 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)) | |
| 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() | |