411 lines
14 KiB
Python
411 lines
14 KiB
Python
"""
|
||
3B BASE 모델 생성 품질 + 반복률 종합 분석 스크립트.
|
||
|
||
Part 1: 10개 프롬프트 × 3 온도 → 자유 생성 텍스트 저장
|
||
Part 2: 파라미터 그리드 서치 → 반복률 분석 JSON 저장
|
||
|
||
BASE 모델용 completion-style 프롬프트 사용.
|
||
|
||
Usage:
|
||
cd /PROJECT/0325120031_A/ghong/taketimes/llm-bang
|
||
python eval/analyze_3b_generation.py \
|
||
--checkpoint checkpoints/korean_3b_fp8_run1/checkpoint-0057000 \
|
||
--device cuda:1
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
from collections import Counter
|
||
|
||
import torch
|
||
import torch.nn.functional as F
|
||
|
||
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||
if str(_PROJECT_ROOT) not in sys.path:
|
||
sys.path.insert(0, str(_PROJECT_ROOT))
|
||
|
||
from model.transformer import LLM
|
||
from tokenizers import Tokenizer
|
||
|
||
try:
|
||
import transformer_engine.pytorch as te
|
||
from transformer_engine.common.recipe import MXFP8BlockScaling
|
||
HAS_TE = True
|
||
except ImportError:
|
||
te = None
|
||
HAS_TE = False
|
||
|
||
|
||
def fp8_inference_context():
|
||
"""Return the appropriate inference context manager for FP8 models."""
|
||
if HAS_TE:
|
||
return te.fp8_autocast(enabled=True, fp8_recipe=MXFP8BlockScaling())
|
||
import contextlib
|
||
return contextlib.nullcontext()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# BASE model completion-style prompts (10 prompts)
|
||
# ---------------------------------------------------------------------------
|
||
BASE_PROMPTS = [
|
||
"대한민국의 수도는",
|
||
"인공지능이란",
|
||
"한국의 전통 음식 중에서",
|
||
"지구 온난화의 주요 원인은",
|
||
"프로그래밍을 배우려면",
|
||
"조선시대에는",
|
||
"물리학에서 에너지란",
|
||
"한국어는 세계에서",
|
||
"경제 성장을 위해서는",
|
||
"우주 탐사의 역사를 보면",
|
||
]
|
||
|
||
# Subset for repetition grid (3 prompts to keep runtime reasonable)
|
||
GRID_PROMPTS = BASE_PROMPTS[:3]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Sampling utilities
|
||
# ---------------------------------------------------------------------------
|
||
def top_p_filtering(logits, top_p=0.9, top_k=0):
|
||
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
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Repetition metrics
|
||
# ---------------------------------------------------------------------------
|
||
def compute_ngram_repetition(tokens: list[str], n: int) -> float:
|
||
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_all_repetition_metrics(text: str) -> dict:
|
||
tokens = text.split()
|
||
return {
|
||
f"{n}gram_rep": compute_ngram_repetition(tokens, n)
|
||
for n in [1, 2, 3, 4]
|
||
}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Generation (greedy or sampling, with optional rep penalty + no_repeat_ngram)
|
||
# ---------------------------------------------------------------------------
|
||
@torch.inference_mode()
|
||
def generate_text(
|
||
model,
|
||
tokenizer,
|
||
prompt: str,
|
||
max_new_tokens: int = 256,
|
||
temperature: float = 0.8,
|
||
top_p: float = 0.9,
|
||
top_k: int = 50,
|
||
repetition_penalty: float = 1.0,
|
||
no_repeat_ngram_size: int = 0,
|
||
device: str = "cuda:1",
|
||
) -> tuple[str, int, bool]:
|
||
"""
|
||
Returns: (generated_text, num_new_tokens, hit_eos)
|
||
MXFP8 requires sequence length divisible by 32; we right-pad before each
|
||
forward pass but use the logit at the true last real position.
|
||
"""
|
||
model.eval()
|
||
raw_ids = tokenizer.encode(prompt).ids
|
||
eos_id = tokenizer.token_to_id("</s>")
|
||
pad_id = tokenizer.token_to_id("<pad>") or 0
|
||
|
||
# Keep an unpadded running sequence; pad only for the forward pass
|
||
real_ids: list[int] = list(raw_ids)
|
||
new_token_ids: list[int] = []
|
||
hit_eos = False
|
||
|
||
ctx = fp8_inference_context()
|
||
with ctx:
|
||
for _ in range(max_new_tokens):
|
||
real_len = len(real_ids)
|
||
# Pad to next multiple of 32 for MXFP8
|
||
pad_to = ((real_len + 31) // 32) * 32
|
||
padded = real_ids + [pad_id] * (pad_to - real_len)
|
||
x = torch.tensor([padded], dtype=torch.long, device=device)
|
||
|
||
logits_all, _ = model(x)
|
||
# Logit at the last REAL token (index real_len - 1)
|
||
logits = logits_all[:, real_len - 1, :].clone() # [1, V]
|
||
|
||
# Repetition penalty
|
||
if repetition_penalty != 1.0:
|
||
for token_id in set(real_ids):
|
||
if logits[0, token_id] > 0:
|
||
logits[0, token_id] /= repetition_penalty
|
||
else:
|
||
logits[0, token_id] *= repetition_penalty
|
||
|
||
# No-repeat n-gram blocking
|
||
if no_repeat_ngram_size > 0 and real_len >= no_repeat_ngram_size:
|
||
for i in range(real_len - no_repeat_ngram_size + 1):
|
||
ngram = tuple(real_ids[i:i + no_repeat_ngram_size - 1])
|
||
last_ngram = tuple(real_ids[-(no_repeat_ngram_size - 1):])
|
||
if ngram == last_ngram:
|
||
logits[0, real_ids[i + no_repeat_ngram_size - 1]] = float("-inf")
|
||
|
||
# Decode strategy
|
||
if temperature == 0.0:
|
||
next_token_id = int(logits.argmax(dim=-1).item())
|
||
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_token_id = int(torch.multinomial(probs, num_samples=1).item())
|
||
|
||
real_ids.append(next_token_id)
|
||
new_token_ids.append(next_token_id)
|
||
|
||
if eos_id is not None and next_token_id == eos_id:
|
||
hit_eos = True
|
||
break
|
||
|
||
generated_text = tokenizer.decode(new_token_ids)
|
||
return generated_text, len(new_token_ids), hit_eos
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Part 1: Free generation (10 prompts × 3 temps)
|
||
# ---------------------------------------------------------------------------
|
||
def run_free_generation(model, tokenizer, device, output_path: Path):
|
||
temperatures = [0.0, 0.7, 1.0]
|
||
results = []
|
||
|
||
print("\n" + "=" * 70)
|
||
print(" PART 1: FREE GENERATION (10 prompts × 3 temperatures)")
|
||
print("=" * 70)
|
||
|
||
for temp in temperatures:
|
||
print(f"\n--- Temperature: {temp} ---")
|
||
for prompt in BASE_PROMPTS:
|
||
t0 = time.time()
|
||
gen_text, n_tokens, hit_eos = generate_text(
|
||
model, tokenizer, prompt,
|
||
max_new_tokens=256,
|
||
temperature=temp,
|
||
top_p=0.9,
|
||
top_k=50,
|
||
device=device,
|
||
)
|
||
elapsed = time.time() - t0
|
||
metrics = compute_all_repetition_metrics(gen_text)
|
||
|
||
entry = {
|
||
"prompt": prompt,
|
||
"temperature": temp,
|
||
"generation": gen_text,
|
||
"n_new_tokens": n_tokens,
|
||
"hit_eos": hit_eos,
|
||
"elapsed_sec": round(elapsed, 2),
|
||
**metrics,
|
||
}
|
||
results.append(entry)
|
||
|
||
# Print summary
|
||
preview = gen_text[:120].replace("\n", "\\n")
|
||
print(f" [{temp}] {prompt!r}")
|
||
print(f" → {preview}...")
|
||
print(f" tokens={n_tokens}, eos={hit_eos}, 3gram_rep={metrics['3gram_rep']*100:.1f}%")
|
||
|
||
# Save text version for easy reading
|
||
txt_path = output_path.parent / "3b_generation_results.txt"
|
||
with open(txt_path, "w", encoding="utf-8") as f:
|
||
for r in results:
|
||
f.write(f"\n{'='*60}\n")
|
||
f.write(f"Temperature: {r['temperature']}\n")
|
||
f.write(f"Prompt: {r['prompt']}\n")
|
||
f.write(f"Generated ({r['n_new_tokens']} tokens, eos={r['hit_eos']}):\n")
|
||
f.write(r["generation"] + "\n")
|
||
f.write(f"3gram_rep={r['3gram_rep']*100:.1f}% | 4gram_rep={r['4gram_rep']*100:.1f}%\n")
|
||
|
||
print(f"\n[Part 1] Saved text to: {txt_path}")
|
||
return results
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Part 2: Repetition parameter grid search
|
||
# ---------------------------------------------------------------------------
|
||
PARAM_GRID = []
|
||
|
||
# Generate grid: temp × rep_penalty × no_repeat_ngram × top_p
|
||
for temp in [0.7, 0.9, 1.0]:
|
||
for rep in [1.0, 1.1, 1.2, 1.3]:
|
||
for ngram in [0, 3, 4]:
|
||
for top_p in [0.9, 0.95]:
|
||
name = f"t{temp}_r{rep}_ng{ngram}_tp{top_p}"
|
||
PARAM_GRID.append({
|
||
"name": name,
|
||
"temperature": temp,
|
||
"repetition_penalty": rep,
|
||
"no_repeat_ngram_size": ngram,
|
||
"top_p": top_p,
|
||
"top_k": 50,
|
||
})
|
||
|
||
|
||
def run_repetition_analysis(model, tokenizer, device, output_path: Path):
|
||
print("\n" + "=" * 70)
|
||
print(f" PART 2: REPETITION ANALYSIS ({len(PARAM_GRID)} configs × {len(GRID_PROMPTS)} prompts)")
|
||
print("=" * 70)
|
||
|
||
all_results = {}
|
||
eos_counts = {}
|
||
|
||
for params in PARAM_GRID:
|
||
name = params["name"]
|
||
rep_scores = {n: [] for n in [1, 2, 3, 4]}
|
||
eos_hits = 0
|
||
token_counts = []
|
||
generations = []
|
||
|
||
for prompt in GRID_PROMPTS:
|
||
gen_text, n_tokens, hit_eos = generate_text(
|
||
model, tokenizer, prompt,
|
||
max_new_tokens=256,
|
||
temperature=params["temperature"],
|
||
top_p=params["top_p"],
|
||
top_k=params["top_k"],
|
||
repetition_penalty=params["repetition_penalty"],
|
||
no_repeat_ngram_size=params["no_repeat_ngram_size"],
|
||
device=device,
|
||
)
|
||
metrics = compute_all_repetition_metrics(gen_text)
|
||
for n in [1, 2, 3, 4]:
|
||
rep_scores[n].append(metrics[f"{n}gram_rep"])
|
||
if hit_eos:
|
||
eos_hits += 1
|
||
token_counts.append(n_tokens)
|
||
generations.append({
|
||
"prompt": prompt,
|
||
"generation": gen_text[:300],
|
||
"n_tokens": n_tokens,
|
||
"hit_eos": hit_eos,
|
||
**{f"{n}gram_rep": round(metrics[f"{n}gram_rep"], 4) for n in [1, 2, 3, 4]},
|
||
})
|
||
|
||
n_prompts = len(GRID_PROMPTS)
|
||
avg_reps = {f"avg_{n}gram_rep": round(sum(rep_scores[n]) / n_prompts, 4) for n in [1, 2, 3, 4]}
|
||
eos_rate = eos_hits / n_prompts
|
||
avg_tokens = sum(token_counts) / n_prompts
|
||
|
||
all_results[name] = {
|
||
"params": {k: v for k, v in params.items() if k != "name"},
|
||
**avg_reps,
|
||
"eos_rate": round(eos_rate, 4),
|
||
"avg_tokens": round(avg_tokens, 1),
|
||
"generations": generations,
|
||
}
|
||
|
||
print(f" {name:<45} 3g={avg_reps['avg_3gram_rep']*100:.1f}% eos={eos_rate:.0%} tok={avg_tokens:.0f}")
|
||
|
||
# Save JSON
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
with open(output_path, "w", encoding="utf-8") as f:
|
||
json.dump(all_results, f, ensure_ascii=False, indent=2)
|
||
|
||
# Print ranked summary
|
||
print(f"\n{'='*70}")
|
||
print(" RANKED BY 3-GRAM REPETITION RATE")
|
||
print(f"{'='*70}")
|
||
print(f" {'Config':<45} {'3gram':>7} {'eos':>6} {'tokens':>7}")
|
||
print(f" {'-'*45} {'-'*7} {'-'*6} {'-'*7}")
|
||
sorted_results = sorted(all_results.items(), key=lambda x: x[1]["avg_3gram_rep"])
|
||
for name, res in sorted_results[:20]: # top 20
|
||
print(
|
||
f" {name:<45} {res['avg_3gram_rep']*100:>6.1f}%"
|
||
f" {res['eos_rate']:>5.0%} {res['avg_tokens']:>7.0f}"
|
||
)
|
||
|
||
print(f"\n[Part 2] Saved JSON to: {output_path}")
|
||
return all_results
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main
|
||
# ---------------------------------------------------------------------------
|
||
def main():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument(
|
||
"--checkpoint",
|
||
default="checkpoints/korean_3b_fp8_run1/checkpoint-0057000",
|
||
)
|
||
parser.add_argument("--device", default="cuda:1")
|
||
parser.add_argument("--output_dir", default="eval/outputs")
|
||
args = parser.parse_args()
|
||
|
||
ckpt = Path(args.checkpoint)
|
||
if not ckpt.is_absolute():
|
||
ckpt = _PROJECT_ROOT / ckpt
|
||
|
||
# Set default CUDA device BEFORE loading — required for TE MXFP8 device routing
|
||
device_id = int(args.device.split(":")[-1]) if ":" in args.device else 0
|
||
torch.cuda.set_device(device_id)
|
||
|
||
print(f"Loading model from: {ckpt}")
|
||
model = LLM.from_pretrained(str(ckpt)).cuda(device_id).to(dtype=torch.bfloat16)
|
||
model.eval()
|
||
n_params = sum(p.numel() for p in model.parameters())
|
||
print(f"Model loaded. Params: {n_params / 1e9:.2f}B")
|
||
|
||
tok_path = ckpt / "tokenizer.json"
|
||
if not tok_path.exists():
|
||
tok_path = _PROJECT_ROOT / "tokenizer" / "korean_sp" / "tokenizer.json"
|
||
print(f"Loading tokenizer from: {tok_path}")
|
||
tokenizer = Tokenizer.from_file(str(tok_path))
|
||
|
||
output_dir = _PROJECT_ROOT / args.output_dir
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# Part 1: free generation
|
||
free_gen_results = run_free_generation(
|
||
model, tokenizer, args.device, output_dir / "3b_generation_results.txt"
|
||
)
|
||
|
||
# Save Part 1 JSON
|
||
gen_json_path = output_dir / "3b_generation_results.json"
|
||
with open(gen_json_path, "w", encoding="utf-8") as f:
|
||
json.dump(free_gen_results, f, ensure_ascii=False, indent=2)
|
||
print(f"[Part 1] JSON saved: {gen_json_path}")
|
||
|
||
# Part 2: repetition analysis
|
||
rep_json_path = output_dir / "3b_repetition_analysis.json"
|
||
run_repetition_analysis(model, tokenizer, args.device, rep_json_path)
|
||
|
||
print("\nDone.")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|