""" 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()