Spaces:
Runtime error
Runtime error
File size: 3,952 Bytes
cc15101 | 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 | """
Phase L - standalone val-set perplexity for Abyss-1.
Walks sequentially through val.bin with non-overlapping windows (no
randomness), ao the reported perplexity is a true full-dataset figure.
Usage (from the_final_product/ directory):
python -m ml.eval
python -m ml.eval --checkpoint ml/checkpoints/ckpt.pt --batch_size 32
"""
import argparse
import math
import os
import numpy as np
import torch
from contextlib import nullcontext
from ml.model import TheAbyssModel
VAL_BIN = os.path.join(os.path.dirname(os.path.abspath(__file__)),
"data", "tinystories", "val.bin")
@torch.no_grad()
def compute_perplexity(model, device, block_size, batch_size, ctx):
model.eval()
data = np.memmap(VAL_BIN, dtype=np.uint16, mode="r")
n = len(data)
total_loss = 0.0
total_tokens = 0
batches_done = 0
pos = 0
while pos + block_size + 1 <= n:
starts = []
while len(starts) < batch_size and pos + block_size + 1 <= n:
starts.append(pos)
pos += block_size
xs = torch.stack([
torch.from_numpy(data[i : i + block_size].astype(np.int64))
for i in starts
])
ys = torch.stack([
torch.from_numpy(data[i + 1 : i + 1 + block_size].astype(np.int64))
for i in starts
])
if device.startswith("cuda"):
xs = xs.pin_memory().to(device, non_blocking=True)
ys = ys.pin_memory().to(device, non_blocking=True)
else:
xs, ys = xs.to(device), ys.to(device)
with ctx:
_, loss = model(xs, ys)
n_tok = xs.numel()
total_loss += loss.item() * n_tok
total_tokens += n_tok
batches_done += 1
if batches_done % 50 == 0:
running_ppl = math.exp(total_loss / total_tokens)
pct = 100.0 * pos / n
print(f" {pct:5.1f}% | {total_tokens:>10,} tok | running ppl {running_ppl:.2f}",
flush=True)
avg_loss = total_loss / total_tokens
return avg_loss, math.exp(avg_loss), total_tokens
def main():
parser = argparse.ArgumentParser(description="Val-set perplexity for Abyss-1.")
parser.add_argument("--checkpoint", default="ml/checkpoints/ckpt.pt")
parser.add_argument("--batch_size", type=int, default=16)
parser.add_argument("--block_size", type=int, default=512)
parser.add_argument("--device",
default="cuda" if torch.cuda.is_available() else "cpu")
args = parser.parse_args()
# Load Checkpoint
print(f"loading {args.checkpoint} ...", flush=True)
ck = torch.load(args.checkpoint, map_location=args.device, weights_only=False)
config = ck["config"]
model = TheAbyssModel(config).to(args.device)
model.load_state_dict(ck["model"])
print(f" iter {ck['iter']} | best_val {ck['best_val']:.4f} "
f" | params {model.num_params():,}", flush=True)
ptdtype = (torch.bfloat16 if args.device == "cuda" and torch.cuda.is_bf16_supported()
else torch.float16 if args.device == "cuda" else torch.float32)
ctx = (nullcontext() if args.device == "cpu"
else torch.autocast(device_type="cuda", dtype=ptdtype))
print(f"evaluating full val set "
f"(block={args.block_size}, batch={args.batch_size}, {args.device}) ...\n",
flush=True)
loss, ppl, n_tok = compute_perplexity(
model, args.device, args.block_size, args.batch_size, ctx
)
bar = "=" * 42
print(f"\n{bar}")
print(f" checkpoint : {args.checkpoint}")
print(f" iter : {ck['iter']}")
print(f" val loss : {loss:.4f}")
print(f" perplexity : {ppl:.2f}")
print(f" tokens : {n_tok:,}")
print(bar)
if __name__ == "__main__":
main()
|