pathcosmos/frankenstallm
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

初始化项目，由ModelHub XC社区提供模型

Browse Source 
Model: pathcosmos/frankenstallm
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-07-14 04:21:16 +08:00
commit d4abdb70fa
516 changed files with 5584477 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
source/eval/tasks/__init__.py Normal file
View File

203
source/eval/tasks/calibration_task.py Normal file
View File

@@ -0,0 +1,203 @@
"""
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

472
source/eval/tasks/generation_task.py Normal file
View File

@@ -0,0 +1,472 @@
"""
generation_task.py — Text generation quality evaluation tasks.
Top-level functions for ProcessPoolExecutor (spawn) compatibility:
- eval_generation(device) -> dict
- eval_repetition_grid(device) -> dict
Helper functions (also top-level, used internally):
- top_p_filtering(logits, top_p, top_k)
- generate_one(model, tokenizer, prompt, temperature, ...)
- compute_ngram_rep(text, n)
"""
from __future__ import annotations
import logging
import os
import sys
import time
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F
logger = logging.getLogger(__name__)
_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"))
# Chat template support for SFT models
USE_CHAT_TEMPLATE = os.environ.get("USE_CHAT_TEMPLATE", "0") == "1"
CHAT_TEMPLATE_FMT = "<|user|>\n{prompt}\n<|assistant|>\n"
DATA_DIR = _PROJECT_ROOT / "data"
SEQ_LEN = 2048
STRIDE = 512
BATCH_SIZE = 32
# ---------------------------------------------------------------------------
# Prompt / temperature constants
# ---------------------------------------------------------------------------
PROMPTS = [
"대한민국의 수도는",
"인공지능이란",
"한국의 전통 음식 중에서",
"지구 온난화의 주요 원인은",
"프로그래밍을 배우려면",
"조선시대에는",
"물리학에서 에너지란",
"한국어는 세계에서",
"경제 성장을 위해서는",
"우주 탐사의 역사를 보면",
"머신러닝과 딥러닝의 차이는",
"한국 문학의 대표적인 작품으로는",
"양자 컴퓨터란",
"건강한 식습관을 위해서는",
"세계 2차 대전 이후",
]
TEMPERATURES = [0.0, 0.5, 0.8, 1.0]
REP_GRID = [
{"name": "greedy", "temperature": 0.0, "repetition_penalty": 1.0},
{"name": "t0.5", "temperature": 0.5, "repetition_penalty": 1.0},
{"name": "t0.5_rep1.1", "temperature": 0.5, "repetition_penalty": 1.1},
{"name": "t0.7", "temperature": 0.7, "repetition_penalty": 1.0},
{"name": "t0.7_rep1.1", "temperature": 0.7, "repetition_penalty": 1.1},
{"name": "t0.7_rep1.2", "temperature": 0.7, "repetition_penalty": 1.2},
{"name": "t0.7_rep1.3", "temperature": 0.7, "repetition_penalty": 1.3},
{"name": "t0.9", "temperature": 0.9, "repetition_penalty": 1.0},
{"name": "t0.9_rep1.1", "temperature": 0.9, "repetition_penalty": 1.1},
{"name": "t0.9_rep1.2", "temperature": 0.9, "repetition_penalty": 1.2},
{"name": "t1.0", "temperature": 1.0, "repetition_penalty": 1.0},
{"name": "t1.0_rep1.1", "temperature": 1.0, "repetition_penalty": 1.1},
]
# ---------------------------------------------------------------------------
# Shared model utilities
# ---------------------------------------------------------------------------
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)
# ---------------------------------------------------------------------------
# Generation helpers (top-level for pickle compatibility)
# ---------------------------------------------------------------------------
def top_p_filtering(logits: torch.Tensor, top_p: float = 0.9, top_k: int = 0) -> torch.Tensor:
"""Apply top-p (nucleus) and/or top-k filtering to a logits tensor.
Args:
logits: Shape (..., vocab_size).
top_p: Nucleus probability threshold in (0, 1). 0 or 1 disables.
top_k: Keep only the top-k tokens. 0 disables.
Returns:
Filtered logits tensor of the same shape.
"""
if logits.dim() == 1:
logits = logits.unsqueeze(0)
squeeze = True
else:
squeeze = False
if top_k > 0:
k = min(top_k, logits.size(-1))
kth = torch.topk(logits, k, dim=-1).values[:, -1, None]
logits = logits.masked_fill(logits < kth, float("-inf"))
if 0.0 < top_p < 1.0:
sorted_logits, sorted_idx = torch.sort(logits, dim=-1, descending=True)
cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
remove = cum_probs - F.softmax(sorted_logits, dim=-1) >= top_p
sorted_logits[remove] = float("-inf")
logits = torch.zeros_like(logits).scatter_(-1, sorted_idx, sorted_logits)
if squeeze:
logits = logits.squeeze(0)
return logits
def generate_one(
model,
tokenizer,
prompt: str,
temperature: float,
top_p: float = 0.9,
top_k: int = 50,
max_new_tokens: int = 256,
device: str = "cuda:0",
repetition_penalty: float = 1.0,
) -> tuple[str, int, bool]:
"""Generate a single continuation for a prompt using the given model.
Args:
model: Pre-loaded language model (eval mode).
tokenizer: Tokenizer with encode/decode methods.
prompt: Input prompt string.
temperature: Sampling temperature. 0.0 = greedy.
top_p: Nucleus filtering threshold.
top_k: Top-k filtering count.
max_new_tokens: Maximum number of tokens to generate.
device: CUDA device string.
repetition_penalty: Penalty > 1.0 discourages token repetition.
Returns:
Tuple of (generated_text, num_new_tokens, hit_eos).
"""
input_ids = torch.tensor(
[tokenizer.encode(prompt).ids], dtype=torch.long, device=device
)
eos_id = tokenizer.token_to_id("</s>")
generated = input_ids
new_ids: list[int] = []
hit_eos = False
for _ in range(max_new_tokens):
logits_all, _ = model(generated)
logits = logits_all[:, -1, :].clone()
if repetition_penalty != 1.0:
for tid in set(generated[0].tolist()):
if logits[0, tid] > 0:
logits[0, tid] /= repetition_penalty
else:
logits[0, tid] *= repetition_penalty
if temperature == 0.0:
next_id = logits.argmax(dim=-1, keepdim=True)
else:
logits = logits / max(temperature, 1e-8)
logits = top_p_filtering(logits, top_p=top_p, top_k=top_k)
probs = F.softmax(logits, dim=-1)
next_id = torch.multinomial(probs, num_samples=1)
generated = torch.cat([generated, next_id], dim=-1)
new_ids.append(next_id.item())
if eos_id is not None and next_id.item() == eos_id:
hit_eos = True
break
text = tokenizer.decode(new_ids)
return text, len(new_ids), hit_eos
def compute_ngram_rep(text: str, n: int) -> float:
"""Compute n-gram repetition rate for a whitespace-tokenized string.
Repetition rate = 1 - (unique n-grams / total n-grams).
A value of 0 means no repeated n-grams; 1 means all n-grams are repeated.
Args:
text: Input text (whitespace-tokenized).
n: N-gram order (1, 2, 3, 4, ...).
Returns:
Float in [0, 1].
"""
tokens = text.split()
if len(tokens) < n:
return 0.0
ngrams = [tuple(tokens[i : i + n]) for i in range(len(tokens) - n + 1)]
if not ngrams:
return 0.0
return 1.0 - len(set(ngrams)) / len(ngrams)
def compute_diversity_metrics(text: str) -> dict:
"""N-gram 반복률을 보완하는 어휘 다양성 메트릭.
- Distinct-n (Li et al., 2016): 고유 n-gram 비율
- Type-Token Ratio: 어휘 풍부도
"""
tokens = text.split()
n = len(tokens)
if n == 0:
return {"distinct_1": 0.0, "distinct_2": 0.0, "distinct_3": 0.0,
"type_token_ratio": 0.0, "vocab_size": 0, "total_tokens": 0}
unigrams = set(tokens)
bigrams = set(zip(tokens, tokens[1:])) if n > 1 else set()
trigrams = set(zip(tokens, tokens[1:], tokens[2:])) if n > 2 else set()
return {
"distinct_1": len(unigrams) / n,
"distinct_2": len(bigrams) / max(n - 1, 1),
"distinct_3": len(trigrams) / max(n - 2, 1),
"type_token_ratio": len(unigrams) / n,
"vocab_size": len(unigrams),
"total_tokens": n,
}
# ---------------------------------------------------------------------------
# Main task functions (must be top-level for pickle / spawn compatibility)
# ---------------------------------------------------------------------------
def eval_generation(device: str) -> dict:
"""Evaluate generation quality: 15 prompts x 4 temperatures.
For each (prompt, temperature) combination:
- Generates up to 256 new tokens
- Computes 1-gram through 4-gram repetition rates
Args:
device: CUDA device string, e.g. "cuda:4".
Returns:
Dict with keys:
- summary: aggregate statistics across all generations
- samples: list of per-generation result dicts
"""
torch.cuda.set_device(int(device.split(":")[-1]))
print(f"[GEN {device}] Loading model...")
model = _load_model(device)
tokenizer = _load_tokenizer()
t0 = time.time()
results: list[dict] = []
total_combinations = len(PROMPTS) * len(TEMPERATURES)
done = 0
if USE_CHAT_TEMPLATE:
print(f"[GEN {device}] Chat template ENABLED", flush=True)
for prompt in PROMPTS:
effective_prompt = CHAT_TEMPLATE_FMT.format(prompt=prompt) if USE_CHAT_TEMPLATE else prompt
for temp in TEMPERATURES:
with torch.inference_mode():
text, n_tokens, hit_eos = generate_one(
model, tokenizer, effective_prompt, temp, device=device
)
rep1 = compute_ngram_rep(text, 1)
rep2 = compute_ngram_rep(text, 2)
rep3 = compute_ngram_rep(text, 3)
rep4 = compute_ngram_rep(text, 4)
diversity = compute_diversity_metrics(text)
entry = {
"prompt": prompt,
"chat_template": USE_CHAT_TEMPLATE,
"effective_prompt": effective_prompt if USE_CHAT_TEMPLATE else prompt,
"temperature": temp,
"generated_tokens": n_tokens,
"hit_eos": hit_eos,
"1gram_rep": round(rep1, 4),
"2gram_rep": round(rep2, 4),
"3gram_rep": round(rep3, 4),
"4gram_rep": round(rep4, 4),
"distinct_1": round(diversity["distinct_1"], 4),
"distinct_2": round(diversity["distinct_2"], 4),
"distinct_3": round(diversity["distinct_3"], 4),
"type_token_ratio": round(diversity["type_token_ratio"], 4),
"text": text[:500], # truncate for readability
}
results.append(entry)
done += 1
label = "greedy" if temp == 0.0 else f"t={temp}"
print(
f"[GEN {device}] ({done}/{total_combinations}) "
f"{prompt[:15]}... ({label}): "
f"{n_tokens}tok, 3gram_rep={rep3:.2%}, eos={hit_eos}"
)
elapsed = time.time() - t0
# Aggregate stats per temperature group
greedy = [r for r in results if r["temperature"] == 0.0]
sampled = [r for r in results if r["temperature"] > 0.0]
if not greedy:
logger.warning("No greedy generation results — all prompts may have failed")
if not sampled:
logger.warning("No sampled generation results")
summary = {
"total_generations": len(results),
"n_prompts": len(PROMPTS),
"temperatures": TEMPERATURES,
"greedy_avg_1gram_rep": round(np.mean([r["1gram_rep"] for r in greedy]), 4) if greedy else 0.0,
"greedy_avg_2gram_rep": round(np.mean([r["2gram_rep"] for r in greedy]), 4) if greedy else 0.0,
"greedy_avg_3gram_rep": round(np.mean([r["3gram_rep"] for r in greedy]), 4) if greedy else 0.0,
"greedy_avg_4gram_rep": round(np.mean([r["4gram_rep"] for r in greedy]), 4) if greedy else 0.0,
"greedy_eos_rate": round(np.mean([r["hit_eos"] for r in greedy]), 4) if greedy else 0.0,
"greedy_avg_tokens": round(np.mean([r["generated_tokens"] for r in greedy]), 1) if greedy else 0.0,
"sampled_avg_3gram_rep": round(np.mean([r["3gram_rep"] for r in sampled]), 4) if sampled else 0.0,
"sampled_eos_rate": round(np.mean([r["hit_eos"] for r in sampled]), 4) if sampled else 0.0,
"sampled_avg_tokens": round(np.mean([r["generated_tokens"] for r in sampled]), 1) if sampled else 0.0,
"greedy_avg_distinct_1": round(float(np.mean([r["distinct_1"] for r in greedy])), 4) if greedy else 0.0,
"greedy_avg_distinct_2": round(float(np.mean([r["distinct_2"] for r in greedy])), 4) if greedy else 0.0,
"greedy_avg_distinct_3": round(float(np.mean([r["distinct_3"] for r in greedy])), 4) if greedy else 0.0,
"sampled_avg_distinct_2": round(float(np.mean([r["distinct_2"] for r in sampled])), 4) if sampled else 0.0,
"token_count_min": int(np.min([r["generated_tokens"] for r in results])) if results else 0,
"token_count_max": int(np.max([r["generated_tokens"] for r in results])) if results else 0,
"token_count_p25": int(np.percentile([r["generated_tokens"] for r in results], 25)) if results else 0,
"token_count_p75": int(np.percentile([r["generated_tokens"] for r in results], 75)) if results else 0,
"elapsed_sec": round(elapsed, 1),
}
print(
f"[GEN {device}] DONE greedy 3gram_rep={summary['greedy_avg_3gram_rep']:.4f}, "
f"eos_rate={summary['greedy_eos_rate']:.2%}, {elapsed:.1f}s"
)
return {"summary": summary, "samples": results}
def eval_repetition_grid(device: str) -> dict:
"""Grid search over 12 generation parameter combinations x 5 prompts.
Evaluates each config (temperature x repetition_penalty) on the first 5
prompts and returns results sorted by average 3-gram repetition rate.
Args:
device: CUDA device string, e.g. "cuda:5".
Returns:
Dict with keys:
- grid_results: list of per-config dicts, sorted by avg_3gram_rep
- best: config with lowest avg_3gram_rep
- elapsed_sec: wall-clock time
"""
torch.cuda.set_device(int(device.split(":")[-1]))
print(f"[REP {device}] Loading model...")
model = _load_model(device)
tokenizer = _load_tokenizer()
t0 = time.time()
rep_prompts = PROMPTS[:5] # first 5 prompts
results: list[dict] = []
total = len(REP_GRID) * len(rep_prompts)
done = 0
if USE_CHAT_TEMPLATE:
print(f"[REP {device}] Chat template ENABLED", flush=True)
for params in REP_GRID:
combo_results: list[dict] = []
for prompt in rep_prompts:
effective_prompt = CHAT_TEMPLATE_FMT.format(prompt=prompt) if USE_CHAT_TEMPLATE else prompt
with torch.inference_mode():
text, n_tokens, hit_eos = generate_one(
model,
tokenizer,
effective_prompt,
temperature=params["temperature"],
repetition_penalty=params["repetition_penalty"],
device=device,
max_new_tokens=256,
)
combo_results.append(
{
"prompt": prompt,
"n_tokens": n_tokens,
"hit_eos": hit_eos,
"1gram_rep": compute_ngram_rep(text, 1),
"2gram_rep": compute_ngram_rep(text, 2),
"3gram_rep": compute_ngram_rep(text, 3),
"4gram_rep": compute_ngram_rep(text, 4),
}
)
done += 1
if not combo_results:
logger.warning("All prompts failed for config %s — skipping", params.get("name", "unknown"))
continue
avg_3gram = float(np.mean([r["3gram_rep"] for r in combo_results]))
avg_4gram = float(np.mean([r["4gram_rep"] for r in combo_results]))
eos_rate = float(np.mean([r["hit_eos"] for r in combo_results]))
avg_tokens = float(np.mean([r["n_tokens"] for r in combo_results]))
entry = {
"params": params["name"],
"temperature": params["temperature"],
"repetition_penalty": params["repetition_penalty"],
"avg_3gram_rep": round(avg_3gram, 4),
"avg_4gram_rep": round(avg_4gram, 4),
"eos_rate": round(eos_rate, 4),
"avg_tokens": round(avg_tokens, 1),
"per_prompt": combo_results,
}
results.append(entry)
print(
f"[REP {device}] {params['name']}: "
f"3gram={avg_3gram:.2%}, 4gram={avg_4gram:.2%}, "
f"eos={eos_rate:.0%}, {avg_tokens:.0f}tok"
)
elapsed = time.time() - t0
# Sort by avg 3-gram repetition (ascending = better)
sorted_results = sorted(results, key=lambda r: r["avg_3gram_rep"])
best = sorted_results[0]
print(
f"[REP {device}] DONE best={best['params']} "
f"(3gram={best['avg_3gram_rep']:.2%}), {elapsed:.1f}s"
)
return {
"grid_results": sorted_results,
"best": {
"params": best["params"],
"temperature": best["temperature"],
"repetition_penalty": best["repetition_penalty"],
"avg_3gram_rep": best["avg_3gram_rep"],
"avg_4gram_rep": best["avg_4gram_rep"],
},
"elapsed_sec": round(elapsed, 1),
}

365
source/eval/tasks/lm_eval_task.py Normal file
View File

@@ -0,0 +1,365 @@
"""
lm_eval_task.py — lm-evaluation-harness integration task.
Top-level function for ProcessPoolExecutor (spawn) compatibility:
- run_lm_eval_tasks(hf_model_path, tasks, device, num_fewshot=0) -> dict
Requires: lm_eval >= 0.4 (installed as lm-eval 0.4.11)
"""
from __future__ import annotations
import logging
import os
import sys
import time
from pathlib import Path
from typing import Any
_PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
if str(_PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(_PROJECT_ROOT))
CHECKPOINT = str(_PROJECT_ROOT / "checkpoints" / "korean_3b_fp8_run1" / "checkpoint-0057000")
TOKENIZER_PATH = str(_PROJECT_ROOT / "tokenizer" / "korean_sp" / "tokenizer.json")
DATA_DIR = _PROJECT_ROOT / "data"
SEQ_LEN = 2048
STRIDE = 512
BATCH_SIZE = 32
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Main task function (must be top-level for pickle / spawn compatibility)
# ---------------------------------------------------------------------------
def run_lm_eval_tasks(
hf_model_path: str,
tasks: list[str],
device: str,
num_fewshot: int = 0,
) -> dict:
"""Run lm-evaluation-harness benchmarks on a HuggingFace-format model.
Isolates a single GPU via CUDA_VISIBLE_DEVICES so the function is safe
to run in a ProcessPoolExecutor worker without VRAM conflicts.
Args:
hf_model_path: Path to a HuggingFace-compatible model directory
(must contain config.json + safetensors/pytorch_model).
tasks: List of lm-eval task names, e.g.
["hellaswag", "arc_easy", "piqa"].
Unknown tasks are skipped with a warning.
device: CUDA device string, e.g. "cuda:7".
The function maps this to CUDA_VISIBLE_DEVICES=7 and
then uses device="cuda:0" inside lm_eval.
num_fewshot: Number of few-shot examples (0 = zero-shot).
Returns:
Dict with keys:
- model_path: hf_model_path as provided
- tasks_requested: original task list
- tasks_evaluated: tasks that were actually run
- tasks_skipped: tasks that were not available / errored
- per_task_metrics: dict mapping task name to metric sub-dict
- raw_results: full results dict from lm_eval.simple_evaluate
- elapsed_sec: wall-clock time for the evaluation
"""
# --- GPU isolation ---
gpu_index = int(device.split(":")[-1])
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_index)
# After this point use cuda:0 since only one GPU is visible
_internal_device = "cuda:0"
print(
f"[LM_EVAL] Starting on {device} "
f"(CUDA_VISIBLE_DEVICES={gpu_index}), tasks={tasks}, "
f"num_fewshot={num_fewshot}"
)
# --- Validate task list ---
try:
import lm_eval # type: ignore[import]
from lm_eval.tasks import TaskManager # type: ignore[import]
task_manager = TaskManager()
available_tasks: set[str] = set(task_manager.all_tasks)
except Exception as exc:
logger.warning(f"[LM_EVAL] Could not enumerate available tasks: {exc}")
available_tasks = set() # will attempt all and catch errors per task
valid_tasks: list[str] = []
skipped_tasks: list[str] = []
for t in tasks:
if (not available_tasks) or (t in available_tasks):
valid_tasks.append(t)
else:
logger.warning(f"[LM_EVAL] Task '{t}' not found in lm_eval registry — skipping.")
skipped_tasks.append(t)
if not valid_tasks:
print("[LM_EVAL] No valid tasks to evaluate.")
return {
"model_path": hf_model_path,
"tasks_requested": tasks,
"tasks_evaluated": [],
"tasks_skipped": skipped_tasks,
"per_task_metrics": {},
"raw_results": {},
"elapsed_sec": 0.0,
}
# --- Run evaluation ---
t0 = time.time()
raw_results: dict[str, Any] = {}
evaluated_tasks: list[str] = []
error_tasks: list[str] = []
# Attempt all valid tasks together first; fall back to per-task on error
try:
print(
f"[LM_EVAL] Evaluating {len(valid_tasks)} task(s) together: {valid_tasks}"
)
raw_results = lm_eval.simple_evaluate(
model="hf",
model_args=(
f"pretrained={hf_model_path},"
f"dtype=bfloat16,"
f"device={_internal_device}"
),
tasks=valid_tasks,
num_fewshot=num_fewshot,
batch_size="auto",
)
evaluated_tasks = list(valid_tasks)
except Exception as exc:
logger.warning(
f"[LM_EVAL] Batch evaluation failed ({exc}). "
"Falling back to per-task evaluation."
)
# Fall back: evaluate one task at a time
for task_name in valid_tasks:
try:
print(f"[LM_EVAL] Evaluating task '{task_name}' individually...")
task_result = lm_eval.simple_evaluate(
model="hf",
model_args=(
f"pretrained={hf_model_path},"
f"dtype=bfloat16,"
f"device={_internal_device}"
),
tasks=[task_name],
num_fewshot=num_fewshot,
batch_size="auto",
device=_internal_device,
)
# Merge per-task results into raw_results
if not raw_results:
raw_results = task_result
else:
if "results" in task_result and "results" in raw_results:
raw_results["results"].update(task_result.get("results", {}))
evaluated_tasks.append(task_name)
except Exception as task_exc:
logger.warning(
f"[LM_EVAL] Task '{task_name}' failed: {task_exc}"
)
error_tasks.append(task_name)
skipped_tasks.extend(error_tasks)
elapsed = time.time() - t0
# --- Extract per-task metrics ---
# Group tasks (e.g. global_mmlu_ko, mmlu) expand to subtasks at eval time.
# Capture ALL result keys, not just the originally requested task names,
# so that subtask-level metrics are available for downstream reporting.
per_task_metrics: dict[str, dict] = {}
lm_results: dict[str, Any] = raw_results.get("results", {})
for task_name, task_data in lm_results.items():
if not isinstance(task_data, dict):
continue
metrics: dict[str, Any] = {}
for key, value in task_data.items():
# Skip non-metric metadata keys
if key in ("alias", "group"):
continue
metrics[key] = value
per_task_metrics[task_name] = metrics
# Warn about any requested tasks that produced no results at all
for task_name in evaluated_tasks:
if task_name not in per_task_metrics:
logger.warning(
f"[LM_EVAL] Task '{task_name}' not found in results dict after evaluation."
)
# --- Summary print ---
print(f"[LM_EVAL] Evaluation complete in {elapsed:.1f}s")
for task_name, metrics in per_task_metrics.items():
# Print the most common accuracy variants
acc = metrics.get("acc,none") or metrics.get("acc") or metrics.get("accuracy")
acc_norm = metrics.get("acc_norm,none") or metrics.get("acc_norm")
if acc is not None:
line = f" {task_name}: acc={acc:.4f}"
if acc_norm is not None:
line += f", acc_norm={acc_norm:.4f}"
print(f"[LM_EVAL] {line}")
else:
print(f"[LM_EVAL] {task_name}: {metrics}")
if skipped_tasks:
print(f"[LM_EVAL] Skipped tasks: {skipped_tasks}")
return {
"model_path": hf_model_path,
"tasks_requested": tasks,
"tasks_evaluated": evaluated_tasks,
"tasks_skipped": skipped_tasks,
"per_task_metrics": per_task_metrics,
"raw_results": raw_results,
"elapsed_sec": round(elapsed, 1),
}
# ---------------------------------------------------------------------------
# Pipeline mode — load model ONCE, run multiple fewshot settings sequentially
# ---------------------------------------------------------------------------
def _extract_per_task_metrics(raw_results: dict) -> dict[str, dict]:
"""Extract per-task metrics from lm_eval raw results."""
per_task_metrics: dict[str, dict] = {}
lm_results: dict[str, Any] = raw_results.get("results", {})
for task_name, task_data in lm_results.items():
if not isinstance(task_data, dict):
continue
metrics = {k: v for k, v in task_data.items() if k not in ("alias", "group")}
per_task_metrics[task_name] = metrics
return per_task_metrics
def run_lm_eval_tasks_pipeline(
hf_model_path: str,
tasks: list[str],
device: str,
fewshot_values: list[int],
output_dir: str = "",
output_prefix: str = "",
) -> dict:
"""Run lm-eval with multiple fewshot settings, loading the model ONCE.
This avoids the overhead of loading the model N times when running
0-shot then 5-shot on the same GPU.
Returns:
Dict with keys like "0shot", "5shot", each containing the same
structure as run_lm_eval_tasks().
"""
import json as _json
import lm_eval # type: ignore[import]
from lm_eval.models.huggingface import HFLM # type: ignore[import]
# --- GPU isolation (same as run_lm_eval_tasks) ---
gpu_index = int(device.split(":")[-1])
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_index)
_internal_device = "cuda:0"
print(
f"[LM_EVAL_PIPELINE] Loading model once on {device} "
f"for fewshot={fewshot_values}, tasks={tasks}",
flush=True,
)
# --- Load model ONCE ---
model_obj = HFLM(
pretrained=hf_model_path,
dtype="bfloat16",
device=_internal_device,
batch_size="auto",
)
# --- Validate tasks ---
try:
from lm_eval.tasks import TaskManager # type: ignore[import]
available_tasks = set(TaskManager().all_tasks)
except Exception:
available_tasks = set()
valid_tasks = [t for t in tasks if (not available_tasks) or (t in available_tasks)]
skipped_tasks = [t for t in tasks if t not in valid_tasks]
if not valid_tasks:
print("[LM_EVAL_PIPELINE] No valid tasks.", flush=True)
empty = {
"model_path": hf_model_path,
"tasks_requested": tasks,
"tasks_evaluated": [],
"tasks_skipped": skipped_tasks,
"per_task_metrics": {},
"raw_results": {},
"elapsed_sec": 0.0,
}
return {f"{n}shot": empty for n in fewshot_values}
# --- Run each fewshot setting, reusing model_obj ---
all_results: dict[str, Any] = {}
for num_fewshot in fewshot_values:
print(
f"[LM_EVAL_PIPELINE] Running {num_fewshot}-shot on {valid_tasks}...",
flush=True,
)
t0 = time.time()
try:
raw_results = lm_eval.simple_evaluate(
model=model_obj,
tasks=valid_tasks,
num_fewshot=num_fewshot,
)
per_task_metrics = _extract_per_task_metrics(raw_results)
elapsed = time.time() - t0
shot_result = {
"model_path": hf_model_path,
"tasks_requested": tasks,
"tasks_evaluated": list(valid_tasks),
"tasks_skipped": list(skipped_tasks),
"per_task_metrics": per_task_metrics,
"raw_results": raw_results,
"elapsed_sec": round(elapsed, 1),
}
print(
f"[LM_EVAL_PIPELINE] {num_fewshot}-shot complete in {elapsed:.1f}s",
flush=True,
)
except Exception as exc:
elapsed = time.time() - t0
shot_result = {
"model_path": hf_model_path,
"tasks_requested": tasks,
"tasks_evaluated": [],
"tasks_skipped": list(tasks),
"per_task_metrics": {},
"raw_results": {},
"elapsed_sec": round(elapsed, 1),
"error": str(exc),
}
print(
f"[LM_EVAL_PIPELINE] {num_fewshot}-shot FAILED: {exc}",
flush=True,
)
all_results[f"{num_fewshot}shot"] = shot_result
# Save intermediate result per fewshot
if output_dir:
shot_path = Path(output_dir) / f"{output_prefix}_{num_fewshot}shot.json"
with open(shot_path, "w", encoding="utf-8") as f:
_json.dump(shot_result, f, ensure_ascii=False, indent=2, default=str)
return all_results

206
source/eval/tasks/ppl_task.py Normal file
View File

@@ -0,0 +1,206 @@
"""
ppl_task.py — Sliding-window perplexity evaluation task.
Top-level functions for ProcessPoolExecutor (spawn) compatibility:
- eval_ppl_single(val_file, device, model=None) -> dict
- eval_ppl_multi(val_files, device) -> list[dict]
"""
from __future__ import annotations
import math
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 perplexity 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 functions (must be top-level for pickle / spawn compatibility)
# ---------------------------------------------------------------------------
def eval_ppl_single(val_file: str, device: str, model=None) -> dict:
"""Compute sliding-window perplexity for a single validation file.
Args:
val_file: Relative path under DATA_DIR, e.g. "3b_val.bin".
device: CUDA device string, e.g. "cuda:0".
model: Optional pre-loaded model. If None, loads from checkpoint.
Returns:
Dict with keys: name, file, n_tokens, n_eval_tokens, ppl,
bits_per_token, avg_nll, elapsed_sec, device.
"""
torch.cuda.set_device(int(device.split(":")[-1]))
data_path = DATA_DIR / val_file
if not data_path.exists():
raise FileNotFoundError(f"Validation file not found: {data_path}")
name = val_file.replace("_val.bin", "").replace(".bin", "")
own_model = model is None
if own_model:
print(f"[PPL {device}] Loading model for {name}...")
model = _load_model(device)
tokens = np.fromfile(str(data_path), dtype=np.uint16)
if len(tokens) == 0:
raise ValueError(f"Validation file is empty (0 tokens): {data_path}")
n_tokens = len(tokens)
print(f"[PPL {device}] {name}: {n_tokens:,} tokens, {n_tokens * 2 / 1e6:.1f} MB")
ds = SlidingWindowDataset(tokens, SEQ_LEN, STRIDE)
dl = DataLoader(
ds,
batch_size=BATCH_SIZE,
shuffle=False,
num_workers=4,
pin_memory=True,
)
total_nll = 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)
loss_flat = F.cross_entropy(
logits.view(-1, logits.size(-1)),
tgt.view(-1),
reduction="none",
)
loss_flat = loss_flat.view(mask.shape)
nll = (loss_flat * mask.float()).sum().item()
cnt = mask.sum().item()
total_nll += nll
total_count += cnt
if (batch_idx + 1) % 50 == 0:
running_ppl = (
math.exp(total_nll / total_count) if total_count > 0 else float("inf")
)
elapsed = time.time() - t0
print(
f"[PPL {device}] {name}: batch {batch_idx + 1}/{len(dl)}, "
f"running PPL={running_ppl:.4f}, {elapsed:.0f}s"
)
avg_nll = total_nll / total_count if total_count > 0 else 0.0
ppl = math.exp(avg_nll)
bpt = avg_nll / math.log(2)
elapsed = time.time() - t0
result: dict = {
"name": name,
"file": val_file,
"n_tokens": int(n_tokens),
"n_eval_tokens": int(total_count),
"ppl": round(ppl, 4),
"bits_per_token": round(bpt, 4),
"avg_nll": round(avg_nll, 6),
"elapsed_sec": round(elapsed, 1),
"device": device,
}
print(
f"[PPL {device}] DONE {name}: PPL={ppl:.4f}, BPT={bpt:.4f}, {elapsed:.1f}s"
)
return result
def eval_ppl_multi(val_files: list[str], device: str) -> list[dict]:
"""Compute PPL for multiple val files on a single GPU, loading model once.
Args:
val_files: List of relative paths under DATA_DIR.
device: CUDA device string.
Returns:
List of result dicts (one per file), in the same order as val_files.
"""
torch.cuda.set_device(int(device.split(":")[-1]))
print(f"[PPL_MULTI {device}] Loading model once for {len(val_files)} files...")
model = _load_model(device)
results: list[dict] = []
for val_file in val_files:
result = eval_ppl_single(val_file, device, model=model)
results.append(result)
return results

250
source/eval/tasks/task_runner.py Normal file
View File

@@ -0,0 +1,250 @@
"""
task_runner.py — Thin CLI entry point for subprocess GPU workers.
Usage:
CUDA_VISIBLE_DEVICES=5 python eval/tasks/task_runner.py \
--task calibration --gpu-id 5 --output /path/to/result.json
"""
import argparse
import json
import os
import sys
import traceback
from pathlib import Path
# ---------------------------------------------------------------------------
# Project root on sys.path
# ---------------------------------------------------------------------------
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
# ---------------------------------------------------------------------------
# NUMA affinity helper
# ---------------------------------------------------------------------------
def _set_numa_affinity(gpu_id: int) -> None:
"""Pin the process to the NUMA node that owns the given GPU.
GPU 0-3 → cores 0-35 (NUMA node 0)
GPU 4-7 → cores 36-71 (NUMA node 1)
"""
try:
import os
if gpu_id <= 3:
cores = list(range(0, 36))
else:
cores = list(range(36, 72))
# os.sched_setaffinity is available on Linux
os.sched_setaffinity(0, cores)
print(
f"[TASK_RUNNER gpu_id={gpu_id}] NUMA affinity set: cores {cores[0]}-{cores[-1]}",
flush=True,
)
except Exception as exc:
# Non-fatal — just warn and continue
print(
f"[TASK_RUNNER gpu_id={gpu_id}] WARNING: could not set NUMA affinity: {exc}",
flush=True,
)
# ---------------------------------------------------------------------------
# Task dispatch
# ---------------------------------------------------------------------------
VALID_TASKS = {
"ppl_single",
"ppl_multi",
"calibration",
"token_nll",
"calib_nll",
"generation",
"repetition_grid",
"lm_eval",
}
def _run_task(args: argparse.Namespace) -> dict:
task = args.task
device = "cuda:0" # CUDA_VISIBLE_DEVICES already set by parent
if task == "ppl_single":
if not args.val_file:
raise ValueError("--val-file is required for ppl_single task")
from eval.tasks.ppl_task import eval_ppl_single
result = eval_ppl_single(args.val_file, device)
elif task == "ppl_multi":
if not args.val_files:
raise ValueError("--val-files is required for ppl_multi task")
val_files_list = [f.strip() for f in args.val_files.split(",") if f.strip()]
from eval.tasks.ppl_task import eval_ppl_multi
result = eval_ppl_multi(val_files_list, device)
elif task == "calibration":
from eval.tasks.calibration_task import eval_calibration
result = eval_calibration(device)
elif task == "token_nll":
from eval.tasks.token_nll_task import eval_token_nll
result = eval_token_nll(device)
elif task == "calib_nll":
from eval.tasks.calibration_task import eval_calibration
from eval.tasks.token_nll_task import eval_token_nll
calib_result = eval_calibration(device)
nll_result = eval_token_nll(device)
result = {"calibration": calib_result, "token_nll": nll_result}
elif task == "generation":
from eval.tasks.generation_task import eval_generation
result = eval_generation(device)
elif task == "repetition_grid":
from eval.tasks.generation_task import eval_repetition_grid
result = eval_repetition_grid(device)
elif task == "lm_eval":
if not args.hf_model_path:
raise ValueError("--hf-model-path is required for lm_eval task")
if not args.lm_eval_tasks:
raise ValueError("--lm-eval-tasks is required for lm_eval task")
tasks_list = [t.strip() for t in args.lm_eval_tasks.split(",") if t.strip()]
if args.fewshot_list:
# Pipeline mode: load model once, run multiple fewshot settings
fewshot_values = [int(x.strip()) for x in args.fewshot_list.split(",")]
from eval.tasks.lm_eval_task import run_lm_eval_tasks_pipeline
result = run_lm_eval_tasks_pipeline(
args.hf_model_path,
tasks_list,
device,
fewshot_values,
output_dir=str(Path(args.output).parent),
output_prefix=Path(args.output).stem,
)
else:
from eval.tasks.lm_eval_task import run_lm_eval_tasks
result = run_lm_eval_tasks(
args.hf_model_path,
tasks_list,
device,
num_fewshot=args.num_fewshot,
)
else:
raise ValueError(f"Unknown task: {task!r}. Valid tasks: {sorted(VALID_TASKS)}")
return result
# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Thin CLI entry point for subprocess GPU eval workers."
)
parser.add_argument(
"--task",
required=True,
choices=sorted(VALID_TASKS),
help="Eval task to run.",
)
parser.add_argument(
"--gpu-id",
type=int,
required=True,
help="Original GPU ID (used for NUMA affinity only).",
)
parser.add_argument(
"--output",
required=True,
help="Path to write JSON result file.",
)
# --- ppl_single ---
parser.add_argument(
"--val-file",
default=None,
help="Single validation filename (for ppl_single).",
)
# --- ppl_multi ---
parser.add_argument(
"--val-files",
default=None,
help="Comma-separated validation filenames (for ppl_multi).",
)
# --- lm_eval ---
parser.add_argument(
"--hf-model-path",
default=None,
help="HuggingFace model directory (for lm_eval).",
)
parser.add_argument(
"--lm-eval-tasks",
default=None,
help="Comma-separated lm-eval task names (for lm_eval).",
)
parser.add_argument(
"--num-fewshot",
type=int,
default=0,
help="Number of few-shot examples (for lm_eval). Default: 0.",
)
parser.add_argument(
"--fewshot-list",
default=None,
help="Comma-separated fewshot values to run sequentially, e.g. '0,5'. "
"Model is loaded once and reused. Overrides --num-fewshot.",
)
return parser.parse_args()
# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main() -> None:
args = _parse_args()
gpu_id = args.gpu_id
task_name = args.task
output_path = args.output
print(f"[TASK_RUNNER gpu_id={gpu_id}] Starting task={task_name}", flush=True)
# Set NUMA affinity early
_set_numa_affinity(gpu_id)
exit_code = 0
try:
result = _run_task(args)
payload = result
except Exception as exc:
tb_str = traceback.format_exc()
print(
f"[TASK_RUNNER gpu_id={gpu_id}] ERROR in task={task_name}:\n{tb_str}",
file=sys.stderr,
flush=True,
)
payload = {"error": str(exc), "traceback": tb_str}
exit_code = 1
# Write result JSON
output_path_obj = Path(output_path)
output_path_obj.parent.mkdir(parents=True, exist_ok=True)
with open(output_path_obj, "w", encoding="utf-8") as fh:
json.dump(payload, fh, ensure_ascii=False, indent=2, default=str)
print(
f"[TASK_RUNNER gpu_id={gpu_id}] Done. Result saved to {output_path}",
flush=True,
)
sys.exit(exit_code)
if __name__ == "__main__":
main()

216
source/eval/tasks/token_nll_task.py Normal file
View File

@@ -0,0 +1,216 @@
"""
token_nll_task.py — Token-level NLL distribution analysis.
Top-level function for ProcessPoolExecutor (spawn) compatibility:
- eval_token_nll(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_token_nll(device: str, n_tokens: int = 50000) -> dict:
"""Analyse the per-token NLL distribution on 3b_val.bin.
Collects the NLL of every valid (unmasked) token and computes summary
statistics and percentile breakdowns, as well as the fraction of
"high-loss" tokens that may indicate out-of-distribution content.
Args:
device: CUDA device string, e.g. "cuda:6".
n_tokens: Number of tokens to process (first n_tokens of 3b_val.bin).
Returns:
Dict with keys:
- n_eval_tokens: number of tokens included in stats
- nll_mean: mean token NLL
- nll_std: standard deviation of token NLL
- nll_median: 50th-percentile NLL
- nll_percentiles: dict mapping percentile label to value
(keys: p5, p25, p75, p95, p99)
- high_loss_fraction_5: fraction of tokens with NLL > 5.0
- high_loss_fraction_10: fraction of tokens with NLL > 10.0
- elapsed_sec: wall-clock time
"""
torch.cuda.set_device(int(device.split(":")[-1]))
print(f"[NLL {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"[NLL {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=4,
pin_memory=True,
)
all_nlls: list[np.ndarray] = []
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)
# Per-token NLL — shape (batch, seq_len)
per_token_nll = F.cross_entropy(
logits.view(-1, logits.size(-1)),
tgt.view(-1),
reduction="none",
ignore_index=-100,
).view(mask.shape)
# Collect only valid (unmasked) positions
valid_nll = per_token_nll[mask].float().cpu().numpy()
if len(valid_nll) > 0:
all_nlls.append(valid_nll)
if (batch_idx + 1) % 50 == 0:
n_collected = sum(len(a) for a in all_nlls)
elapsed = time.time() - t0
print(
f"[NLL {device}] batch {batch_idx + 1}/{len(dl)}, "
f"tokens collected={n_collected:,}, {elapsed:.0f}s"
)
elapsed = time.time() - t0
if all_nlls:
nll_arr = np.concatenate(all_nlls)
else:
nll_arr = np.array([], dtype=np.float32)
n_eval = len(nll_arr)
if n_eval > 0:
nll_mean = float(np.mean(nll_arr))
nll_std = float(np.std(nll_arr))
nll_median = float(np.median(nll_arr))
percentiles = {
"p5": round(float(np.percentile(nll_arr, 5)), 4),
"p25": round(float(np.percentile(nll_arr, 25)), 4),
"p75": round(float(np.percentile(nll_arr, 75)), 4),
"p95": round(float(np.percentile(nll_arr, 95)), 4),
"p99": round(float(np.percentile(nll_arr, 99)), 4),
}
high_loss_5 = float(np.mean(nll_arr > 5.0))
high_loss_10 = float(np.mean(nll_arr > 10.0))
else:
nll_mean = nll_std = nll_median = 0.0
percentiles = {"p5": 0.0, "p25": 0.0, "p75": 0.0, "p95": 0.0, "p99": 0.0}
high_loss_5 = high_loss_10 = 0.0
result: dict = {
"n_eval_tokens": int(n_eval),
"nll_mean": round(nll_mean, 4),
"nll_std": round(nll_std, 4),
"nll_median": round(nll_median, 4),
"nll_percentiles": {k: round(v, 4) for k, v in percentiles.items()},
"high_loss_fraction_5": round(high_loss_5, 6),
"high_loss_fraction_10": round(high_loss_10, 6),
"elapsed_sec": round(elapsed, 1),
}
print(
f"[NLL {device}] DONE n={n_eval:,}, "
f"mean={nll_mean:.4f}, std={nll_std:.4f}, "
f"median={nll_median:.4f}, "
f"high_loss(>5)={high_loss_5:.2%}, "
f"high_loss(>10)={high_loss_10:.2%}, "
f"{elapsed:.1f}s"
)
return result