File size: 6,749 Bytes
7da1a03 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | """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()
|