"""Pretrain from scratch or fine-tune Eather-500M on prepared tokens. Run from the project folder. A GPU is strongly recommended. """ import argparse import json import math import random import time from pathlib import Path import numpy as np import torch from model import Eather500M, ModelConfig def main(): p = argparse.ArgumentParser() p.add_argument("--config", default="config.json") p.add_argument("--data-dir", required=True) p.add_argument("--out-dir", default="runs/pretrain") p.add_argument("--steps", type=int, default=10000) p.add_argument("--batch-size", type=int, default=1) p.add_argument("--seq-len", type=int, default=256) p.add_argument("--grad-accum", type=int, default=16) p.add_argument("--lr", type=float, default=3e-4) p.add_argument("--min-lr", type=float, default=3e-5) p.add_argument("--warmup", type=int, default=200) p.add_argument("--save-every", type=int, default=500) p.add_argument("--log-every", type=int, default=10) p.add_argument("--resume", default=None, help="Checkpoint to resume or fine-tune") p.add_argument("--fine-tune", action="store_true", help="Load model weights only, reset optimizer and step") p.add_argument("--seed", type=int, default=42) p.add_argument("--no-checkpointing", action="store_true") p.add_argument("--precision", choices=["auto", "bf16", "fp16", "fp32"], default="auto") args = p.parse_args() if min(args.steps, args.batch_size, args.seq_len, args.grad_accum) <= 0: raise ValueError("Steps, batch size, seq len, and grad accum must be positive") torch.manual_seed(args.seed) np.random.seed(args.seed) random.seed(args.seed) cfg = ModelConfig.load(args.config) if args.seq_len > cfg.max_seq_len: raise ValueError("seq-len exceeds configured max_seq_len") data_dir = Path(args.data_dir) tokens = np.memmap(data_dir / "train.bin", dtype=np.uint32, mode="r") mask = np.memmap(data_dir / "mask.bin", dtype=np.uint8, mode="r") if len(tokens) != len(mask) or len(tokens) <= args.seq_len: raise ValueError("Data/mask length mismatch or dataset smaller than seq-len + 1") if int(tokens.max()) >= cfg.vocab_size: raise ValueError("Token ID outside model vocabulary. Check tokenizer/config.") device = "cuda" if torch.cuda.is_available() else "cpu" # ROCm GPUs appear as cuda in PyTorch. if device == "cpu": print("WARNING: CPU training for 500M parameters will be extremely slow.", flush=True) supported_bf16 = device == "cuda" and torch.cuda.is_bf16_supported() precision = ("bf16" if supported_bf16 else "fp16") if args.precision == "auto" and device == "cuda" else args.precision if precision == "auto": precision = "fp32" if device == "cpu" and precision != "fp32": raise ValueError("CPU training requires fp32 in this script") dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision] amp = device == "cuda" and precision != "fp32" print("Building Eather model...", flush=True) model = Eather500M(cfg).to(device) model.gradient_checkpointing = not args.no_checkpointing n_params = sum(p.numel() for p in model.parameters()) print(f"{n_params:,} parameters | device={device} | precision={precision} | tokens={len(tokens):,}", flush=True) optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.1) # Scale gradients only when using float16. scaler = torch.amp.GradScaler("cuda", enabled=(precision == "fp16" and device == "cuda")) start_step = 0 if args.resume: ck = torch.load(args.resume, map_location="cpu", weights_only=False) model.load_state_dict(ck["model"]) if not args.fine_tune: optimizer.load_state_dict(ck["optimizer"]) start_step = int(ck["step"]) if "scaler" in ck: scaler.load_state_dict(ck["scaler"]) del ck print(f"Loaded {args.resume}; starting at step {start_step}", flush=True) out = Path(args.out_dir) out.mkdir(parents=True, exist_ok=True) (out / "config.json").write_text(json.dumps(vars(cfg), indent=2)) def get_batch(): # Sample contiguous token blocks from the tokenized corpus. hi = len(tokens) - args.seq_len - 1 ix = np.random.randint(0, hi + 1, size=args.batch_size) x = np.stack([tokens[i:i + args.seq_len] for i in ix]).astype(np.int64) y = np.stack([tokens[i + 1:i + args.seq_len + 1] for i in ix]).astype(np.int64) m = np.stack([mask[i + 1:i + args.seq_len + 1] for i in ix]).astype(np.float32) return (torch.from_numpy(x).to(device), torch.from_numpy(y).to(device), torch.from_numpy(m).to(device)) def save(step): dest = out / f"step_{step:07d}.pt" torch.save({ "model": model.state_dict(), "optimizer": optimizer.state_dict(), "scaler": scaler.state_dict(), "step": step, "config": vars(cfg), }, dest) print(f"Saved {dest}", flush=True) model.train() t0 = time.monotonic() for step in range(start_step, args.steps): if step < args.warmup: lr = args.lr * (step + 1) / max(args.warmup, 1) else: progress = min(1.0, (step - args.warmup) / max(args.steps - args.warmup, 1)) lr = args.min_lr + 0.5 * (args.lr - args.min_lr) * (1 + math.cos(math.pi * progress)) for group in optimizer.param_groups: group["lr"] = lr optimizer.zero_grad(set_to_none=True) running = 0.0 for _ in range(args.grad_accum): x, y, m = get_batch() with torch.autocast(device_type=device, dtype=dtype, enabled=amp): _, loss = model(x, labels=y, loss_mask=m) # Skip an all-zero mask sample without leaving an invalid gradient. if m.sum().item() == 0: continue running += loss.item() / args.grad_accum scaler.scale(loss / args.grad_accum).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update() if (step + 1) % args.log_every == 0 or step == start_step: elapsed = max(time.monotonic() - t0, 1e-6) tokens_seen = (step + 1 - start_step) * args.batch_size * args.grad_accum * args.seq_len print(f"step={step+1} loss={running:.4f} lr={lr:.2e} " f"tokens/s={tokens_seen / elapsed:.1f}", flush=True) if (step + 1) % args.save_every == 0 or (step + 1) == args.steps: save(step + 1) if __name__ == "__main__": main()