diff --git a/examples/vlm-evaluation/README.md b/examples/vlm-evaluation/README.md index d8a919e..0fb31d5 100644 --- a/examples/vlm-evaluation/README.md +++ b/examples/vlm-evaluation/README.md @@ -127,6 +127,36 @@ accelerate launch --num_processes 4 examples/vlm-evaluation/vlm_eval_harness.py | `--batch-size` | `1` | requests decoded together (grouped by `gen_kwargs`) | | `--max-new-tokens` | `128` | fallback only; task `gen_kwargs` override it | +## Experiment tracking + +Results can be logged through the framework's metrics backends +(`kempnerforge/metrics/tracker.py`) using the same config flags as training, +forwarded as dotted overrides — any unrecognized `--section.key=value` argument +is layered over `--config` (unknown keys fail fast) and applies to both the +evaluated model (e.g. `--video.max_frames=8`) and experiment tracking: + +```bash +uv run python examples/vlm-evaluation/vlm_eval_harness.py \ + --config configs/train/vlm_jd.toml \ + --checkpoint checkpoints/vlm/step_10000 \ + --tasks mmmu_val \ + --metrics.enable_wandb=true --metrics.wandb_project=vlm-eval +``` + +- Results land in the checkpoint's training run: the harness reads the + `wandb_run_id` saved in `train_state.pt` and resumes that run. Without one, a + fresh run named `-` starts with a warning; target a specific run + with `--metrics.wandb_run_id=`. +- Metrics are logged at the checkpoint's training step: + `eval/benchmarks/agg/` (aggregate, normalized into [0, 1]), + `eval/benchmarks/raw//` (every metric, unnormalized), and + `eval/benchmarks/throughput/overall/...` (per-invocation, so comparable + across runs only for the same task set). +- Per-benchmark metric names and scales live in + [`benchmark_manifest.py`](benchmark_manifest.py); an unregistered benchmark + warns with the exact registry line to add. +- A tracking failure never fails a completed eval — it warns and moves on. + ## Video evaluation When `--config` is a **video checkpoint** (its TOML has a `[video]` section), the diff --git a/examples/vlm-evaluation/adapter.py b/examples/vlm-evaluation/adapter.py index 9760db2..068b210 100644 --- a/examples/vlm-evaluation/adapter.py +++ b/examples/vlm-evaluation/adapter.py @@ -159,14 +159,19 @@ def _log_checkpoint_metadata(ckpt_path: Path) -> None: ) -def _load_config(config_path: str) -> JobConfig: - config = load_config(config_path, cli_args=[]) - if not config.is_vlm: +def _load_config(config: str | JobConfig) -> JobConfig: + """A TOML path, or an already-loaded ``JobConfig`` (the harness passes its + CLI-merged config so ``--section.key=value`` overrides reach the model); + the VLM-only guard applies either way. + """ + loaded = config if isinstance(config, JobConfig) else load_config(config, cli_args=[]) + if not loaded.is_vlm: + source = repr(config) if isinstance(config, str) else "the provided JobConfig" raise ValueError( - f"{config_path!r} is not a VLM config (config.vlm is None); this evaluation " + f"{source} is not a VLM config (config.vlm is None); this evaluation " f"path is VLM-only. Use scripts/eval.py for text-model loss/perplexity." ) - return config + return loaded def _check_generative(vlm_config: VLMConfig) -> None: @@ -553,7 +558,8 @@ class KempnerForgeVLM(lmms): directly and passes the instance to ``simple_evaluate``): - ``config`` (required): path to the KempnerForge TOML the checkpoint was - trained with. + trained with, or an already-loaded ``JobConfig`` (the harness passes its + CLI-merged config so ``--section.key=value`` overrides reach the model). - ``checkpoint`` (required): DCP checkpoint directory (a run dir or a specific ``step_N`` dir). - ``device`` (default ``"cuda"``), ``dtype`` (default: the checkpoint config's @@ -568,7 +574,7 @@ class KempnerForgeVLM(lmms): def __init__( self, - config: str, + config: str | JobConfig, checkpoint: str, device: str = "cuda", dtype: str | None = None, diff --git a/examples/vlm-evaluation/benchmark_manifest.py b/examples/vlm-evaluation/benchmark_manifest.py new file mode 100644 index 0000000..b884850 --- /dev/null +++ b/examples/vlm-evaluation/benchmark_manifest.py @@ -0,0 +1,289 @@ +"""Benchmark manifest: per-benchmark metric knowledge for VLM eval. + +Each entry records which metric in an lmms-eval results dict is a benchmark's +authoritative aggregate and how to map it into [0, 1]. Unregistered benchmarks +fall back to a metadata-driven guess with a loud warning. + +``build_eval_metrics`` flattens a ``simple_evaluate`` results dict for logging: + + eval/benchmarks/agg/ normalized [0, 1] aggregate + eval/benchmarks/raw//[/...] every numeric metric, raw + eval/benchmarks/throughput/overall/ native lmms-eval throughput (per-invocation) + eval/benchmarks/efficiency// only when the run logged samples +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass + +logger = logging.getLogger(__name__) + +_KEY_PREFIX = "eval/benchmarks" + + +@dataclass(frozen=True) +class MetricSpec: + """How to read one benchmark's authoritative aggregate from an lmms-eval result. + + ``metric`` is the name before the ``,`` suffix; ``None`` = no local + score (submission-only task). ``scale`` divides the raw value into [0, 1]; + ``None`` = infer by magnitude (raw > 1.0 -> /100, else /1). ``filter`` + disambiguates multi-filter metrics; ``subkey`` indexes a dict-valued one. + """ + + metric: str | None + scale: float | None = 1.0 + filter: str | None = None + subkey: str | None = None + + +# Authoritative aggregate per benchmark, keyed by the lmms-eval top-level task +# name. Neither the metric nor its scale can be inferred (``accuracy`` is 0-100 +# for mmvu_val but 0-1 for perceptiontest_val_mc), so entries are explicit. +BENCHMARK_METRICS: dict[str, MetricSpec] = { + # --- image --- + "hallusion_bench_image": MetricSpec("aAcc", scale=100.0), + "mmvu_val": MetricSpec("accuracy", scale=100.0), + "mmstar": MetricSpec("average"), + "blink": MetricSpec("blink_acc"), + "realworldqa": MetricSpec("exact_match"), + "mmmu_pro_standard": MetricSpec("mmmu_acc"), + # --- video --- + "mlvu_dev": MetricSpec("mlvu_percetion_score", scale=100.0), # upstream typo; real key + "tempcompass": MetricSpec("avg_accuracy", scale=100.0), # main entry empty -> subtask mean + "tvbench": MetricSpec("tvbench_acc"), # main entry empty -> subtask mean + "vsibench": MetricSpec("vsibench_overall"), + "tomato": MetricSpec("tomato_score"), + "nextqa_mc_test": MetricSpec("exact_match"), + "perceptiontest_val_mc": MetricSpec("accuracy"), + "videoevalpro": MetricSpec("videoevalpro_score", subkey="overall"), # dict-valued aggregate + "lvbench": MetricSpec("lvbench_score"), + "videomme_v2": MetricSpec("videomme_v2_overall_acc", scale=100.0), + "egoschema": MetricSpec(None), # submission-only; no local score + # --- vdc caption splits (upstream double-l typo) --- + "camera_test": MetricSpec("llmms_eval_acc"), + "background_test": MetricSpec("llmms_eval_acc"), + "detailed_test": MetricSpec("llmms_eval_acc"), + "main_object_test": MetricSpec("llmms_eval_acc"), + "short_test": MetricSpec("llmms_eval_acc"), + # --- text --- + "gsm8k_cot_zeroshot": MetricSpec("exact_match", filter="flexible-extract"), + "mmlu_flan_cot_zeroshot": MetricSpec("exact_match", filter="flexible-extract"), + "gpqa_main_cot_zeroshot": MetricSpec("exact_match", filter="flexible-extract"), +} + + +def _lookup_metric( + entry: dict, metric: str, filter: str | None = None, subkey: str | None = None +) -> float | None: + """Value of ``metric`` in one result entry, matched on the name before the comma. + + A requested ``filter`` must match exactly (filter variants like strict-match + vs flexible-extract are materially different numbers, so no other variant may + stand in for a registered one); without a ``filter``, prefers ``"none"``, then + the first match. Skips ``alias``/stderr columns. Returns ``None`` when absent + or non-numeric. + """ + matches: list[tuple[str, object]] = [] + for key, value in entry.items(): + if key == "alias" or "stderr" in key: + continue + base, _, filt = key.partition(",") + if base == metric: + matches.append((filt, value)) + if not matches: + return None + if filter is not None: + chosen = next((v for f, v in matches if f == filter), None) + else: + chosen = next((v for f, v in matches if f == "none"), matches[0][1]) + if subkey is not None and isinstance(chosen, dict): + chosen = chosen.get(subkey) + if isinstance(chosen, bool) or not isinstance(chosen, (int, float)): + return None + return float(chosen) + + +def _available_filters(entry: dict, metric: str) -> list[str]: + """Filter variants under which ``metric`` appears in one result entry.""" + filters: list[str] = [] + for key in entry: + if key == "alias" or "stderr" in key: + continue + base, _, filt = key.partition(",") + if base == metric: + filters.append(filt) + return filters + + +def _resolve_spec(task: str, higher_is_better: dict) -> MetricSpec | None: + """Registry entry for ``task``, else a loud ``higher_is_better``-driven + fallback, else ``None`` (skip). + """ + spec = BENCHMARK_METRICS.get(task) + if spec is not None: + return spec # registered (metric may be None -> caller skips silently) + candidates = list((higher_is_better.get(task) or {}).keys()) + if len(candidates) == 1: + logger.warning( + "Benchmark %r is not registered in BENCHMARK_METRICS; falling back to its sole " + "higher_is_better metric %r with a magnitude-inferred scale (raw > 1.0 is treated " + "as 0-100, else 0-1). Register it in benchmark_manifest.py to make this explicit:\n" + " %r: MetricSpec(%r), # add scale=100.0 if the metric is a 0-100 percentage", + task, + candidates[0], + task, + candidates[0], + ) + return MetricSpec(candidates[0], scale=None) + logger.warning( + "Benchmark %r is not registered in BENCHMARK_METRICS and its higher_is_better names " + "%d candidate metric(s) %s — cannot choose an aggregate, so it is skipped (raw " + "metrics are still logged). Register it in benchmark_manifest.py:\n" + ' %r: MetricSpec(""),', + task, + len(candidates), + candidates, + task, + ) + return None + + +def benchmark_aggregates(results: dict, tasks: list[str]) -> dict[str, float]: + """Map each requested benchmark to its aggregate as a [0, 1] fraction. + + Only the requested tasks are resolved (``group_subtasks`` also lists nested + intermediate groups). A grouped task whose main entry carries no value is + the mean of the metric over its subtasks. + """ + task_results: dict[str, dict] = results.get("results", {}) + group_subtasks: dict[str, list] = results.get("group_subtasks", {}) + higher_is_better: dict = results.get("higher_is_better", {}) + scores: dict[str, float] = {} + for task in tasks: + spec = _resolve_spec(task, higher_is_better) + if spec is None or spec.metric is None: + continue # unresolved (already warned) or submission-only + raw = _lookup_metric(task_results.get(task) or {}, spec.metric, spec.filter, spec.subkey) + if raw is None: + subtask_values = [ + value + for name in group_subtasks.get(task) or [] + for value in ( + _lookup_metric( + task_results.get(name) or {}, spec.metric, spec.filter, spec.subkey + ), + ) + if value is not None + ] + if subtask_values: + raw = sum(subtask_values) / len(subtask_values) + if raw is None: + available = sorted( + { + filt + for name in (task, *(group_subtasks.get(task) or [])) + for filt in _available_filters(task_results.get(name) or {}, spec.metric) + } + ) + if spec.filter is not None and available: + logger.warning( + "Benchmark %r: registered metric %r is present only under filter(s) %s, not " + "the registered filter %r; skipping its aggregate rather than silently " + "logging a different variant (raw metrics are still logged). If lmms-eval " + "renamed the filter, update benchmark_manifest.py:\n" + " %r: MetricSpec(%r, filter=),", + task, + spec.metric, + available, + spec.filter, + task, + spec.metric, + available, + ) + else: + logger.warning( + "Benchmark %r: registered metric %r is absent from the results and no " + "subtask supplied it; skipping its aggregate (raw metrics are still logged).", + task, + spec.metric, + ) + continue + divisor = spec.scale if spec.scale is not None else (100.0 if raw > 1.0 else 1.0) + scores[task] = raw / divisor + return scores + + +def _raw_metrics(results: dict) -> dict[str, float]: + """Flatten every numeric metric of every task and subtask, unnormalized. + + Keys are ``eval/benchmarks/raw//``, with non-``none`` filters + and dict keys appended; ``alias``/stderr columns are skipped. + """ + metrics: dict[str, float] = {} + for task, entry in (results.get("results") or {}).items(): + if not isinstance(entry, dict): + continue + for key, value in entry.items(): + if key == "alias" or "stderr" in key: + continue + base, _, filt = key.partition(",") + parts = [_KEY_PREFIX, "raw", task, base] + if filt and filt != "none": + parts.append(filt) + if isinstance(value, dict): + for sub, sub_value in value.items(): + if isinstance(sub_value, (int, float)) and not isinstance(sub_value, bool): + metrics["/".join([*parts, str(sub)])] = float(sub_value) + elif isinstance(value, (int, float)) and not isinstance(value, bool): + metrics["/".join(parts)] = float(value) + return metrics + + +_THROUGHPUT_KEYS = ("avg_speed", "total_gen_tokens", "total_elapsed_time", "avg_latency") +_EFFICIENCY_KEYS = ( + "total_output_tokens", + "total_input_tokens", + "total_tokens", + "avg_output_tokens_per_sample", + "tokens_per_correct_answer", +) + + +def _perf_metrics(results: dict) -> dict[str, float]: + """Native lmms-eval throughput/efficiency summaries as scalars. + + ``results["throughput"]`` is per-invocation (one ``generate_until`` call + spans every task), so it is logged once under ``overall`` — comparable + across runs only for the same task set. ``efficiency["by_task"]`` is + genuinely per-task and exists only when samples were logged. + """ + metrics: dict[str, float] = {} + throughput = results.get("throughput") or {} + for key in _THROUGHPUT_KEYS: + value = throughput.get(key) + if isinstance(value, (int, float)) and not isinstance(value, bool): + metrics[f"{_KEY_PREFIX}/throughput/overall/{key}"] = float(value) + by_task = (results.get("efficiency") or {}).get("by_task") or {} + for task, summary in by_task.items(): + if not isinstance(summary, dict): + continue + for key in _EFFICIENCY_KEYS: + value = summary.get(key) + if isinstance(value, (int, float)) and not isinstance(value, bool): + metrics[f"{_KEY_PREFIX}/efficiency/{task}/{key}"] = float(value) + return metrics + + +def build_eval_metrics(results: dict, tasks: list[str]) -> dict[str, float]: + """One flat ``{key: float}`` dict for logging: normalized aggregates, the + full raw flatten, and native throughput/efficiency. + """ + metrics = { + f"{_KEY_PREFIX}/agg/{bench}": score + for bench, score in benchmark_aggregates(results, tasks).items() + } + metrics.update(_raw_metrics(results)) + metrics.update(_perf_metrics(results)) + return metrics diff --git a/examples/vlm-evaluation/tests/unit/test_adapter.py b/examples/vlm-evaluation/tests/unit/test_adapter.py index f1f7381..d5ce4f3 100644 --- a/examples/vlm-evaluation/tests/unit/test_adapter.py +++ b/examples/vlm-evaluation/tests/unit/test_adapter.py @@ -801,6 +801,15 @@ def test_vlm_config_accepted(self, monkeypatch, tiny_vlm_configs): monkeypatch.setattr("adapter.load_config", lambda _p, cli_args=None: job) assert _load_config("ignored.toml") is job + def test_jobconfig_instance_passes_through(self, tiny_vlm_configs): + """The harness passes its CLI-merged JobConfig directly; no re-parse.""" + job = _vlm_job_config(tiny_vlm_configs) + assert _load_config(job) is job + + def test_non_vlm_jobconfig_instance_rejected(self, tiny_job_config): + with pytest.raises(ValueError, match="not a VLM config"): + _load_config(tiny_job_config) + # --------------------------------------------------------------------------- # _load_weights (path resolution / missing checkpoint) diff --git a/examples/vlm-evaluation/tests/unit/test_benchmark_manifest.py b/examples/vlm-evaluation/tests/unit/test_benchmark_manifest.py new file mode 100644 index 0000000..dd0f8c0 --- /dev/null +++ b/examples/vlm-evaluation/tests/unit/test_benchmark_manifest.py @@ -0,0 +1,304 @@ +"""Unit tests for the benchmark manifest. Fixtures mirror the result shapes +``simple_evaluate`` produces (leaf tasks, grouped tasks with empty main +entries, nested groups, dict-valued aggregates).""" + +from __future__ import annotations + +import logging + +import pytest +from benchmark_manifest import ( + BENCHMARK_METRICS, + MetricSpec, + benchmark_aggregates, + build_eval_metrics, +) + +# --------------------------------------------------------------------------- # +# Aggregate resolution — registered benchmarks +# --------------------------------------------------------------------------- # + + +def test_registered_scale_normalizes_to_unit_interval(): + """A 0-100 metric (hallusion aAcc) is divided into [0, 1].""" + results = { + "results": { + "hallusion_bench_image": { + "alias": "hallusion_bench_image", + "aAcc,none": 63.0, + "aAcc_stderr,none": 1.2, + } + }, + "group_subtasks": {"hallusion_bench_image": []}, + } + scores = benchmark_aggregates(results, ["hallusion_bench_image"]) + assert scores == pytest.approx({"hallusion_bench_image": 0.63}) + + +def test_grouped_task_averages_subtasks_when_main_entry_empty(): + """tempcompass-shape: the group entry has no numeric metric -> mean over subtasks.""" + results = { + "results": { + "tempcompass": {" ": " ", "alias": "tempcompass"}, + "tempcompass_mc": {"alias": " - mc", "avg_accuracy,none": 40.0}, + "tempcompass_yn": {"alias": " - yn", "avg_accuracy,none": 60.0}, + }, + "group_subtasks": {"tempcompass": ["tempcompass_mc", "tempcompass_yn"]}, + } + scores = benchmark_aggregates(results, ["tempcompass"]) + assert scores == pytest.approx({"tempcompass": 0.5}) # mean(40, 60) / 100 + + +def test_filter_disambiguates_multi_filter_metric(): + """gsm8k reports exact_match under two filters; the registered one wins.""" + results = { + "results": { + "gsm8k_cot_zeroshot": { + "alias": "gsm8k_cot_zeroshot", + "exact_match,strict-match": 0.10, + "exact_match,flexible-extract": 0.42, + } + }, + "group_subtasks": {"gsm8k_cot_zeroshot": []}, + } + scores = benchmark_aggregates(results, ["gsm8k_cot_zeroshot"]) + assert scores == pytest.approx({"gsm8k_cot_zeroshot": 0.42}) + + +def test_registered_filter_absent_warns_and_skips(caplog): + """The registered filter must match exactly: another variant (strict-match) + must never silently stand in for the registered one (flexible-extract).""" + results = { + "results": { + "gsm8k_cot_zeroshot": { + "alias": "gsm8k_cot_zeroshot", + "exact_match,strict-match": 0.10, + } + }, + "group_subtasks": {"gsm8k_cot_zeroshot": []}, + } + with caplog.at_level(logging.WARNING, logger="benchmark_manifest"): + scores = benchmark_aggregates(results, ["gsm8k_cot_zeroshot"]) + assert scores == {} + warnings = [r.getMessage() for r in caplog.records] + assert any( + "'flexible-extract'" in w and "strict-match" in w and "skipping its aggregate" in w + for w in warnings + ) + + +def test_no_registered_filter_still_falls_back(): + """Specs without a filter keep the lenient chain: 'none' first, else the + first available filter variant.""" + results = { + "results": {"blink": {"alias": "blink", "blink_acc,custom-filter": 0.6}}, + "group_subtasks": {"blink": []}, + } + scores = benchmark_aggregates(results, ["blink"]) + assert scores == pytest.approx({"blink": 0.6}) + + +def test_subkey_selects_from_dict_valued_metric(): + """videoevalpro's aggregate is a dict keyed by task type plus 'overall'.""" + results = { + "results": { + "videoevalpro": { + "alias": "videoevalpro", + "videoevalpro_score,none": {"Local Perception": 0.30, "overall": 0.45}, + } + }, + "group_subtasks": {"videoevalpro": []}, + } + scores = benchmark_aggregates(results, ["videoevalpro"]) + assert scores == pytest.approx({"videoevalpro": 0.45}) + + +def test_submission_only_benchmark_skipped_silently(caplog): + """egoschema is registered with metric=None -> no aggregate, no warning.""" + results = { + "results": {"egoschema": {"alias": "egoschema", "submission,none": None}}, + "group_subtasks": {"egoschema": []}, + } + with caplog.at_level(logging.WARNING, logger="benchmark_manifest"): + scores = benchmark_aggregates(results, ["egoschema"]) + assert scores == {} + assert not caplog.records + + +def test_registered_metric_absent_warns_and_skips(caplog): + """A registered metric missing from the results warns instead of guessing.""" + results = { + "results": {"blink": {"alias": "blink", "some_other_metric,none": 0.5}}, + "group_subtasks": {"blink": []}, + } + with caplog.at_level(logging.WARNING, logger="benchmark_manifest"): + scores = benchmark_aggregates(results, ["blink"]) + assert scores == {} + assert any("absent from the results" in r.getMessage() for r in caplog.records) + + +# --------------------------------------------------------------------------- # +# Aggregate resolution — unregistered fallback +# --------------------------------------------------------------------------- # + + +def test_unregistered_sole_metric_fallback_infers_scale(caplog): + """Sole higher_is_better metric is used; magnitude picks the divisor; warns loudly.""" + results = { + "results": { + "newbench_pct": {"alias": "newbench_pct", "acc,none": 73.2}, + "newbench_frac": {"alias": "newbench_frac", "acc,none": 0.73}, + }, + "group_subtasks": {"newbench_pct": [], "newbench_frac": []}, + "higher_is_better": {"newbench_pct": {"acc": True}, "newbench_frac": {"acc": True}}, + } + with caplog.at_level(logging.WARNING, logger="benchmark_manifest"): + scores = benchmark_aggregates(results, ["newbench_pct", "newbench_frac"]) + assert scores == pytest.approx({"newbench_pct": 0.732, "newbench_frac": 0.73}) + warnings = [r.getMessage() for r in caplog.records] + assert sum("not registered in BENCHMARK_METRICS" in w for w in warnings) == 2 + # The warning embeds a paste-ready registry line for the offending task. + assert any("'newbench_pct': MetricSpec('acc')" in w for w in warnings) + + +def test_unregistered_ambiguous_task_skipped_with_warning(caplog): + """Several candidate metrics -> no way to choose an aggregate -> skip + warn.""" + results = { + "results": {"newbench": {"alias": "newbench", "a,none": 0.1, "b,none": 0.2}}, + "group_subtasks": {"newbench": []}, + "higher_is_better": {"newbench": {"a": True, "b": True}}, + } + with caplog.at_level(logging.WARNING, logger="benchmark_manifest"): + scores = benchmark_aggregates(results, ["newbench"]) + assert scores == {} + assert any("cannot choose an aggregate" in r.getMessage() for r in caplog.records) + + +def test_nested_group_names_never_resolved(caplog): + """mmlu-shape: intermediate groups (stem, humanities) appear in group_subtasks and + results but must not be resolved as benchmarks of their own.""" + results = { + "results": { + "mmlu_flan_cot_zeroshot": { + "alias": "mmlu_flan_cot_zeroshot", + "exact_match,flexible-extract": 0.31, + }, + "stem": {"alias": " - stem", "exact_match,flexible-extract": 0.28}, + "humanities": {"alias": " - humanities", "exact_match,flexible-extract": 0.33}, + "mmlu_abstract_algebra": {"alias": " - aa", "exact_match,flexible-extract": 0.2}, + }, + "group_subtasks": { + "mmlu_flan_cot_zeroshot": ["stem", "humanities"], + "stem": ["mmlu_abstract_algebra"], + "humanities": [], + }, + "higher_is_better": {"stem": {"exact_match": True}, "humanities": {"exact_match": True}}, + } + with caplog.at_level(logging.WARNING, logger="benchmark_manifest"): + scores = benchmark_aggregates(results, ["mmlu_flan_cot_zeroshot"]) + assert scores == pytest.approx({"mmlu_flan_cot_zeroshot": 0.31}) + assert not caplog.records # stem/humanities never hit the fallback + + +# --------------------------------------------------------------------------- # +# build_eval_metrics — the flat log dict +# --------------------------------------------------------------------------- # + + +def test_build_eval_metrics_key_scheme(): + results = { + "results": { + "realworldqa": { + "alias": "realworldqa", + "exact_match,none": 0.44, + "exact_match_stderr,none": 0.01, + } + }, + "group_subtasks": {"realworldqa": []}, + "throughput": {"avg_speed": 12.5, "total_gen_tokens": 300, "total_elapsed_time": 24.0}, + } + metrics = build_eval_metrics(results, ["realworldqa"]) + assert metrics["eval/benchmarks/agg/realworldqa"] == pytest.approx(0.44) + assert metrics["eval/benchmarks/raw/realworldqa/exact_match"] == pytest.approx(0.44) + assert metrics["eval/benchmarks/throughput/overall/avg_speed"] == pytest.approx(12.5) + assert metrics["eval/benchmarks/throughput/overall/total_gen_tokens"] == pytest.approx(300) + assert not any("stderr" in k for k in metrics) + + +def test_throughput_is_run_level_never_per_task(): + """The invocation-wide throughput summary is logged once under ``overall``, + never duplicated under each requested task (a two-task run would otherwise + attribute the combined run's totals to both tasks).""" + results = { + "results": { + "realworldqa": {"alias": "realworldqa", "exact_match,none": 0.44}, + "mmstar": {"alias": "mmstar", "average,none": 0.5}, + }, + "group_subtasks": {"realworldqa": [], "mmstar": []}, + "throughput": {"avg_speed": 12.5, "total_elapsed_time": 24.0}, + } + metrics = build_eval_metrics(results, ["realworldqa", "mmstar"]) + assert metrics["eval/benchmarks/throughput/overall/total_elapsed_time"] == pytest.approx(24.0) + per_task = [k for k in metrics if "/throughput/" in k and "/throughput/overall/" not in k] + assert per_task == [] + + +def test_build_eval_metrics_raw_includes_subtasks_and_filters(): + """Subtask rows and non-'none' filters are flattened; dicts flatten one level.""" + results = { + "results": { + "tvbench": {" ": " ", "alias": "tvbench"}, + "tvbench_action": {"alias": " - action", "tvbench_acc,none": 0.2}, + "gsm8k_cot_zeroshot": { + "alias": "gsm8k", + "exact_match,strict-match": 0.10, + "exact_match,flexible-extract": 0.42, + }, + "videoevalpro": { + "alias": "videoevalpro", + "videoevalpro_score,none": {"overall": 0.45}, + }, + }, + "group_subtasks": {"tvbench": ["tvbench_action"]}, + } + metrics = build_eval_metrics(results, ["tvbench"]) + raw = {k.removeprefix("eval/benchmarks/raw/"): v for k, v in metrics.items() if "/raw/" in k} + assert raw["tvbench_action/tvbench_acc"] == pytest.approx(0.2) + assert raw["gsm8k_cot_zeroshot/exact_match/strict-match"] == pytest.approx(0.10) + assert raw["gsm8k_cot_zeroshot/exact_match/flexible-extract"] == pytest.approx(0.42) + assert raw["videoevalpro/videoevalpro_score/overall"] == pytest.approx(0.45) + + +def test_build_eval_metrics_efficiency_only_when_present(): + results = { + "results": {"realworldqa": {"alias": "realworldqa", "exact_match,none": 0.4}}, + "group_subtasks": {"realworldqa": []}, + "efficiency": { + "by_task": { + "realworldqa": { + "total_output_tokens": 300.0, + "tokens_per_correct_answer": None, # None must be skipped + } + }, + "overall": {}, + }, + } + metrics = build_eval_metrics(results, ["realworldqa"]) + eff_prefix = "eval/benchmarks/efficiency/realworldqa" + assert metrics[f"{eff_prefix}/total_output_tokens"] == pytest.approx(300.0) + assert f"{eff_prefix}/tokens_per_correct_answer" not in metrics + + +# --------------------------------------------------------------------------- # +# Registry sanity +# --------------------------------------------------------------------------- # + + +def test_registry_specs_are_well_formed(): + """Every registered spec has a positive explicit scale (or None/submission-only).""" + for task, spec in BENCHMARK_METRICS.items(): + assert isinstance(spec, MetricSpec) + if spec.scale is not None: + assert spec.scale > 0, f"{task}: non-positive scale" + if spec.metric is None: + assert spec.filter is None and spec.subkey is None, f"{task}: dead fields" diff --git a/examples/vlm-evaluation/tests/unit/test_harness_tracking.py b/examples/vlm-evaluation/tests/unit/test_harness_tracking.py new file mode 100644 index 0000000..fe3a43d --- /dev/null +++ b/examples/vlm-evaluation/tests/unit/test_harness_tracking.py @@ -0,0 +1,142 @@ +"""Unit tests for the harness's experiment-tracking helpers. +``load_train_state_extras`` is monkeypatched at its source module — the +harness re-imports it per call.""" + +from __future__ import annotations + +import json +import logging +from pathlib import Path + +import vlm_eval_harness + +import kempnerforge.checkpoint +from kempnerforge.config.schema import JobConfig + +# --------------------------------------------------------------------------- # +# _checkpoint_step +# --------------------------------------------------------------------------- # + + +def test_checkpoint_step_from_metadata(tmp_path: Path): + ckpt = tmp_path / "step_500" + ckpt.mkdir() + (ckpt / "metadata.json").write_text(json.dumps({"step": 10_000, "tokens_seen": 1})) + + assert vlm_eval_harness._checkpoint_step(ckpt) == 10_000 + + +def test_checkpoint_step_falls_back_to_dir_name(tmp_path: Path): + ckpt = tmp_path / "step_500" + ckpt.mkdir() + + assert vlm_eval_harness._checkpoint_step(ckpt) == 500 + + +def test_checkpoint_step_unknown_warns_and_returns_zero(tmp_path: Path, caplog): + ckpt = tmp_path / "final" + ckpt.mkdir() + + with caplog.at_level(logging.WARNING, logger="vlm_eval_harness"): + step = vlm_eval_harness._checkpoint_step(ckpt) + + assert step == 0 + assert any("training step" in r.getMessage() for r in caplog.records) + + +# --------------------------------------------------------------------------- # +# _resolve_run_id +# --------------------------------------------------------------------------- # + + +def _config(**metrics_overrides) -> JobConfig: + config = JobConfig() + for key, value in metrics_overrides.items(): + setattr(config.metrics, key, value) + return config + + +def test_explicit_run_id_wins_without_reading_checkpoint(tmp_path: Path, monkeypatch): + def _boom(_ckpt_dir): + raise AssertionError("checkpoint must not be consulted when an id is explicit") + + monkeypatch.setattr(kempnerforge.checkpoint, "load_train_state_extras", _boom) + config = _config(wandb_run_id="explicit-id") + + vlm_eval_harness._resolve_run_id(config, tmp_path / "step_1") + + assert config.metrics.wandb_run_id == "explicit-id" + + +def test_checkpoint_run_id_adopted(tmp_path: Path, monkeypatch): + monkeypatch.setattr( + kempnerforge.checkpoint, + "load_train_state_extras", + lambda _ckpt_dir: {"wandb_run_id": "ckpt-run"}, + ) + config = _config() + + vlm_eval_harness._resolve_run_id(config, tmp_path / "step_1") + + assert config.metrics.wandb_run_id == "ckpt-run" + assert config.metrics.wandb_run_name is None # attach path never renames the run + + +def test_no_run_id_warns_and_derives_fresh_run_name(tmp_path: Path, monkeypatch, caplog): + monkeypatch.setattr(kempnerforge.checkpoint, "load_train_state_extras", lambda _d: {}) + config = _config() + ckpt_dir = tmp_path / "vlm_run" / "step_1000" + + with caplog.at_level(logging.WARNING, logger="vlm_eval_harness"): + vlm_eval_harness._resolve_run_id(config, ckpt_dir) + + assert config.metrics.wandb_run_id == "" + assert config.metrics.wandb_run_name == "vlm_run-step_1000" + assert any("no saved wandb_run_id" in r.getMessage() for r in caplog.records) + + +def test_no_run_id_respects_explicit_run_name(tmp_path: Path, monkeypatch): + monkeypatch.setattr(kempnerforge.checkpoint, "load_train_state_extras", lambda _d: {}) + config = _config(wandb_run_name="my-eval-run") + + vlm_eval_harness._resolve_run_id(config, tmp_path / "run" / "step_1") + + assert config.metrics.wandb_run_name == "my-eval-run" + + +def test_unreadable_train_state_degrades_to_fresh_run(tmp_path: Path, monkeypatch, caplog): + def _foreign(_ckpt_dir): + raise PermissionError("Refusing to load: owned by uid=999") + + monkeypatch.setattr(kempnerforge.checkpoint, "load_train_state_extras", _foreign) + config = _config() + ckpt_dir = tmp_path / "run" / "step_2000" + + with caplog.at_level(logging.WARNING, logger="vlm_eval_harness"): + vlm_eval_harness._resolve_run_id(config, ckpt_dir) + + assert config.metrics.wandb_run_id == "" + assert config.metrics.wandb_run_name == "run-step_2000" + messages = [r.getMessage() for r in caplog.records] + assert any("Could not read" in m for m in messages) + assert any("no saved wandb_run_id" in m for m in messages) + + +# --------------------------------------------------------------------------- # +# _track_eval glue +# --------------------------------------------------------------------------- # + + +def test_track_eval_failure_never_raises(tmp_path: Path, monkeypatch, caplog): + """Any tracking-side exception is downgraded to a warning.""" + + def _boom(_ckpt_dir): + raise RuntimeError("backend exploded") + + monkeypatch.setattr(kempnerforge.checkpoint, "load_train_state_extras", _boom) + monkeypatch.setattr(vlm_eval_harness, "_resolve_run_id", None) # force a TypeError inside + + with caplog.at_level(logging.WARNING, logger="vlm_eval_harness"): + vlm_eval_harness._track_eval(_config(), {"results": {}}, ["t"], tmp_path) + + assert any("Experiment tracking failed" in r.getMessage() for r in caplog.records) diff --git a/examples/vlm-evaluation/vlm_eval_harness.py b/examples/vlm-evaluation/vlm_eval_harness.py index 444684a..99c0cb2 100644 --- a/examples/vlm-evaluation/vlm_eval_harness.py +++ b/examples/vlm-evaluation/vlm_eval_harness.py @@ -30,6 +30,13 @@ --checkpoint checkpoints/vlm/step_10000 \ --tasks mmmu_val,mmbench_en_dev \ --limit 4 + + # Log results to the checkpoint's training run (see README: Experiment tracking) + uv run python examples/vlm-evaluation/vlm_eval_harness.py \ + --config configs/train/vlm_jd.toml \ + --checkpoint checkpoints/vlm/step_10000 \ + --tasks mmmu_val \ + --metrics.enable_wandb=true --metrics.wandb_project=vlm-eval """ from __future__ import annotations @@ -40,6 +47,10 @@ import sys import tempfile from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from kempnerforge.config.schema import JobConfig logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") logger = logging.getLogger(__name__) @@ -57,9 +68,90 @@ def _limit_type(value: str) -> int | float: raise argparse.ArgumentTypeError("--limit must be an integer count, or a fraction < 1.0") +def _resolve_checkpoint(checkpoint: str) -> Path: + """Resolve a checkpoint arg (run dir or step_N dir) to a concrete step directory. + + Mirrors how the adapter loads weights: a run directory resolves to its + latest ``step_N`` via ``resolve_resume_path``. + """ + from kempnerforge.resilience.elastic import resolve_resume_path + + return (resolve_resume_path(checkpoint) or Path(checkpoint)).resolve() + + +def _checkpoint_step(ckpt_dir: Path) -> int: + """The checkpoint's training step: metadata.json, else the step_N dir name, else 0.""" + meta_file = ckpt_dir / "metadata.json" + if meta_file.exists(): + try: + return int(json.loads(meta_file.read_text())["step"]) + except (OSError, json.JSONDecodeError, KeyError, TypeError, ValueError): + logger.warning(f"Could not read step from {meta_file}; falling back to the dir name") + step_suffix = ckpt_dir.name.removeprefix("step_") + if ckpt_dir.name.startswith("step_") and step_suffix.isdigit(): + return int(step_suffix) + logger.warning(f"Could not determine {ckpt_dir}'s training step; logging eval at step 0") + return 0 + + +def _resolve_run_id(config: JobConfig, ckpt_dir: Path) -> None: + """Point ``config.metrics`` at the run these results belong to: an explicit + ``--metrics.wandb_run_id`` override wins, else the id training saved into + the checkpoint, else a fresh run named after the checkpoint. + """ + mc = config.metrics + if mc.wandb_run_id: + return # explicit override (or TOML) wins + from kempnerforge.checkpoint import load_train_state_extras + + run_id = None + try: + run_id = load_train_state_extras(ckpt_dir).get("wandb_run_id") + except Exception as exc: # foreign-owned or corrupt train_state.pt — never fatal here + logger.warning(f"Could not read {ckpt_dir / 'train_state.pt'} ({exc})") + if run_id: + mc.wandb_run_id = run_id + logger.info(f"Attaching eval metrics to the checkpoint's training run ({run_id})") + return + logger.warning( + f"{ckpt_dir} has no saved wandb_run_id — starting a fresh run " + f"(attach to an existing one with --metrics.wandb_run_id=)" + ) + if mc.wandb_run_name is None: + mc.wandb_run_name = f"{ckpt_dir.parent.name}-{ckpt_dir.name}" + + +def _track_eval(config: JobConfig, results: dict, tasks: list[str], ckpt_dir: Path) -> None: + """Log eval metrics through the framework's MetricsTracker backends. + + ``gpu_peak_tflops`` is a nonzero sentinel: eval never computes MFU, and + ``None``/``0.0`` would trigger the GPU probe. A tracking failure never + fails a completed eval. + """ + try: + from benchmark_manifest import build_eval_metrics + + from kempnerforge.metrics.tracker import MetricsTracker + + _resolve_run_id(config, ckpt_dir) + tracker = MetricsTracker(config, num_gpus=1, gpu_peak_tflops=1.0) + tracker.init_backends(config) + tracker.log_eval(build_eval_metrics(results, tasks), step=_checkpoint_step(ckpt_dir)) + tracker.close() + except Exception as exc: # tracking must never fail a completed eval + logger.warning(f"Experiment tracking failed (eval results are unaffected): {exc}") + + def main() -> None: parser = argparse.ArgumentParser( description="Run lmms-eval on a KempnerForge VLM checkpoint", + epilog=( + "Unrecognized --section.key=value arguments are forwarded to the KempnerForge " + "config loader as dotted overrides on --config and apply to both the evaluated " + "model (e.g. --video.max_frames=8) and experiment tracking, which is enabled " + "that way, e.g. --metrics.enable_wandb=true --metrics.wandb_project=vlm-eval " + "(see 'Experiment tracking' in README.md)." + ), ) parser.add_argument( "--config", @@ -99,7 +191,16 @@ def main() -> None: default=128, help="Fallback max new tokens; task gen_kwargs override it (default: 128)", ) - args = parser.parse_args() + args, extra_overrides = parser.parse_known_args() + + # Forwarded --section.key=value overrides layer over the checkpoint TOML; + # unknown keys raise here, before the expensive model build. The merged + # config object is passed to the adapter below, so overrides reach the + # evaluated model, not just experiment tracking. + from kempnerforge.config.loader import load_config + + config = load_config(args.config, cli_args=extra_overrides) + ckpt_dir = _resolve_checkpoint(args.checkpoint) # lmms-eval is optional and undeclared; import lazily with a helpful error. try: @@ -128,7 +229,7 @@ def main() -> None: # the checkpoint config (train.param_dtype). dtype_kwargs = {"dtype": args.dtype} if args.dtype is not None else {} model = KempnerForgeVLM( - config=args.config, + config=config, checkpoint=args.checkpoint, device=args.device, batch_size=args.batch_size, @@ -179,6 +280,12 @@ def main() -> None: json.dump(results, f, indent=2, default=str) logger.info(f"Results saved to {output_path}") + # --- Experiment tracking (opt-in) --- + mc = config.metrics + track = mc.enable_wandb or mc.enable_tensorboard + if track and results is not None and "results" in results: + _track_eval(config, results, args.tasks.split(","), ckpt_dir) + if __name__ == "__main__": main() diff --git a/kempnerforge/checkpoint/__init__.py b/kempnerforge/checkpoint/__init__.py index c8fde27..25ab5ca 100644 --- a/kempnerforge/checkpoint/__init__.py +++ b/kempnerforge/checkpoint/__init__.py @@ -4,10 +4,11 @@ - CheckpointManager: Save/load/cleanup distributed checkpoints - AsyncCheckpointer: Non-blocking checkpoint saves - build_train_state / restore_train_state: State assembly + - load_train_state_extras: Read-only accessor for checkpoint ``extra`` metadata """ from kempnerforge.checkpoint.async_save import AsyncCheckpointer -from kempnerforge.checkpoint.manager import CheckpointManager +from kempnerforge.checkpoint.manager import CheckpointManager, load_train_state_extras from kempnerforge.checkpoint.state import ( build_train_state, get_rng_state, @@ -20,6 +21,7 @@ "CheckpointManager", "build_train_state", "get_rng_state", + "load_train_state_extras", "restore_train_state", "set_rng_state", ] diff --git a/kempnerforge/checkpoint/manager.py b/kempnerforge/checkpoint/manager.py index 7252a0c..a228998 100644 --- a/kempnerforge/checkpoint/manager.py +++ b/kempnerforge/checkpoint/manager.py @@ -28,7 +28,11 @@ ) from kempnerforge.checkpoint.async_save import AsyncCheckpointer -from kempnerforge.checkpoint.state import build_train_state, restore_train_state +from kempnerforge.checkpoint.state import ( + TRAIN_STATE_STANDARD_KEYS, + build_train_state, + restore_train_state, +) from kempnerforge.config.schema import AsyncCheckpointMode, CheckpointConfig logger = logging.getLogger(__name__) @@ -110,6 +114,23 @@ def _load_train_state(path: Path) -> dict[str, Any]: return torch.load(path, map_location="cpu", weights_only=False) +def load_train_state_extras(checkpoint_dir: str | Path) -> dict[str, Any]: + """Read the caller-supplied ``extra`` keys from a checkpoint's train_state.pt. + + Read-only accessor for the metadata training saves via + ``build_train_state(extra=...)``, e.g. ``wandb_run_id``; unlike + ``restore_train_state`` it never applies the saved RNG state. Expects a + concrete ``step_N`` directory and returns ``{}`` when there is no + ``train_state.pt``. Raises ``PermissionError`` for a foreign-owned file + (see ``_load_train_state``). + """ + path = Path(checkpoint_dir) / _TRAIN_STATE_FILE + if not path.exists(): + return {} + state = _load_train_state(path) + return {k: v for k, v in state.items() if k not in TRAIN_STATE_STANDARD_KEYS} + + class CheckpointManager: """Manages save/load/cleanup of distributed checkpoints. diff --git a/kempnerforge/checkpoint/state.py b/kempnerforge/checkpoint/state.py index a02eecb..c3b63ef 100644 --- a/kempnerforge/checkpoint/state.py +++ b/kempnerforge/checkpoint/state.py @@ -17,6 +17,9 @@ logger = logging.getLogger(__name__) +# Keys build_train_state always owns; anything else in the dict is caller ``extra``. +TRAIN_STATE_STANDARD_KEYS = frozenset({"step", "tokens_seen", "rng", "scheduler", "dataloader"}) + def get_rng_state() -> dict[str, Any]: """Capture all RNG states for reproducibility on resume.""" @@ -113,7 +116,6 @@ def restore_train_state( dataloader.load_state_dict(state["dataloader"]) logger.info("Restored dataloader state") - _standard_keys = {"step", "tokens_seen", "rng", "scheduler", "dataloader"} - extra = {k: v for k, v in state.items() if k not in _standard_keys} + extra = {k: v for k, v in state.items() if k not in TRAIN_STATE_STANDARD_KEYS} return step, tokens_seen, extra diff --git a/tests/unit/test_checkpoint_extras.py b/tests/unit/test_checkpoint_extras.py new file mode 100644 index 0000000..73187f1 --- /dev/null +++ b/tests/unit/test_checkpoint_extras.py @@ -0,0 +1,83 @@ +"""Tests for ``load_train_state_extras``: extras round-trip, the UID trust +boundary, and that reading never applies the saved RNG state.""" + +from __future__ import annotations + +import os +import random +from pathlib import Path + +import pytest +import torch + +import kempnerforge.checkpoint.manager as manager_mod +from kempnerforge.checkpoint import build_train_state, load_train_state_extras + + +class _Payload: + """Pickle-time side effect. If ``__reduce__`` runs, the marker file appears.""" + + def __init__(self, marker: Path) -> None: + self._marker = marker + + def __reduce__(self): + return (os.system, (f"touch {self._marker}",)) + + +def _write_train_state(ckpt_dir: Path, extra: dict | None = None) -> None: + ckpt_dir.mkdir(parents=True, exist_ok=True) + state = build_train_state(step=7, tokens_seen=128, extra=extra) + torch.save(state, ckpt_dir / "train_state.pt") + + +class TestLoadTrainStateExtras: + def test_returns_extras_and_strips_standard_keys(self, tmp_path: Path) -> None: + _write_train_state(tmp_path / "step_7", extra={"wandb_run_id": "abc123", "phase_idx": 1}) + + extras = load_train_state_extras(tmp_path / "step_7") + + assert extras == {"wandb_run_id": "abc123", "phase_idx": 1} + + def test_no_extras_returns_empty(self, tmp_path: Path) -> None: + _write_train_state(tmp_path / "step_7") + + assert load_train_state_extras(tmp_path / "step_7") == {} + + def test_missing_train_state_returns_empty(self, tmp_path: Path) -> None: + (tmp_path / "step_7").mkdir() + + assert load_train_state_extras(tmp_path / "step_7") == {} + + def test_does_not_touch_global_rng(self, tmp_path: Path) -> None: + """The saved RNG state must be ignored, not applied (restore_train_state applies it).""" + _write_train_state(tmp_path / "step_7", extra={"wandb_run_id": "abc"}) + + random.seed(1234) # move to a state different from the one captured at save time + before_python = random.getstate() + before_torch = torch.random.get_rng_state() + + load_train_state_extras(tmp_path / "step_7") + + assert random.getstate() == before_python + assert torch.equal(torch.random.get_rng_state(), before_torch) + + def test_foreign_owned_file_raises_before_unpickling(self, tmp_path: Path) -> None: + """The UID gate fires before torch.load, so a planted payload never executes.""" + ckpt_dir = tmp_path / "step_42" + ckpt_dir.mkdir() + marker = tmp_path / "rce_marker" + torch.save( + {"step": 42, "tokens_seen": 0, "rng": {}, "payload": _Payload(marker)}, + ckpt_dir / "train_state.pt", + ) + + real_uid = os.getuid() + orig_getuid = manager_mod.os.getuid + try: + manager_mod.os.getuid = lambda: real_uid + 12345 + with pytest.raises(PermissionError, match="Refusing to load"): + load_train_state_extras(ckpt_dir) + finally: + manager_mod.os.getuid = orig_getuid + + assert not marker.exists(), "payload fired despite the ownership gate"