531 lines
20 KiB
Python
531 lines
20 KiB
Python
"""
|
|
FRANKENSTALLM 3B — SFT Evaluation Pipeline Orchestrator
|
|
=========================================================
|
|
|
|
Evaluates the SFT checkpoint across 6 dimensions and generates a
|
|
comparison report against the Base model results.
|
|
|
|
Runs 4 phases sequentially:
|
|
Phase 0 — Convert SFT checkpoint to HuggingFace format
|
|
Phase 1 — Internal evaluation across 8 GPUs (PPL, Calibration, Generation)
|
|
Phase 2 — Standard benchmarks via lm-eval-harness (8 GPU parallel)
|
|
Phase 3 — Base vs SFT comparison report generation
|
|
|
|
Usage:
|
|
python eval/sft_eval_pipeline.py
|
|
python eval/sft_eval_pipeline.py --dry-run
|
|
python eval/sft_eval_pipeline.py --skip-phase0 --skip-phase2
|
|
python eval/sft_eval_pipeline.py --skip-phase0 --hf-model-path eval/outputs/hf_3b_sft_best
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import multiprocessing as mp
|
|
import os
|
|
import sys
|
|
import time
|
|
import traceback
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Project root
|
|
# ---------------------------------------------------------------------------
|
|
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
|
if str(_PROJECT_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(_PROJECT_ROOT))
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# SFT checkpoint and Base results paths
|
|
# ---------------------------------------------------------------------------
|
|
SFT_CHECKPOINT = str(
|
|
_PROJECT_ROOT / "checkpoints" / "korean_3b_sft_v2" / "checkpoint-best"
|
|
)
|
|
# SFT tokenizer lives alongside the SFT checkpoint
|
|
SFT_TOKENIZER = str(
|
|
_PROJECT_ROOT / "checkpoints" / "korean_3b_sft_v2" / "tokenizer.json"
|
|
)
|
|
# Fallback tokenizer if SFT-specific one doesn't exist
|
|
_FALLBACK_TOKENIZER = str(
|
|
_PROJECT_ROOT / "tokenizer" / "korean_sp" / "tokenizer.json"
|
|
)
|
|
BASE_RESULTS_DIR = _PROJECT_ROOT / "eval" / "outputs" / "3b_reeval_20260305_1451"
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Import shared infrastructure from full_eval_pipeline
|
|
# ---------------------------------------------------------------------------
|
|
from eval.full_eval_pipeline import (
|
|
_bar,
|
|
_build_phase1_tasks,
|
|
_build_phase2_tasks,
|
|
_fmt_seconds,
|
|
_make_output_dir,
|
|
_NUMA_CORES,
|
|
_print_banner,
|
|
_print_phase_header,
|
|
_save_json,
|
|
_spawn_task,
|
|
_wait_and_collect,
|
|
run_phase0,
|
|
SEQ_LEN,
|
|
STRIDE,
|
|
BATCH_SIZE,
|
|
DATA_DIR,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Logging
|
|
# ---------------------------------------------------------------------------
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s [%(levelname)s] %(message)s",
|
|
datefmt="%Y-%m-%d %H:%M:%S",
|
|
)
|
|
logger = logging.getLogger("sft_eval")
|
|
|
|
|
|
# ===========================================================================
|
|
# Override: spawn tasks with SFT environment variables
|
|
# ===========================================================================
|
|
|
|
def _spawn_sft_task(
|
|
task_name: str,
|
|
gpu_id: int,
|
|
output_path: Path,
|
|
label: str,
|
|
checkpoint: str,
|
|
tokenizer: str,
|
|
use_chat_template: bool = False,
|
|
extra_args: Optional[Dict[str, str]] = None,
|
|
) -> tuple:
|
|
"""Spawn a subprocess task with SFT checkpoint via environment variables."""
|
|
cmd = [
|
|
sys.executable,
|
|
str(_PROJECT_ROOT / "eval" / "tasks" / "task_runner.py"),
|
|
"--task", task_name,
|
|
"--gpu-id", str(gpu_id),
|
|
"--output", str(output_path),
|
|
]
|
|
if extra_args:
|
|
for k, v in extra_args.items():
|
|
cmd.extend([k, v])
|
|
|
|
env = os.environ.copy()
|
|
env["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
|
|
env["EVAL_CHECKPOINT"] = checkpoint
|
|
env["EVAL_TOKENIZER"] = tokenizer
|
|
if use_chat_template:
|
|
env["USE_CHAT_TEMPLATE"] = "1"
|
|
|
|
import subprocess
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
log_path = output_path.with_suffix(".log")
|
|
log_file = open(log_path, "w")
|
|
|
|
logger.info(" Spawning: %s (GPU %d) [SFT]", label, gpu_id)
|
|
proc = subprocess.Popen(
|
|
cmd,
|
|
stdout=log_file,
|
|
stderr=subprocess.STDOUT,
|
|
env=env,
|
|
cwd=str(_PROJECT_ROOT),
|
|
)
|
|
return proc, label, output_path, log_file
|
|
|
|
|
|
# ===========================================================================
|
|
# Phase 1 — Internal Evaluation (SFT variant)
|
|
# ===========================================================================
|
|
|
|
def run_sft_phase1(
|
|
output_dir: Path,
|
|
gpu_ids: List[int],
|
|
checkpoint: str,
|
|
tokenizer: str,
|
|
) -> Dict[str, Any]:
|
|
"""Run internal eval tasks with SFT checkpoint, chat template enabled for gen tasks."""
|
|
task_descriptors = _build_phase1_tasks(gpu_ids)
|
|
processes = []
|
|
|
|
for desc in task_descriptors:
|
|
is_gen_task = desc["task"] in ("generation", "repetition_grid")
|
|
out_path = output_dir / f"phase1_{desc['task']}_gpu{desc['gpu_id']}.json"
|
|
proc_info = _spawn_sft_task(
|
|
task_name=desc["task"],
|
|
gpu_id=desc["gpu_id"],
|
|
output_path=out_path,
|
|
label=desc["label"],
|
|
checkpoint=checkpoint,
|
|
tokenizer=tokenizer,
|
|
use_chat_template=is_gen_task,
|
|
extra_args=desc.get("extra_args"),
|
|
)
|
|
processes.append(proc_info)
|
|
|
|
results = _wait_and_collect(processes)
|
|
|
|
phase1_out = output_dir / "phase1_results.json"
|
|
_save_json(results, phase1_out)
|
|
logger.info(" Phase 1 results saved: %s", phase1_out)
|
|
|
|
# Save generation samples separately
|
|
gen_samples: Dict[str, Any] = {}
|
|
for label, result in results.items():
|
|
if isinstance(result, dict) and "error" not in result:
|
|
if "Generation" in label:
|
|
gen_samples["generation"] = result
|
|
elif "Repetition" in label:
|
|
gen_samples["repetition_grid"] = result
|
|
if gen_samples:
|
|
gen_out = output_dir / "generation_samples.json"
|
|
_save_json(gen_samples, gen_out)
|
|
logger.info(" Generation samples saved: %s", gen_out)
|
|
|
|
return results
|
|
|
|
|
|
# ===========================================================================
|
|
# Phase 2 — lm-eval Benchmarks (SFT variant — uses HF model)
|
|
# ===========================================================================
|
|
|
|
def _spawn_sft_phase2_batch(
|
|
hf_model_path: Path,
|
|
output_dir: Path,
|
|
gpu_task_list: list,
|
|
num_fewshot: int,
|
|
label_suffix: str,
|
|
checkpoint: str,
|
|
tokenizer: str,
|
|
) -> Dict[str, Any]:
|
|
"""Spawn Phase 2 subprocesses with SFT environment."""
|
|
processes = []
|
|
|
|
for gpu_id, task_names, label in gpu_task_list:
|
|
fewshot_label = f"[{num_fewshot}-shot] {label}"
|
|
out_path = output_dir / f"phase2_gpu{gpu_id}_{num_fewshot}shot{label_suffix}.json"
|
|
proc_info = _spawn_sft_task(
|
|
task_name="lm_eval",
|
|
gpu_id=gpu_id,
|
|
output_path=out_path,
|
|
label=fewshot_label,
|
|
checkpoint=checkpoint,
|
|
tokenizer=tokenizer,
|
|
extra_args={
|
|
"--hf-model-path": str(hf_model_path),
|
|
"--lm-eval-tasks": ",".join(task_names),
|
|
"--num-fewshot": str(num_fewshot),
|
|
},
|
|
)
|
|
processes.append(proc_info)
|
|
|
|
return _wait_and_collect(processes)
|
|
|
|
|
|
def run_sft_phase2(
|
|
hf_model_path: Path,
|
|
output_dir: Path,
|
|
gpu_ids: List[int],
|
|
checkpoint: str,
|
|
tokenizer: str,
|
|
) -> Dict[str, Any]:
|
|
"""Run lm-eval benchmarks for SFT model (0-shot + 5-shot)."""
|
|
gpu_task_list = _build_phase2_tasks(gpu_ids)
|
|
|
|
logger.info(" Running 0-shot benchmarks on %d GPUs ...", len(gpu_ids))
|
|
results = _spawn_sft_phase2_batch(
|
|
hf_model_path, output_dir, gpu_task_list, 0, "",
|
|
checkpoint, tokenizer,
|
|
)
|
|
logger.info(" Phase 2 (0-shot) complete.")
|
|
|
|
# 5-shot
|
|
logger.info(" Attempting 5-shot benchmarks ...")
|
|
try:
|
|
five_shot_results = _spawn_sft_phase2_batch(
|
|
hf_model_path, output_dir, gpu_task_list, 5, "_5shot",
|
|
checkpoint, tokenizer,
|
|
)
|
|
logger.info(" Phase 2 (5-shot) complete.")
|
|
except Exception:
|
|
logger.warning(" 5-shot failed (non-fatal): %s", traceback.format_exc())
|
|
five_shot_results = {"error": traceback.format_exc()}
|
|
results["5shot"] = five_shot_results
|
|
|
|
phase2_out = output_dir / "phase2_results.json"
|
|
_save_json(results, phase2_out)
|
|
logger.info(" Phase 2 results saved: %s", phase2_out)
|
|
return results
|
|
|
|
|
|
# ===========================================================================
|
|
# Phase 3 — Comparison Report
|
|
# ===========================================================================
|
|
|
|
def run_sft_phase3(
|
|
phase1_results: Dict[str, Any],
|
|
phase2_results: Dict[str, Any],
|
|
output_dir: Path,
|
|
base_results_dir: Path,
|
|
total_elapsed_sec: float,
|
|
) -> Optional[Path]:
|
|
"""Generate Base vs SFT comparison report."""
|
|
try:
|
|
from eval.report_generator import generate_comparison_report
|
|
|
|
report_path = generate_comparison_report(
|
|
base_results_dir=base_results_dir,
|
|
sft_phase1_results=phase1_results,
|
|
sft_phase2_results=phase2_results,
|
|
output_path=_PROJECT_ROOT / "reports" / f"{datetime.now().strftime('%Y-%m-%d')}_3B_SFT_V2_EVALUATION_REPORT.md",
|
|
sft_output_dir=output_dir,
|
|
total_elapsed_sec=total_elapsed_sec,
|
|
)
|
|
logger.info(" Comparison report saved: %s", report_path)
|
|
return report_path
|
|
except Exception:
|
|
logger.error(" Phase 3 report generation failed:\n%s", traceback.format_exc())
|
|
|
|
# Fallback: dump raw JSON
|
|
fallback = output_dir / "sft_eval_summary.json"
|
|
_save_json({
|
|
"phase1": phase1_results,
|
|
"phase2": phase2_results,
|
|
}, fallback)
|
|
logger.info(" Fallback summary saved: %s", fallback)
|
|
return None
|
|
|
|
|
|
# ===========================================================================
|
|
# CLI
|
|
# ===========================================================================
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="FRANKENSTALLM 3B — SFT Evaluation Pipeline",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
)
|
|
parser.add_argument("--dry-run", action="store_true")
|
|
parser.add_argument("--skip-phase0", action="store_true",
|
|
help="Skip HF conversion (reuse existing).")
|
|
parser.add_argument("--skip-phase1", action="store_true",
|
|
help="Skip internal eval.")
|
|
parser.add_argument("--skip-phase2", action="store_true",
|
|
help="Skip lm-eval benchmarks.")
|
|
parser.add_argument("--checkpoint", type=str, default=None,
|
|
help=f"Override SFT checkpoint (default: {SFT_CHECKPOINT})")
|
|
parser.add_argument("--hf-model-path", type=str, default=None,
|
|
help="Pre-converted HF model path (skips Phase 0).")
|
|
parser.add_argument("--output-dir", type=str, default=None,
|
|
help="Override output directory.")
|
|
parser.add_argument("--base-results", type=str, default=None,
|
|
help=f"Base eval results dir (default: {BASE_RESULTS_DIR})")
|
|
parser.add_argument("--gpus", type=str, default=None,
|
|
help="Comma-separated GPU IDs (default: 0-7).")
|
|
return parser.parse_args()
|
|
|
|
|
|
# ===========================================================================
|
|
# Main
|
|
# ===========================================================================
|
|
|
|
def main() -> None:
|
|
try:
|
|
mp.set_start_method("spawn", force=True)
|
|
except RuntimeError:
|
|
pass
|
|
|
|
args = parse_args()
|
|
|
|
# Resolve paths
|
|
checkpoint = args.checkpoint or SFT_CHECKPOINT
|
|
tokenizer = SFT_TOKENIZER if Path(SFT_TOKENIZER).exists() else _FALLBACK_TOKENIZER
|
|
base_results_dir = Path(args.base_results) if args.base_results else BASE_RESULTS_DIR
|
|
|
|
# Output directory
|
|
if args.output_dir:
|
|
output_dir = Path(args.output_dir)
|
|
else:
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M")
|
|
output_dir = _PROJECT_ROOT / "eval" / "outputs" / f"3b_sft_eval_{timestamp}"
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# GPU IDs
|
|
gpu_ids = sorted([int(g.strip()) for g in args.gpus.split(",")]) if args.gpus else list(range(8))
|
|
|
|
# Dry run
|
|
if args.dry_run:
|
|
_print_banner("DRY RUN — SFT Eval Pipeline")
|
|
logger.info(" SFT Checkpoint : %s", checkpoint)
|
|
logger.info(" Tokenizer : %s", tokenizer)
|
|
logger.info(" Base Results : %s", base_results_dir)
|
|
logger.info(" Output dir : %s", output_dir)
|
|
logger.info(" GPUs : %s", gpu_ids)
|
|
logger.info(" Chat template : ENABLED for generation tasks")
|
|
logger.info("")
|
|
|
|
phase1_tasks = _build_phase1_tasks(gpu_ids)
|
|
logger.info(" Phase 1 Tasks (%d):", len(phase1_tasks))
|
|
for desc in phase1_tasks:
|
|
is_gen = desc["task"] in ("generation", "repetition_grid")
|
|
chat_mark = " [CHAT]" if is_gen else ""
|
|
logger.info(" GPU %d — %s%s", desc["gpu_id"], desc["label"], chat_mark)
|
|
|
|
phase2_tasks = _build_phase2_tasks(gpu_ids)
|
|
logger.info(" Phase 2 Tasks (%d):", len(phase2_tasks))
|
|
for gpu_id, tasks, label in phase2_tasks:
|
|
logger.info(" GPU %d — %s", gpu_id, label)
|
|
|
|
# Check Base results exist
|
|
if base_results_dir.exists():
|
|
p1_file = base_results_dir / "phase1_results.json"
|
|
p2_file = base_results_dir / "phase2_results.json"
|
|
logger.info(" Base phase1_results.json: %s", "OK" if p1_file.exists() else "MISSING")
|
|
logger.info(" Base phase2_results.json: %s", "OK" if p2_file.exists() else "MISSING")
|
|
else:
|
|
logger.warning(" Base results dir NOT FOUND: %s", base_results_dir)
|
|
|
|
sys.exit(0)
|
|
|
|
# -----------------------------------------------------------------------
|
|
# Banner
|
|
# -----------------------------------------------------------------------
|
|
_print_banner("FRANKENSTALLM 3B — SFT Evaluation Pipeline")
|
|
logger.info(" SFT Checkpoint : %s", checkpoint)
|
|
logger.info(" Tokenizer : %s", tokenizer)
|
|
logger.info(" Base Results : %s", base_results_dir)
|
|
logger.info(" Output dir : %s", output_dir)
|
|
logger.info(" GPUs : %s", gpu_ids)
|
|
logger.info(" Phases : phase0=%s phase1=%s phase2=%s",
|
|
"skip" if args.skip_phase0 else "run",
|
|
"skip" if args.skip_phase1 else "run",
|
|
"skip" if args.skip_phase2 else "run")
|
|
|
|
# Preflight checks
|
|
if not Path(checkpoint).exists():
|
|
logger.error("SFT checkpoint not found: %s", checkpoint)
|
|
sys.exit(1)
|
|
if not Path(tokenizer).exists():
|
|
logger.error("Tokenizer not found: %s", tokenizer)
|
|
sys.exit(1)
|
|
logger.info(" Preflight OK: checkpoint=%s, tokenizer=%s", checkpoint, tokenizer)
|
|
|
|
pipeline_start = time.time()
|
|
phase1_results: Dict[str, Any] = {}
|
|
phase2_results: Dict[str, Any] = {}
|
|
hf_model_path: Optional[Path] = None
|
|
|
|
# -----------------------------------------------------------------------
|
|
# Phase 0 — HF Conversion
|
|
# -----------------------------------------------------------------------
|
|
_print_phase_header("PHASE 0", "SFT Checkpoint → HuggingFace Conversion")
|
|
if args.hf_model_path:
|
|
hf_model_path = Path(args.hf_model_path)
|
|
logger.info(" Using pre-converted HF model: %s", hf_model_path)
|
|
elif args.skip_phase0:
|
|
# Search for existing HF conversion
|
|
candidate = output_dir / "hf_3b_sft_best"
|
|
outputs_dir = _PROJECT_ROOT / "eval" / "outputs"
|
|
if candidate.exists():
|
|
hf_model_path = candidate
|
|
else:
|
|
candidates = list(outputs_dir.glob("hf_3b_sft*"))
|
|
if candidates:
|
|
hf_model_path = candidates[0]
|
|
if hf_model_path:
|
|
logger.info(" Skipping Phase 0 — reusing: %s", hf_model_path)
|
|
else:
|
|
logger.warning(" No HF model found. Phase 2 will be skipped.")
|
|
else:
|
|
t0 = time.time()
|
|
try:
|
|
hf_output = output_dir / "hf_3b_sft_best"
|
|
hf_output.mkdir(parents=True, exist_ok=True)
|
|
|
|
import subprocess
|
|
convert_script = _PROJECT_ROOT / "scripts" / "convert_to_hf.py"
|
|
cmd = [
|
|
sys.executable, str(convert_script),
|
|
"--checkpoint", checkpoint,
|
|
"--output", str(hf_output),
|
|
"--tokenizer", tokenizer,
|
|
]
|
|
logger.info(" Running: %s", " ".join(cmd))
|
|
subprocess.run(cmd, check=True, cwd=str(_PROJECT_ROOT))
|
|
hf_model_path = hf_output
|
|
logger.info(" Phase 0 complete in %s.", _fmt_seconds(time.time() - t0))
|
|
except Exception:
|
|
logger.error(" Phase 0 FAILED:\n%s", traceback.format_exc())
|
|
|
|
# -----------------------------------------------------------------------
|
|
# Phase 1 — Internal Evaluation (8 GPU)
|
|
# -----------------------------------------------------------------------
|
|
_print_phase_header("PHASE 1", f"SFT Internal Evaluation — {len(gpu_ids)} GPU Parallel")
|
|
if args.skip_phase1:
|
|
logger.info(" Skipping Phase 1.")
|
|
phase1_out = output_dir / "phase1_results.json"
|
|
if phase1_out.exists():
|
|
with open(phase1_out, encoding="utf-8") as f:
|
|
phase1_results = json.load(f)
|
|
logger.info(" Loaded existing Phase 1 results.")
|
|
else:
|
|
t0 = time.time()
|
|
try:
|
|
phase1_results = run_sft_phase1(output_dir, gpu_ids, checkpoint, tokenizer)
|
|
logger.info(" Phase 1 complete in %s.", _fmt_seconds(time.time() - t0))
|
|
except Exception:
|
|
logger.error(" Phase 1 FAILED:\n%s", traceback.format_exc())
|
|
|
|
# -----------------------------------------------------------------------
|
|
# Phase 2 — lm-eval Benchmarks (8 GPU)
|
|
# -----------------------------------------------------------------------
|
|
_print_phase_header("PHASE 2", f"SFT Benchmarks — {len(gpu_ids)} GPU Parallel")
|
|
if args.skip_phase2:
|
|
logger.info(" Skipping Phase 2.")
|
|
phase2_out = output_dir / "phase2_results.json"
|
|
if phase2_out.exists():
|
|
with open(phase2_out, encoding="utf-8") as f:
|
|
phase2_results = json.load(f)
|
|
logger.info(" Loaded existing Phase 2 results.")
|
|
elif hf_model_path is None:
|
|
logger.warning(" Phase 2 skipped — HF model unavailable.")
|
|
else:
|
|
t0 = time.time()
|
|
try:
|
|
phase2_results = run_sft_phase2(
|
|
hf_model_path, output_dir, gpu_ids, checkpoint, tokenizer,
|
|
)
|
|
logger.info(" Phase 2 complete in %s.", _fmt_seconds(time.time() - t0))
|
|
except Exception:
|
|
logger.error(" Phase 2 FAILED:\n%s", traceback.format_exc())
|
|
|
|
# -----------------------------------------------------------------------
|
|
# Phase 3 — Comparison Report
|
|
# -----------------------------------------------------------------------
|
|
_print_phase_header("PHASE 3", "Base vs SFT Comparison Report")
|
|
t0 = time.time()
|
|
report_path = run_sft_phase3(
|
|
phase1_results, phase2_results, output_dir, base_results_dir,
|
|
total_elapsed_sec=time.time() - pipeline_start,
|
|
)
|
|
logger.info(" Phase 3 complete in %s.", _fmt_seconds(time.time() - t0))
|
|
|
|
# -----------------------------------------------------------------------
|
|
# Final Summary
|
|
# -----------------------------------------------------------------------
|
|
total_elapsed = time.time() - pipeline_start
|
|
_print_banner("SFT EVALUATION PIPELINE COMPLETE")
|
|
logger.info(" Total time : %s", _fmt_seconds(total_elapsed))
|
|
logger.info(" Output dir : %s", output_dir)
|
|
logger.info(" Phase 1 results : %s", output_dir / "phase1_results.json")
|
|
logger.info(" Phase 2 results : %s", output_dir / "phase2_results.json")
|
|
logger.info(" Report : %s", report_path or "N/A")
|
|
logger.info(_bar())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|