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