366 lines
13 KiB
Python
366 lines
13 KiB
Python
|
"""
|
||
|
|
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
|