File size: 7,285 Bytes
5b1ff4d | 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 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 | """
calibration_task.py — Top-k accuracy and entropy calibration evaluation.
Top-level function for ProcessPoolExecutor (spawn) compatibility:
- eval_calibration(device, n_tokens=50000) -> dict
"""
from __future__ import annotations
import sys
import time
from pathlib import Path
import os
import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
_PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
if str(_PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(_PROJECT_ROOT))
_DEFAULT_CHECKPOINT = str(_PROJECT_ROOT / "checkpoints" / "korean_3b_fp8_run1" / "checkpoint-0057000")
CHECKPOINT = os.environ.get("EVAL_CHECKPOINT", _DEFAULT_CHECKPOINT)
TOKENIZER_PATH = os.environ.get("EVAL_TOKENIZER", str(_PROJECT_ROOT / "tokenizer" / "korean_sp" / "tokenizer.json"))
DATA_DIR = _PROJECT_ROOT / "data"
SEQ_LEN = 2048
STRIDE = 512
BATCH_SIZE = 32
# ---------------------------------------------------------------------------
# Shared dataset / model utilities
# ---------------------------------------------------------------------------
class SlidingWindowDataset(Dataset):
"""Sliding-window tokenized dataset for evaluation."""
def __init__(self, tokens: np.ndarray, seq_len: int, stride: int) -> None:
self.tokens = tokens
self.seq_len = seq_len
self.stride = stride
self.n_windows = max(0, (len(tokens) - seq_len + stride - 1) // stride)
def __len__(self) -> int:
return self.n_windows
def __getitem__(self, idx: int):
start = idx * self.stride
end = start + self.seq_len
actual_end = min(end, len(self.tokens))
chunk_len = actual_end - start
input_ids = torch.zeros(self.seq_len, dtype=torch.long)
targets = torch.full((self.seq_len,), fill_value=-100, dtype=torch.long)
loss_mask = torch.zeros(self.seq_len, dtype=torch.bool)
if chunk_len > 1:
toks = torch.from_numpy(self.tokens[start:actual_end].astype(np.int64))
input_ids[:chunk_len] = toks
targets[:chunk_len - 1] = toks[1:]
new_start = 0 if idx == 0 else self.stride
if chunk_len > 1:
for pos in range(new_start, chunk_len - 1):
loss_mask[pos] = True
return input_ids, targets, loss_mask
def _load_model(device: str):
"""Load FRANKENSTALLM 3B from checkpoint onto the given device."""
from model.transformer import LLM # type: ignore[import]
model = LLM.from_pretrained(CHECKPOINT)
model = model.to(device=device, dtype=torch.bfloat16)
model.eval()
return model
def _load_tokenizer():
"""Load the Korean SentencePiece tokenizer."""
from tokenizers import Tokenizer # type: ignore[import]
return Tokenizer.from_file(TOKENIZER_PATH)
# ---------------------------------------------------------------------------
# Main task function (must be top-level for pickle / spawn compatibility)
# ---------------------------------------------------------------------------
def eval_calibration(device: str, n_tokens: int = 50000) -> dict:
"""Compute top-k accuracy and entropy calibration on 3b_val.bin.
Measures how well the model's probability distribution is calibrated:
- Top-1/5/10 next-token prediction accuracy
- Mean probability assigned to the correct next token
- Mean Shannon entropy of the predictive distribution
Args:
device: CUDA device string, e.g. "cuda:3".
n_tokens: Number of tokens to evaluate (first n_tokens of 3b_val.bin).
Returns:
Dict with keys: n_eval_tokens, top1_accuracy, top5_accuracy,
top10_accuracy, mean_correct_prob, mean_entropy, elapsed_sec.
"""
torch.cuda.set_device(int(device.split(":")[-1]))
print(f"[CALIB {device}] Loading model...")
model = _load_model(device)
val_path = DATA_DIR / "3b_val.bin"
if not val_path.exists():
raise FileNotFoundError(f"Validation file not found: {val_path}")
tokens = np.fromfile(str(val_path), dtype=np.uint16)
if len(tokens) == 0:
raise ValueError(f"Validation file is empty (0 tokens): {val_path}")
tokens = tokens[: min(n_tokens, len(tokens))]
print(f"[CALIB {device}] Using {len(tokens):,} tokens from 3b_val.bin")
ds = SlidingWindowDataset(tokens, SEQ_LEN, STRIDE)
dl = DataLoader(
ds,
batch_size=BATCH_SIZE,
shuffle=False,
num_workers=2,
pin_memory=True,
)
top1_correct = 0
top5_correct = 0
top10_correct = 0
total_entropy = 0.0
total_prob = 0.0
total_count = 0
t0 = time.time()
with torch.inference_mode():
for batch_idx, (inp, tgt, mask) in enumerate(dl):
inp = inp.to(device)
tgt = tgt.to(device)
mask = mask.to(device)
logits, _ = model(inp)
probs = F.softmax(logits, dim=-1)
valid = mask & (tgt != -100)
if valid.sum() == 0:
continue
flat_logits = logits[valid]
flat_tgt = tgt[valid]
flat_probs = probs[valid]
# Top-k accuracy
_, top1_pred = flat_logits.topk(1, dim=-1)
_, top5_pred = flat_logits.topk(5, dim=-1)
_, top10_pred = flat_logits.topk(10, dim=-1)
top1_correct += (top1_pred.squeeze(-1) == flat_tgt).sum().item()
top5_correct += (
(top5_pred == flat_tgt.unsqueeze(-1)).any(dim=-1).sum().item()
)
top10_correct += (
(top10_pred == flat_tgt.unsqueeze(-1)).any(dim=-1).sum().item()
)
# Mean probability of correct token
correct_probs = flat_probs[torch.arange(len(flat_tgt), device=device), flat_tgt]
total_prob += correct_probs.sum().item()
# Shannon entropy: H = -sum(p * log(p))
log_probs = torch.log(torch.clamp(flat_probs, min=1e-7))
entropy = -(flat_probs * log_probs).sum(dim=-1)
total_entropy += entropy.sum().item()
total_count += valid.sum().item()
if (batch_idx + 1) % 50 == 0:
elapsed = time.time() - t0
print(
f"[CALIB {device}] batch {batch_idx + 1}/{len(dl)}, "
f"tokens so far={total_count:,}, {elapsed:.0f}s"
)
elapsed = time.time() - t0
result: dict = {
"n_eval_tokens": int(total_count),
"top1_accuracy": round(top1_correct / total_count, 4) if total_count > 0 else 0.0,
"top5_accuracy": round(top5_correct / total_count, 4) if total_count > 0 else 0.0,
"top10_accuracy": round(top10_correct / total_count, 4) if total_count > 0 else 0.0,
"mean_correct_prob": round(total_prob / total_count, 4) if total_count > 0 else 0.0,
"mean_entropy": round(total_entropy / total_count, 4) if total_count > 0 else 0.0,
"elapsed_sec": round(elapsed, 1),
}
print(
f"[CALIB {device}] DONE top1={result['top1_accuracy']:.4f}, "
f"top5={result['top5_accuracy']:.4f}, "
f"top10={result['top10_accuracy']:.4f}, "
f"entropy={result['mean_entropy']:.4f}, {elapsed:.1f}s"
)
return result
|