diff --git a/BUILD.bazel b/BUILD.bazel index 517ada40..fa67818c 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -802,6 +802,7 @@ cc_library( ":basics", ":configs", ":mat", + ":matmul_env", ":threading_context", "//compression:types", "//io", @@ -1206,6 +1207,22 @@ cc_binary( ], ) +cc_library( + name = "model_comparison", + srcs = ["evals/model_comparison.cc"], + hdrs = ["evals/model_comparison.h"], + deps = ["@nlohmann_json//:json"], +) + +cc_test( + name = "model_comparison_test", + srcs = ["evals/model_comparison_test.cc"], + deps = [ + ":model_comparison", + "@googletest//:gtest_main", + ], +) + cc_binary( name = "gemma_mmlu", srcs = ["evals/run_mmlu.cc"], @@ -1213,6 +1230,7 @@ cc_binary( ":args", ":benchmark_helper", ":gemma_lib", + ":model_comparison", "//io", "@highway//:hwy", "@highway//:profiler", diff --git a/CMakeLists.txt b/CMakeLists.txt index d2dd6fe8..0bbe75fd 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -351,6 +351,14 @@ target_link_libraries(benchmarks libgemma hwy hwy_contrib nlohmann_json::nlohman add_executable(debug_prompt evals/debug_prompt.cc) target_link_libraries(debug_prompt libgemma hwy hwy_contrib nlohmann_json::nlohmann_json) +add_library(model_comparison evals/model_comparison.cc) +target_include_directories(model_comparison PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}) +target_link_libraries(model_comparison PRIVATE nlohmann_json::nlohmann_json) + +add_executable(gemma_mmlu evals/run_mmlu.cc) +target_link_libraries(gemma_mmlu PRIVATE libgemma model_comparison hwy hwy_contrib + nlohmann_json::nlohmann_json) + ## Tests set(GEMMA_ENABLE_TESTS OFF CACHE BOOL "Enable Gemma tests") if (GEMMA_ENABLE_TESTS) @@ -359,6 +367,23 @@ enable_testing() find_package(GTest REQUIRED) include(GoogleTest) +# Gemma's SIMD tests need Highway's test helpers even when its own test suite +# is disabled (enabling both suites introduces duplicate target names). +if(NOT TARGET hwy_test) + add_library(hwy_test ${highway_SOURCE_DIR}/hwy/tests/test_util.cc) + target_link_libraries(hwy_test PUBLIC hwy) +endif() + +add_executable(model_comparison_test evals/model_comparison_test.cc) +target_link_libraries(model_comparison_test PRIVATE model_comparison GTest::Main) +gtest_discover_tests(model_comparison_test) +find_package(Python3 3.10 COMPONENTS Interpreter QUIET) +if(Python3_Interpreter_FOUND) + add_test(NAME model_comparison_python_test + COMMAND ${Python3_EXECUTABLE} -B -m unittest discover + -s ${CMAKE_CURRENT_SOURCE_DIR}/evals -p compare*_test.py) +endif() + # Tests use HWY_IS_TEST=1 with `hwy/foreach_target.h`, which expands the test # TU for every attainable SIMD target (including EMU128). libgemma's # foreach_target sources (flash_attention.cc, attention.cc, gemma.cc, vit.cc) diff --git a/deepseek/run_dsv4.cc b/deepseek/run_dsv4.cc index 3ab31fb9..2a66c1cf 100644 --- a/deepseek/run_dsv4.cc +++ b/deepseek/run_dsv4.cc @@ -161,7 +161,7 @@ int Main(int argc, char** argv) { consumed.AbortIfUnconsumed(); ThreadingContext ctx(args.threading); - MatMulEnv env(ctx); + MatMulEnv env(ctx, args.inference.MatMulSchedule()); Gemma gemma(args, ctx); KVCache kv_cache(gemma.Config(), args.inference, ctx.allocator); diff --git a/evals/README.md b/evals/README.md new file mode 100644 index 00000000..82459260 --- /dev/null +++ b/evals/README.md @@ -0,0 +1,152 @@ +# Model comparison and MMLU + +`gemma_mmlu` scores multiple-choice questions and emits machine-readable results. +`compare_models.py` runs a root model and one or more targets **serially**, then +writes `comparison.json` and `comparison.md`. It accepts arbitrary weight files, +CLI arguments, and environment overrides; no quantization method is required. +Models must use the same tokenizer and vocabulary for full-vocabulary KL. + +## Build and test + +From the repository root, with the usual CMake dependencies and GoogleTest: + +```sh +cmake -S . -B build -DCMAKE_BUILD_TYPE=Release \ + -DGEMMA_ENABLE_TESTS=ON -DHWY_ENABLE_TESTS=OFF +cmake --build build --target gemma_mmlu model_comparison_test -j2 +ctest --test-dir build -R 'ModelComparison|model_comparison_python' --output-on-failure +python3 -B -m unittest discover -s evals -p 'compare*_test.py' +``` + +Python 3.10+ is required; the scripts use only the standard library. CMake +registers Python tests when an interpreter is available. Optional cross-entropy +runs also require `cmake --build build --target single_benchmark -j2`. +Bazel targets `//:gemma_mmlu` and `//:model_comparison_test` are also provided. + +## Compare models + +Save a configuration such as `comparison.json`: + +```json +{ + "mmlu": "gemma/evals/mmlu.json", + "max_questions": 0, + "root": { + "name": "root", + "weights": "root.sbs", + "args": ["--num_threads", "6", "--pin", "1"], + "env": {} + }, + "targets": [ + { + "name": "target", + "weights": "target.sbs", + "args": ["--num_threads", "6", "--pin", "1"], + "env": {} + } + ] +} +``` + +Paths for weights, MMLU, and optional `cross_entropy` text are relative to the +configuration file. Extra CLI arguments retain their ordinary interpretation +relative to the runner's working directory. Names must be unique and contain +only letters, digits, underscores, dots, or hyphens, starting with a letter, +digit, or underscore. The runner owns input, output, weights, question-limit, +cross-entropy, and verbosity flags; these cannot be overridden in `args`. +Environment overrides extend the calling process's environment. Record or clear +any inherited settings relevant to the experiment, and use equivalent runtime +settings across models when comparing performance. + +```sh +python3 -B evals/compare_models.py --config comparison.json \ + --build_dir build --output_dir eval-results/run-1 +``` + +The output directory must be new or empty. Each model retains its stdout and +stderr. The root additionally writes `.root-kl.jsonl`, reused by all +targets. A failed process or invalid result stops the series without producing a +successful comparison report; existing logs remain available for diagnosis. +`max_questions: 0` runs all samples; a positive value selects the first N. +To measure run-to-run variation, use the same weights and settings for root +and target. For reproducible MatMul schedules, add `--matmul_schedule fixed` to +both models' `args` arrays. The available modes are: + +- `auto` (default): measure candidates and choose the fastest. +- `fixed`: choose the first legal candidate without comparing timings. +- `fixed_min_k`: choose the legal candidate with the fewest K partitions, + retaining candidate order to break ties. This reduces intermediate output + rounding but may be slower. + +Fixed modes also fix activation-conversion scheduling and, when enabled, the +oneDNN BRGeMM candidate. A single candidate is measured once, then reused. +The policy is set when `MatMulEnv` is constructed and cannot change within that +environment. C++ callers can pass `MMSchedule::kFixed` or +`MMSchedule::kFixedMinK` as its second constructor argument. + +For a self-comparison, keep the binary, SIMD target, thread/topology settings, +batch sizes, weights, and prompts identical. Fixed scheduling removes variation +from timing-based MatMul candidate selection; it does not promise identical +results across hardware or builds. `--deterministic` separately controls sampling. +The comparison report retains each model's scheduling option in its arguments. + +## Scoring and metrics + +Each question uses the dataset prompt, the runner's answer-only instruction, +and the model's chat template. The runner scores the **first generated token**. +It accepts single-token encodings of `A`–`D`, with or without a leading space, +and uses the highest logit among each label's available spellings. A label with +no single-token encoding causes an error. Equal label scores select A before B, +then C, then D. This replaces the previous runner's free-form answer decoding; +results from the two protocols should not be mixed. + +- `MMLU_RESULT`: question ID, expected/predicted label, correctness, four label + logits, their softmax probabilities, and the winning margin. These + probabilities normalize the four label scores; they are not full-vocabulary + token probabilities. Target runs also report `full_vocab_kl`. +- `MMLU_SUMMARY`: number answered, number correct, and accuracy as a fraction. +- `MMLU_KL_SUMMARY`: mean, median, p95, and maximum of + `KL(root || target)` in nats over the **whole vocabulary at the first answer + token**, after the model's logit soft cap. Percentiles use linear interpolation. + This is not KL averaged over every token of all possible answer strings. +- **Flips**: correct-to-incorrect plus incorrect-to-correct transitions. + Wrong-to-wrong changes are separate; `answer_changes` counts all changed + labels. Aggregate accuracy can hide changes in individual answers. +- `MMLU_TIMING`: time in `Generate`, time inside the evaluation sampling callback, + and their difference (`inference_seconds`). The latter includes MatMul + candidate measurements; it excludes prompt construction, reference + serialization/parsing, and result printing. +- **Wall time** includes process startup, model loading, evaluation, and reference + I/O. It is not directly comparable between writing and reading references. + Peak RSS is sampled from Linux `/proc` every 20 ms and can miss short-lived + peaks; unavailable measurements are shown as null/—. +- Optional **cross entropy** comes from `single_benchmark`: total bits divided + by input tokens, including its uniform first-token cost. The report uses the + final cumulative entropy and throughput, and compares bits/token to root. + +The bundled 83-question fixture is a small regression dataset, not the complete +MMLU benchmark. A single run is not a statistically reliable speed measurement. + +## Reference files + +References are versioned JSON Lines, with one metadata header followed by one +record per question. The header contains the exact dataset text and hex-encoded +serialized tokenizer bytes. Records store IDs, expected labels, tokenized +prompts, and float32 logits that round-trip without precision loss. Exact +comparisons avoid hash collisions and reject mismatched datasets, tokenizers, +question limits, or chat-template tokenization. Truncated, extra, malformed, or +non-finite records fail explicitly. + +Text is inspectable and portable, but larger than binary; expect hundreds of +MiB for an 83-question, 262k-vocabulary run. Targets stream one question at a +time. Reference I/O is kept outside the reported `Generate` time. + +For a standalone root/target pair: + +```sh +build/gemma_mmlu --weights root.sbs --input gemma/evals/mmlu.json \ + --reference_out root.jsonl > root.out +build/gemma_mmlu --weights target.sbs --input gemma/evals/mmlu.json \ + --reference_in root.jsonl > target.out +python3 -B evals/compare_mmlu.py root.out target.out +``` diff --git a/evals/benchmark_helper.cc b/evals/benchmark_helper.cc index eabf351d..9a2c6eaa 100644 --- a/evals/benchmark_helper.cc +++ b/evals/benchmark_helper.cc @@ -37,7 +37,9 @@ namespace gcpp { GemmaEnv::GemmaEnv(const GemmaArgs& args) - : ctx_(args.threading), env_(ctx_), gemma_(args, ctx_) { + : ctx_(args.threading), + env_(ctx_, args.inference.MatMulSchedule()), + gemma_(args, ctx_) { const ModelConfig& config = gemma_.Config(); if (args.inference.verbosity >= 2) { diff --git a/evals/compare_mmlu.py b/evals/compare_mmlu.py new file mode 100755 index 00000000..2449f065 --- /dev/null +++ b/evals/compare_mmlu.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python3 +# Copyright 2026 Google LLC +# SPDX-License-Identifier: Apache-2.0 +"""Compare baseline and target gemma_mmlu outputs. + +Each input is the stdout captured from gemma_mmlu and may contain unrelated +lines. Only lines beginning with ``MMLU_RESULT `` are parsed. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + + +RESULT_PREFIX = "MMLU_RESULT " + + +def load_results(path: Path) -> dict[int, dict[str, Any]]: + results: dict[int, dict[str, Any]] = {} + with path.open(encoding="utf-8") as source: + for line_number, line in enumerate(source, start=1): + if not line.startswith(RESULT_PREFIX): + continue + try: + result = json.loads(line[len(RESULT_PREFIX) :]) + if not isinstance(result, dict): + raise ValueError("result must be an object") + question_id = result["id"] + if type(question_id) is not int: + raise ValueError("id must be an integer") + if result["expected"] not in ("A", "B", "C", "D") or result[ + "predicted" + ] not in ("A", "B", "C", "D"): + raise ValueError("expected/predicted must be A, B, C, or D") + if type(result["correct"]) is not bool or result["correct"] != ( + result["expected"] == result["predicted"] + ): + raise ValueError("correct must match expected == predicted") + except (KeyError, TypeError, ValueError, json.JSONDecodeError) as error: + raise ValueError( + f"{path}:{line_number}: invalid MMLU_RESULT: {error}" + ) from error + if question_id in results: + raise ValueError(f"{path}:{line_number}: duplicate id {question_id}") + results[question_id] = result + + if not results: + raise ValueError(f"{path}: no {RESULT_PREFIX.strip()} lines found") + return results + + +def compare_results( + baseline: dict[int, dict[str, Any]], variant: dict[int, dict[str, Any]] +) -> dict[str, int | float]: + if not baseline or not variant: + raise ValueError("cannot compare empty results") + baseline_ids = set(baseline) + variant_ids = set(variant) + if baseline_ids != variant_ids: + missing = sorted(baseline_ids - variant_ids) + extra = sorted(variant_ids - baseline_ids) + raise ValueError( + "result IDs differ: " + f"missing from variant={missing[:10]}, extra in variant={extra[:10]}" + ) + + correct_to_incorrect = 0 + incorrect_to_correct = 0 + wrong_to_wrong_changes = 0 + answer_changes = 0 + baseline_correct = 0 + variant_correct = 0 + + for question_id in sorted(baseline_ids): + base = baseline[question_id] + changed = variant[question_id] + if base["expected"] != changed["expected"]: + raise ValueError( + f"id {question_id}: expected answers differ: " + f"{base['expected']!r} != {changed['expected']!r}" + ) + + base_correct = base["correct"] + changed_correct = changed["correct"] + baseline_correct += int(base_correct) + variant_correct += int(changed_correct) + answer_changed = base["predicted"] != changed["predicted"] + answer_changes += int(answer_changed) + + if base_correct and not changed_correct: + correct_to_incorrect += 1 + elif not base_correct and changed_correct: + incorrect_to_correct += 1 + elif not base_correct and not changed_correct and answer_changed: + wrong_to_wrong_changes += 1 + + samples = len(baseline_ids) + flips = correct_to_incorrect + incorrect_to_correct + return { + "samples": samples, + "baseline_correct": baseline_correct, + "variant_correct": variant_correct, + "baseline_accuracy": baseline_correct / samples, + "variant_accuracy": variant_correct / samples, + "accuracy_delta": (variant_correct - baseline_correct) / samples, + "correct_to_incorrect": correct_to_incorrect, + "incorrect_to_correct": incorrect_to_correct, + "flips": flips, + "flips_fraction": flips / samples, + "flips_percent": 100.0 * flips / samples, + "wrong_to_wrong_changes": wrong_to_wrong_changes, + "answer_changes": answer_changes, + "answer_changes_fraction": answer_changes / samples, + "answer_changes_percent": 100.0 * answer_changes / samples, + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Compare baseline and variant gemma_mmlu output streams." + ) + parser.add_argument("baseline", type=Path, help="baseline gemma_mmlu stdout") + parser.add_argument("variant", type=Path, help="variant gemma_mmlu stdout") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + try: + metrics = compare_results( + load_results(args.baseline), load_results(args.variant) + ) + except (OSError, ValueError) as error: + print(f"error: {error}", file=sys.stderr) + return 1 + + print(f"MMLU_FLIPS {json.dumps(metrics, sort_keys=True)}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/evals/compare_mmlu_test.py b/evals/compare_mmlu_test.py new file mode 100755 index 00000000..07f3b02d --- /dev/null +++ b/evals/compare_mmlu_test.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +# Copyright 2026 Google LLC +# SPDX-License-Identifier: Apache-2.0 + +import json +import tempfile +import unittest +from pathlib import Path + +from compare_mmlu import compare_results, load_results + + +def result(question_id: int, expected: str, predicted: str) -> dict[str, object]: + return { + "id": question_id, + "expected": expected, + "predicted": predicted, + "correct": expected == predicted, + } + + +class CompareMmluTest(unittest.TestCase): + def test_flip_counts(self) -> None: + baseline = { + 1: result(1, "A", "A"), + 2: result(2, "A", "B"), + 3: result(3, "A", "C"), + 4: result(4, "D", "D"), + } + variant = { + 1: result(1, "A", "B"), + 2: result(2, "A", "A"), + 3: result(3, "A", "D"), + 4: result(4, "D", "D"), + } + + metrics = compare_results(baseline, variant) + + self.assertEqual(metrics["correct_to_incorrect"], 1) + self.assertEqual(metrics["incorrect_to_correct"], 1) + self.assertEqual(metrics["flips"], 2) + self.assertEqual(metrics["flips_percent"], 50.0) + self.assertEqual(metrics["wrong_to_wrong_changes"], 1) + self.assertEqual(metrics["answer_changes"], 3) + self.assertEqual(metrics["answer_changes_percent"], 75.0) + self.assertEqual(metrics["accuracy_delta"], 0.0) + + def test_loads_prefixed_results_and_ignores_other_lines(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "run.log" + rows = [result(7, "B", "B"), result(8, "C", "A")] + path.write_text( + "startup noise\n" + + "\n".join(f"MMLU_RESULT {json.dumps(row)}" for row in rows) + + "\nMMLU_SUMMARY {}\n", + encoding="utf-8", + ) + + loaded = load_results(path) + + self.assertEqual(set(loaded), {7, 8}) + self.assertTrue(loaded[7]["correct"]) + self.assertFalse(loaded[8]["correct"]) + + def test_requires_matching_question_ids(self) -> None: + with self.assertRaisesRegex(ValueError, "result IDs differ"): + compare_results({1: result(1, "A", "A")}, {2: result(2, "A", "A")}) + + def test_requires_matching_expected_answers(self) -> None: + with self.assertRaisesRegex(ValueError, "expected answers differ"): + compare_results({1: result(1, "A", "A")}, {1: result(1, "B", "B")}) + + def test_rejects_malformed_results(self) -> None: + good = result(1, "A", "B") + bad_rows = [ + [], + {}, + {**good, "id": True}, + {**good, "id": "1"}, + {**good, "correct": "false"}, + {**good, "correct": True}, + {**good, "predicted": "?"}, + {**good, "expected": "E"}, + ] + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "results.log" + for row in bad_rows: + with self.subTest(row=row): + path.write_text("MMLU_RESULT " + json.dumps(row) + "\n") + with self.assertRaises(ValueError): + load_results(path) + path.write_text(("MMLU_RESULT " + json.dumps(good) + "\n") * 2) + with self.assertRaisesRegex(ValueError, "duplicate"): + load_results(path) + path.write_text("startup only\n") + with self.assertRaisesRegex(ValueError, "no MMLU_RESULT"): + load_results(path) + + def test_rejects_empty_comparison(self) -> None: + with self.assertRaisesRegex(ValueError, "empty"): + compare_results({}, {}) + + +if __name__ == "__main__": + unittest.main() diff --git a/evals/compare_models.py b/evals/compare_models.py new file mode 100644 index 00000000..5d8f2ca2 --- /dev/null +++ b/evals/compare_models.py @@ -0,0 +1,472 @@ +#!/usr/bin/env python3 +# Copyright 2026 Google LLC +# SPDX-License-Identifier: Apache-2.0 +"""Run generic root-vs-target model comparisons and render a report. + +The configuration contains arbitrary model weights, Gemma CLI arguments, and +environment variables. No optimization (W8A8 or otherwise) is special-cased. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import re +import subprocess +import sys +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from compare_mmlu import compare_results, load_results + + +@dataclass(frozen=True) +class ModelSpec: + name: str + weights: Path + args: tuple[str, ...] + env: dict[str, str] + + +@dataclass(frozen=True) +class RunMetrics: + wall_seconds: float + peak_rss_kib: int | None + + +def _resolve_path(value: str, base: Path) -> Path: + path = Path(value) + return path if path.is_absolute() else (base / path).resolve() + + +def parse_model_spec(data: dict[str, Any], base: Path) -> ModelSpec: + if not isinstance(data, dict): + raise ValueError("model must be an object") + try: + name = data["name"] + if not isinstance(data["weights"], str) or not data["weights"]: + raise ValueError("weights must be a nonempty path string") + weights = _resolve_path(data["weights"], base) + except KeyError as error: + raise ValueError(f"model is missing {error.args[0]!r}") from error + if not isinstance(name, str) or not re.fullmatch( + r"[A-Za-z0-9_][A-Za-z0-9_.-]*", name + ): + raise ValueError(f"invalid model name {name!r}") + raw_args = data.get("args", []) + if not isinstance(raw_args, list): + raise ValueError(f"{name}: args must be an array") + if any(type(arg) not in (str, int, float) for arg in raw_args): + raise ValueError(f"{name}: args must contain strings or numbers") + args = tuple(str(arg) for arg in raw_args) + owned = { + "--weights", + "--input", + "--reference_in", + "--reference_out", + "--max_questions", + "--cross_entropy", + "--verbosity", + } + if any(arg.split("=", 1)[0] in owned for arg in args): + raise ValueError(f"{name}: args override a runner-owned option") + raw_env = data.get("env", {}) + if not isinstance(raw_env, dict): + raise ValueError(f"{name}: env must be an object") + if any( + not isinstance(key, str) + or not key + or "=" in key + or "\0" in key + or type(value) not in (str, int, float) + or "\0" in str(value) + for key, value in raw_env.items() + ): + raise ValueError(f"{name}: invalid environment override") + env = {key: str(value) for key, value in raw_env.items()} + return ModelSpec(name=name, weights=weights, args=args, env=env) + + +def _read_rss_kib(pid: int) -> int | None: + try: + status = Path(f"/proc/{pid}/status").read_text(encoding="utf-8") + except OSError: + return None + values: dict[str, int] = {} + for line in status.splitlines(): + if line.startswith(("VmHWM:", "VmRSS:")): + key, value, *_ = line.split() + values[key.rstrip(":")] = int(value) + return values.get("VmHWM", values.get("VmRSS")) + + +def run_command( + command: list[str], + env_updates: dict[str, str], + stdout_path: Path, + stderr_path: Path, +) -> RunMetrics: + env = os.environ.copy() + env.update(env_updates) + start = time.perf_counter() + peak_rss_kib: int | None = None + with stdout_path.open("w", encoding="utf-8") as stdout, stderr_path.open( + "w", encoding="utf-8" + ) as stderr: + process = subprocess.Popen(command, stdout=stdout, stderr=stderr, env=env) + try: + while process.poll() is None: + rss = _read_rss_kib(process.pid) + if rss is not None: + peak_rss_kib = max(peak_rss_kib or 0, rss) + time.sleep(0.02) + finally: + # Interrupting the runner must not leave model processes running. + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + return_code = process.returncode + wall_seconds = time.perf_counter() - start + if return_code != 0: + tail = "\n".join( + stderr_path.read_text(encoding="utf-8", errors="replace").splitlines()[-20:] + ) + raise RuntimeError( + f"command failed ({return_code}): {' '.join(command)}\n{tail}" + ) + return RunMetrics(wall_seconds=wall_seconds, peak_rss_kib=peak_rss_kib) + + +def parse_prefixed_json(path: Path, prefix: str) -> dict[str, Any]: + found: dict[str, Any] | None = None + for line in path.read_text(encoding="utf-8").splitlines(): + if line.startswith(prefix): + if found is not None: + raise ValueError(f"{path}: duplicate {prefix.strip()} line") + found = json.loads(line[len(prefix) :]) + if not isinstance(found, dict): + raise ValueError(f"{path}: {prefix.strip()} must be an object") + if found is None: + raise ValueError(f"{path}: no {prefix.strip()} line") + return found + + +def parse_entropy(path: Path) -> dict[str, float | int]: + text = path.read_text(encoding="utf-8") + token_matches = re.findall(r"Number of input tokens: (\d+)", text) + speed_matches = re.findall(r"\[([0-9.eE+-]+) tokens / sec\]", text) + entropy_matches = re.findall( + r"Total cross entropy: [0-9.eE+-]+ \[cumulative: ([0-9.eE+-]+)\]", + text, + ) + if not token_matches or not speed_matches or not entropy_matches: + raise ValueError(f"{path}: incomplete cross-entropy output") + tokens = int(token_matches[-1]) + if tokens == 0: + raise ValueError(f"{path}: cross-entropy input has no tokens") + total_bits = float(entropy_matches[-1]) + speed = float(speed_matches[-1]) + if ( + not math.isfinite(total_bits) + or total_bits < 0 + or not math.isfinite(speed) + or speed <= 0 + ): + raise ValueError(f"{path}: invalid cross-entropy metrics") + return { + "tokens": tokens, + "total_bits": total_bits, + "bits_per_token": total_bits / tokens, + "tokens_per_second": speed, + } + + +def render_table(rows: list[dict[str, Any]]) -> str: + root = rows[0] + root_entropy = root.get("entropy") + lines = [ + "| Model | Entropy bits/token | Δ entropy | tok/s | Speedup | " + "MMLU accuracy | Flips | Mean KL | p95 KL | Peak RSS | Wall s | Generate s* | Generate speedup |", + "|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|", + ] + for row in rows: + entropy = row.get("entropy") + if entropy and root_entropy: + entropy_delta = ( + 100.0 + * (entropy["bits_per_token"] / root_entropy["bits_per_token"] - 1.0) + if root_entropy["bits_per_token"] + else None + ) + speedup = 100.0 * ( + entropy["tokens_per_second"] / root_entropy["tokens_per_second"] - 1.0 + ) + entropy_text = f"{entropy['bits_per_token']:.4f}" + delta_text = "—" if entropy_delta is None else f"{entropy_delta:+.3f}%" + speed_text = f"{entropy['tokens_per_second']:.2f}" + speedup_text = f"{speedup:+.1f}%" + else: + entropy_text = delta_text = speed_text = speedup_text = "—" + flips = row.get("flips") + kl = row.get("kl") + rss = row.get("peak_rss_kib") + inference = row.get("mmlu_inference_seconds") + root_inference = root.get("mmlu_inference_seconds") + inference_text = "—" if not inference else f"{inference:.3f}" + inference_speedup = ( + "—" + if not inference or not root_inference + else f"{root_inference / inference:.3f}x" + ) + lines.append( + "| {name} | {entropy} | {delta} | {speed} | {speedup} | " + "{accuracy:.1f}% | {flips} | {mean_kl} | {p95_kl} | {rss} | " + "{wall:.3f} | {inference} | {infer_speedup} |".format( + wall=row["mmlu_wall_seconds"], + inference=inference_text, + infer_speedup=inference_speedup, + name=row["name"], + entropy=entropy_text, + delta=delta_text, + speed=speed_text, + speedup=speedup_text, + accuracy=100.0 * row["mmlu"]["accuracy"], + flips="—" if flips is None else f"{flips['flips_percent']:.2f}%", + mean_kl="—" if kl is None else f"{kl['mean']:.6g}", + p95_kl="—" if kl is None else f"{kl['p95']:.6g}", + rss="—" if rss is None else f"{rss / 1024.0:.1f} MiB", + ) + ) + return "\n".join(lines) + ( + "\n\n*Generate time excludes the evaluation sampling callback; includes " + "MatMul autotuning. Wall time includes startup and reference I/O.\n" + "Timing-dependent MatMul schedules can change outputs between runs, " + "including self-comparisons.\n" + "Flips count changes in correctness; wrong-to-wrong answer changes " + "are reported separately in JSON. KL is full-vocabulary root || target " + "at the first answer token, in nats. Peak RSS is sampled on Linux " + "and may miss short-lived peaks.\n" + ) + + +def run_evaluation( + spec: ModelSpec, + build_dir: Path, + output_dir: Path, + mmlu: Path, + max_questions: int, + reference: Path, + is_root: bool, + entropy_path: Path | None, +) -> tuple[dict[str, Any], Path]: + mmlu_out = output_dir / f"{spec.name}.mmlu.out" + mmlu_err = output_dir / f"{spec.name}.mmlu.err" + command = [ + str(build_dir / "gemma_mmlu"), + "--weights", + str(spec.weights), + "--input", + str(mmlu), + "--verbosity", + "0", + ] + if max_questions: + command.extend(["--max_questions", str(max_questions)]) + command.extend(["--reference_out" if is_root else "--reference_in", str(reference)]) + command.extend(spec.args) + mmlu_run = run_command(command, spec.env, mmlu_out, mmlu_err) + mmlu_summary = parse_prefixed_json(mmlu_out, "MMLU_SUMMARY ") + results = load_results(mmlu_out) + correct = sum(row["correct"] for row in results.values()) + if ( + mmlu_summary.get("answers") != len(results) + or mmlu_summary.get("correct") != correct + or not math.isclose(mmlu_summary.get("accuracy", -1), correct / len(results)) + ): + raise ValueError(f"{mmlu_out}: summary does not match question results") + timing = parse_prefixed_json(mmlu_out, "MMLU_TIMING ") + for key in ("generate_seconds", "sample_seconds", "inference_seconds"): + value = timing[key] + if type(value) not in (int, float) or not math.isfinite(value) or value < 0: + raise ValueError(f"{mmlu_out}: invalid {key}") + kl_summary = None if is_root else parse_prefixed_json(mmlu_out, "MMLU_KL_SUMMARY ") + + if kl_summary is not None: + if ( + kl_summary.get("samples") != len(results) + or kl_summary.get("direction") != "root||target" + or kl_summary.get("unit") != "nats" + ): + raise ValueError(f"{mmlu_out}: invalid KL summary") + for value in [kl_summary[key] for key in ("mean", "median", "p95", "max")] + [ + row["full_vocab_kl"] for row in results.values() + ]: + if type(value) not in (int, float) or not math.isfinite(value) or value < 0: + raise ValueError(f"{mmlu_out}: invalid KL value") + + entropy: dict[str, float | int] | None = None + entropy_run: RunMetrics | None = None + if entropy_path is not None: + entropy_out = output_dir / f"{spec.name}.entropy.out" + entropy_err = output_dir / f"{spec.name}.entropy.err" + entropy_command = [ + str(build_dir / "single_benchmark"), + "--weights", + str(spec.weights), + "--cross_entropy", + str(entropy_path), + "--verbosity", + "0", + *spec.args, + ] + entropy_run = run_command(entropy_command, spec.env, entropy_out, entropy_err) + entropy = parse_entropy(entropy_out) + + peak_values = [mmlu_run.peak_rss_kib] + if entropy_run is not None: + peak_values.append(entropy_run.peak_rss_kib) + peak_rss = max((value for value in peak_values if value is not None), default=None) + return ( + { + "name": spec.name, + "weights": str(spec.weights), + "args": list(spec.args), + "env": spec.env, + "mmlu": mmlu_summary, + "kl": kl_summary, + "entropy": entropy, + "mmlu_timing": timing, + "mmlu_inference_seconds": timing["inference_seconds"], + "mmlu_wall_seconds": mmlu_run.wall_seconds, + "entropy_wall_seconds": ( + None if entropy_run is None else entropy_run.wall_seconds + ), + "peak_rss_kib": peak_rss, + }, + mmlu_out, + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Compare arbitrary target models against a root model." + ) + parser.add_argument("--config", required=True, type=Path) + parser.add_argument("--build_dir", type=Path, default=Path("build")) + parser.add_argument("--output_dir", type=Path) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + try: + config_path = args.config.resolve() + config = json.loads(config_path.read_text(encoding="utf-8")) + base = config_path.parent + root = parse_model_spec(config["root"], base) + if not isinstance(config["targets"], list): + raise ValueError("targets must be an array") + targets = [parse_model_spec(item, base) for item in config["targets"]] + if not targets: + raise ValueError("targets must contain at least one model") + names = [root.name, *(target.name for target in targets)] + if len(names) != len(set(names)): + raise ValueError("model names must be unique") + mmlu = _resolve_path(str(config["mmlu"]), base) + entropy_path = ( + None + if not config.get("cross_entropy") + else _resolve_path(str(config["cross_entropy"]), base) + ) + max_questions = config.get("max_questions", 0) + if type(max_questions) is not int or max_questions < 0: + raise ValueError("max_questions must be non-negative") + build_dir = args.build_dir.resolve() + required_paths = [mmlu, root.weights] + required_paths.extend(target.weights for target in targets) + if entropy_path is not None: + required_paths.append(entropy_path) + missing = [str(path) for path in required_paths if not path.is_file()] + if missing: + raise ValueError("file does not exist: " + ", ".join(missing)) + required_programs = [build_dir / "gemma_mmlu"] + if entropy_path is not None: + required_programs.append(build_dir / "single_benchmark") + missing_programs = [ + str(path) for path in required_programs if not path.is_file() + ] + if missing_programs: + raise ValueError( + "build executable does not exist: " + ", ".join(missing_programs) + ) + output_dir = ( + args.output_dir.resolve() + if args.output_dir + else (base / f"{config_path.stem}-results").resolve() + ) + if output_dir.exists() and any(output_dir.iterdir()): + raise ValueError("output directory must be empty; choose a new output_dir") + output_dir.mkdir(parents=True, exist_ok=True) + reference = output_dir / f"{root.name}.root-kl.jsonl" + + root_row, root_output = run_evaluation( + root, + build_dir, + output_dir, + mmlu, + max_questions, + reference, + True, + entropy_path, + ) + root_row["flips"] = None + root_results = load_results(root_output) + rows = [root_row] + for target in targets: + row, target_output = run_evaluation( + target, + build_dir, + output_dir, + mmlu, + max_questions, + reference, + False, + entropy_path, + ) + row["flips"] = compare_results(root_results, load_results(target_output)) + rows.append(row) + + report = { + "schema_version": 1, + "max_questions": max_questions, + "mmlu": str(mmlu), + "cross_entropy": None if entropy_path is None else str(entropy_path), + "reference": str(reference), + "models": rows, + } + (output_dir / "comparison.json").write_text( + json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + table = render_table(rows) + (output_dir / "comparison.md").write_text(table, encoding="utf-8") + print(table, end="") + return 0 + except (KeyError, TypeError, OSError, ValueError, RuntimeError) as error: + print(f"error: {error}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except KeyboardInterrupt: + print("error: comparison interrupted", file=sys.stderr) + raise SystemExit(130) diff --git a/evals/compare_models_test.py b/evals/compare_models_test.py new file mode 100644 index 00000000..dadcd67c --- /dev/null +++ b/evals/compare_models_test.py @@ -0,0 +1,319 @@ +#!/usr/bin/env python3 +# Copyright 2026 Google LLC +# SPDX-License-Identifier: Apache-2.0 + +import contextlib +import io +import json +import sys +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch, Mock + +from compare_models import ( + parse_entropy, + parse_model_spec, + parse_prefixed_json, + render_table, + main, + RunMetrics, + run_command, + run_evaluation, +) + + +class CompareModelsTest(unittest.TestCase): + def test_parse_model_spec_resolves_paths_and_environment(self) -> None: + spec = parse_model_spec( + { + "name": "target.v1", + "weights": "models/model.sbs", + "args": ["--num_threads", 4], + "env": {"EXAMPLE_MODE": 1}, + }, + Path("/work"), + ) + + self.assertEqual(spec.name, "target.v1") + self.assertEqual(spec.weights, Path("/work/models/model.sbs")) + self.assertEqual(spec.args, ("--num_threads", "4")) + self.assertEqual(spec.env, {"EXAMPLE_MODE": "1"}) + + def test_parse_prefixed_json_rejects_duplicate_summary(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "output.log" + path.write_text("MMLU_SUMMARY {}\nMMLU_SUMMARY {}\n") + with self.assertRaisesRegex(ValueError, "duplicate"): + parse_prefixed_json(path, "MMLU_SUMMARY ") + + def test_parse_entropy(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "entropy.log" + path.write_text( + "Number of input tokens: 8\n" + "Took 1.0 s [8.0 tokens / sec]\n" + "Total cross entropy: 12.0 [cumulative: 12.0]\n", + encoding="utf-8", + ) + parsed = parse_entropy(path) + + self.assertEqual(parsed["tokens"], 8) + self.assertEqual(parsed["total_bits"], 12.0) + self.assertEqual(parsed["bits_per_token"], 1.5) + self.assertEqual(parsed["tokens_per_second"], 8.0) + + def test_render_table(self) -> None: + root = { + "name": "root", + "mmlu_inference_seconds": 8.0, + "mmlu_wall_seconds": 12.0, + "mmlu": {"accuracy": 0.5}, + "entropy": { + "total_bits": 20.0, + "bits_per_token": 2.0, + "tokens_per_second": 10.0, + }, + "flips": None, + "kl": None, + "peak_rss_kib": 1024, + } + target = { + "name": "target", + "mmlu_inference_seconds": 4.0, + "mmlu_wall_seconds": 20.0, + "mmlu": {"accuracy": 0.75}, + "entropy": { + "total_bits": 22.0, + "bits_per_token": 2.2, + "tokens_per_second": 12.0, + }, + "flips": {"flips_percent": 25.0}, + "kl": {"mean": 0.01, "p95": 0.03}, + "peak_rss_kib": 2048, + } + + table = render_table([root, target]) + + self.assertIn("| 20.000 | 4.000 | 2.000x |", table) + self.assertIn("| root | 2.0000 | +0.000% | 10.00 | +0.0%", table) + self.assertIn( + "| target | 2.2000 | +10.000% | 12.00 | +20.0% | " + "75.0% | 25.00% | 0.01 | 0.03 | 2.0 MiB |", + table, + ) + + def test_rejects_unsafe_report_name(self) -> None: + with self.assertRaisesRegex(ValueError, "invalid model name"): + parse_model_spec( + {"name": "../target", "weights": "model.sbs"}, Path("/work") + ) + + def test_rejects_string_model_args(self) -> None: + with self.assertRaisesRegex(ValueError, "args must be an array"): + parse_model_spec( + {"name": "target", "weights": "model.sbs", "args": "--foo"}, + Path("/work"), + ) + + def test_rejects_names_and_owned_options(self) -> None: + for name in (".", "..", "", "../x", "a/b", "a|b", None): + with self.subTest(name=name), self.assertRaises(ValueError): + parse_model_spec({"name": name, "weights": "x"}, Path("/tmp")) + for option in ("--input", "--reference_out=x", "--weights", "--max_questions"): + with self.subTest(option=option), self.assertRaises(ValueError): + parse_model_spec( + {"name": "ok", "weights": "x", "args": [option]}, Path("/tmp") + ) + + def test_child_process_success_and_failure(self) -> None: + with tempfile.TemporaryDirectory() as directory: + out = Path(directory) / "out" + err = Path(directory) / "err" + metrics = run_command( + [sys.executable, "-c", "import os; print(os.environ['EVAL_TEST'])"], + {"EVAL_TEST": "value"}, + out, + err, + ) + self.assertEqual(out.read_text().strip(), "value") + self.assertGreater(metrics.wall_seconds, 0) + with self.assertRaisesRegex(RuntimeError, r"failed \(3\)"): + run_command([sys.executable, "-c", "raise SystemExit(3)"], {}, out, err) + + def test_interrupt_terminates_child(self) -> None: + child = Mock() + child.poll.return_value = None + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) + with patch("compare_models.subprocess.Popen", return_value=child), patch( + "compare_models._read_rss_kib", side_effect=KeyboardInterrupt + ), self.assertRaises(KeyboardInterrupt): + run_command(["fake"], {}, path / "out", path / "err") + child.terminate.assert_called_once() + child.wait.assert_called_once_with(timeout=5) + + def test_runner_report_and_dotted_names(self) -> None: + with tempfile.TemporaryDirectory() as directory: + base = Path(directory) + (base / "weights.sbs").touch() + (base / "mmlu.json").write_text("{}") + (base / "gemma_mmlu").touch() + config = { + "mmlu": "mmlu.json", + "max_questions": 1, + "root": {"name": "root.v1", "weights": "weights.sbs"}, + "targets": [ + {"name": "target.v1", "weights": "weights.sbs"}, + {"name": "target.v2", "weights": "weights.sbs"}, + ], + } + config_path = base / "config.json" + config_path.write_text(json.dumps(config)) + output = base / "results" + calls = [] + + def fake_run(command, env, stdout_path, stderr_path): + calls.append(stdout_path.name) + is_root = "--reference_out" in command + ref_flag = "--reference_out" if is_root else "--reference_in" + reference = Path(command[command.index(ref_flag) + 1]) + if is_root: + reference.write_text("reference") + else: + self.assertEqual(reference.read_text(), "reference") + self.assertEqual( + command[command.index("--weights") + 1], str(base / "weights.sbs") + ) + result = { + "id": 1, + "expected": "A", + "predicted": "A" if is_root else "B", + "correct": is_root, + } + records = { + "MMLU_RESULT": result, + "MMLU_SUMMARY": { + "answers": 1, + "correct": int(is_root), + "accuracy": float(is_root), + }, + "MMLU_TIMING": { + "generate_seconds": 2, + "sample_seconds": 1, + "inference_seconds": 1, + }, + } + if not is_root: + result["full_vocab_kl"] = 0.1 + records["MMLU_KL_SUMMARY"] = { + "samples": 1, + "mean": 0.1, + "median": 0.1, + "p95": 0.1, + "max": 0.1, + "direction": "root||target", + "unit": "nats", + } + stdout_path.write_text( + "".join(k + " " + json.dumps(v) + "\n" for k, v in records.items()) + ) + stderr_path.touch() + return RunMetrics(3.0, 1024) + + argv = [ + "compare_models.py", + "--config", + str(config_path), + "--build_dir", + str(base), + "--output_dir", + str(output), + ] + with patch.object(sys, "argv", argv), patch( + "compare_models.run_command", side_effect=fake_run + ), contextlib.redirect_stdout(io.StringIO()): + self.assertEqual(main(), 0) + self.assertEqual( + calls, ["root.v1.mmlu.out", "target.v1.mmlu.out", "target.v2.mmlu.out"] + ) + report = json.loads((output / "comparison.json").read_text()) + self.assertEqual(report["models"][1]["flips"]["correct_to_incorrect"], 1) + self.assertEqual(report["models"][2]["kl"]["mean"], 0.1) + table = (output / "comparison.md").read_text().splitlines() + self.assertEqual(len(table[0].split("|")), len(table[1].split("|"))) + # Reruns cannot silently overwrite a report or mix stale artifacts. + with patch.object(sys, "argv", argv), contextlib.redirect_stderr( + io.StringIO() + ): + self.assertEqual(main(), 1) + + def test_rejects_incomplete_or_invalid_run(self) -> None: + with tempfile.TemporaryDirectory() as directory: + base = Path(directory) + spec = parse_model_spec({"name": "test", "weights": "model"}, base) + records = { + "MMLU_RESULT": { + "id": 1, + "expected": "A", + "predicted": "A", + "correct": True, + }, + "MMLU_SUMMARY": {"answers": 1, "correct": 1, "accuracy": 1.0}, + "MMLU_TIMING": { + "generate_seconds": 2, + "sample_seconds": 1, + "inference_seconds": 1, + }, + } + for change in ( + {"MMLU_SUMMARY": None}, + {"MMLU_SUMMARY": {"answers": 2, "correct": 1, "accuracy": 0.5}}, + { + "MMLU_TIMING": { + "generate_seconds": 2, + "sample_seconds": 1, + "inference_seconds": float("nan"), + } + }, + ): + changed = {**records, **change} + + def fake_run(command, env, out, err): + out.write_text( + "".join( + key + " " + json.dumps(value) + "\n" + for key, value in changed.items() + if value is not None + ) + ) + return RunMetrics(3, None) + + with self.subTest(change=change), patch( + "compare_models.run_command", side_effect=fake_run + ), self.assertRaises(ValueError): + run_evaluation( + spec, base, base, base / "mmlu", 1, base / "ref", True, None + ) + + def test_entropy_uses_final_cumulative_values(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "entropy.log" + path.write_text( + "Number of input tokens: 8\n" + "Took 1 s [4 tokens / sec]\nTotal cross entropy: 3 [cumulative: 3]\n" + "Took 2 s [4 tokens / sec]\nTotal cross entropy: 5 [cumulative: 8]\n" + ) + self.assertEqual(parse_entropy(path)["bits_per_token"], 1) + for text in ( + "", + "Number of input tokens: 0\n[1 tokens / sec]\nTotal cross entropy: 0 [cumulative: 0]", + "Number of input tokens: 8\n[0 tokens / sec]\nTotal cross entropy: 1 [cumulative: 1]", + ): + path.write_text(text) + with self.assertRaises(ValueError): + parse_entropy(path) + + +if __name__ == "__main__": + unittest.main() diff --git a/evals/model_comparison.cc b/evals/model_comparison.cc new file mode 100644 index 00000000..117020ba --- /dev/null +++ b/evals/model_comparison.cc @@ -0,0 +1,238 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 + +#include "evals/model_comparison.h" + +#include +#include +#include +#include + +#include "nlohmann/json.hpp" + +namespace gcpp { +namespace { + +using Json = nlohmann::json; + +// Keep the normalizer relative to the maximum: adding log(sum) to a very +// large maximum would lose the normalization, even in double precision. +double MaxLogit(const float* logits, size_t size) { + if (size == 0) throw std::invalid_argument("empty logits"); + for (size_t i = 0; i < size; ++i) { + if (!std::isfinite(logits[i])) { + throw std::invalid_argument("non-finite logit"); + } + } + return *std::max_element(logits, logits + size); +} + +double LogSumRelative(const float* logits, size_t size, double maximum) { + double sum = 0.0; + for (size_t i = 0; i < size; ++i) { + sum += std::exp(static_cast(logits[i]) - maximum); + } + return std::log(sum); +} + +template +T Integer(const Json& value) { + if (!value.is_number_integer() || value < std::numeric_limits::lowest() || + value > std::numeric_limits::max()) { + throw std::invalid_argument("invalid integer in KL reference"); + } + return value.get(); +} + +Json ReadLine(std::ifstream& stream) { + std::string line; + if (!std::getline(stream, line)) { + throw std::runtime_error("truncated KL reference"); + } + return Json::parse(line); +} + +void ValidateRecord(const ModelComparisonRecord& record, size_t vocab_size) { + if (record.expected_label < 0 || record.expected_label > 3 || + record.prompt.empty() || record.logits.size() != vocab_size) { + throw std::invalid_argument("invalid KL reference record"); + } + for (int token : record.prompt) { + if (token < 0 || static_cast(token) >= vocab_size) { + throw std::invalid_argument("invalid KL reference prompt token"); + } + } + MaxLogit(record.logits.data(), record.logits.size()); +} + +} // namespace + +std::string ModelComparisonHex(const std::string& bytes) { + constexpr char digits[] = "0123456789abcdef"; + std::string text; + text.reserve(bytes.size() * 2); + for (unsigned char byte : bytes) { + text.push_back(digits[byte >> 4]); + text.push_back(digits[byte & 15]); + } + return text; +} + +double FullVocabKLDivergence(const std::vector& root_logits, + const float* target_logits, size_t size) { + if (root_logits.size() != size) { + throw std::invalid_argument("root/target vocabulary size mismatch"); + } + const double root_max = MaxLogit(root_logits.data(), size); + const double target_max = MaxLogit(target_logits, size); + const double root_norm = LogSumRelative(root_logits.data(), size, root_max); + const double target_norm = LogSumRelative(target_logits, size, target_max); + double kl = 0.0; + for (size_t i = 0; i < size; ++i) { + const double p = + (static_cast(root_logits[i]) - root_max) - root_norm; + const double q = + (static_cast(target_logits[i]) - target_max) - target_norm; + kl += std::exp(p) * (p - q); + } + // The exact result is nonnegative; roundoff can produce tiny negatives. + return std::max(0.0, kl); +} + +MmluAnswer ScoreMmluAnswer(const float* logits, size_t size, + const std::vector>& tokens) { + MaxLogit(logits, size); // Reject invalid model outputs before reporting. + MmluAnswer result; + result.logits.fill(-std::numeric_limits::infinity()); + std::array best_tokens = {-1, -1, -1, -1}; + for (const auto& [token, label] : tokens) { + if (token < 0 || static_cast(token) >= size || label < 0 || + label > 3) { + throw std::invalid_argument("invalid MMLU answer token/label"); + } + if (logits[token] > result.logits[label]) { + result.logits[label] = logits[token]; + best_tokens[label] = token; + } + } + for (int token : best_tokens) { + if (token == -1) throw std::invalid_argument("missing MMLU answer label"); + } + result.label = static_cast( + std::max_element(result.logits.begin(), result.logits.end()) - + result.logits.begin()); + result.token = best_tokens[result.label]; + const double maximum = result.logits[result.label]; + double sum = 0.0; + double second = -std::numeric_limits::infinity(); + for (int label = 0; label < 4; ++label) { + result.probabilities[label] = + std::exp(static_cast(result.logits[label]) - maximum); + sum += result.probabilities[label]; + if (label != result.label) + second = std::max(second, static_cast(result.logits[label])); + } + for (double& probability : result.probabilities) probability /= sum; + result.margin = maximum - second; + return result; +} + +ModelComparisonWriter::ModelComparisonWriter( + const std::string& path, const ModelComparisonMetadata& metadata) + : stream_(path), metadata_(metadata) { + if (!stream_) throw std::runtime_error("cannot create KL reference: " + path); + if (metadata.vocab_size == 0 || metadata.sample_count == 0) { + throw std::invalid_argument("empty KL reference metadata"); + } + const Json header = {{"format", "gemma-mmlu-reference"}, + {"version", 1}, + {"vocab_size", metadata.vocab_size}, + {"sample_count", metadata.sample_count}, + {"dataset", metadata.dataset}, + {"tokenizer_hex", metadata.tokenizer_hex}}; + stream_ << header.dump() << '\n'; + if (!stream_) throw std::runtime_error("failed to write KL reference header"); +} + +void ModelComparisonWriter::Write(const ModelComparisonRecord& record) { + if (finished_ || records_written_ >= metadata_.sample_count) { + throw std::runtime_error("too many KL reference writes"); + } + ValidateRecord(record, metadata_.vocab_size); + const Json row = {{"id", record.sample_id}, + {"expected", record.expected_label}, + {"prompt", record.prompt}, + {"logits", record.logits}}; + stream_ << row.dump() << '\n'; + if (!stream_) throw std::runtime_error("failed to write KL reference logits"); + ++records_written_; +} + +void ModelComparisonWriter::Finish() { + if (finished_) return; + if (records_written_ != metadata_.sample_count) { + throw std::runtime_error("KL reference record count mismatch"); + } + stream_.close(); + if (!stream_) throw std::runtime_error("failed to finish KL reference"); + finished_ = true; +} + +ModelComparisonReader::ModelComparisonReader(const std::string& path) + : stream_(path) { + if (!stream_) throw std::runtime_error("cannot open KL reference: " + path); + const Json header = ReadLine(stream_); + if (header.at("format") != "gemma-mmlu-reference" || + Integer(header.at("version")) != 1) { + throw std::runtime_error("unsupported KL reference format/version"); + } + metadata_.vocab_size = Integer(header.at("vocab_size")); + metadata_.sample_count = Integer(header.at("sample_count")); + metadata_.dataset = header.at("dataset").get(); + metadata_.tokenizer_hex = header.at("tokenizer_hex").get(); + if (metadata_.vocab_size == 0 || metadata_.sample_count == 0) { + throw std::runtime_error("empty KL reference metadata"); + } +} + +void ModelComparisonReader::Validate( + const ModelComparisonMetadata& expected) const { + if (metadata_.vocab_size != expected.vocab_size || + metadata_.sample_count != expected.sample_count || + metadata_.dataset != expected.dataset || + metadata_.tokenizer_hex != expected.tokenizer_hex) { + throw std::runtime_error( + "KL reference metadata mismatch (vocabulary, sample count, dataset, or " + "tokenizer)"); + } +} + +ModelComparisonRecord ModelComparisonReader::Read() { + if (records_read_ >= metadata_.sample_count) { + throw std::runtime_error("too many KL reference reads"); + } + const Json row = ReadLine(stream_); + ModelComparisonRecord record; + record.sample_id = Integer(row.at("id")); + record.expected_label = Integer(row.at("expected")); + if (!row.at("prompt").is_array()) { + throw std::invalid_argument("invalid KL reference prompt"); + } + for (const auto& token : row.at("prompt")) + record.prompt.push_back(Integer(token)); + record.logits = row.at("logits").get>(); + ValidateRecord(record, metadata_.vocab_size); + ++records_read_; + return record; +} + +void ModelComparisonReader::Finish() { + if (records_read_ != metadata_.sample_count) { + throw std::runtime_error("unread KL reference records"); + } + if (stream_.peek() != std::ifstream::traits_type::eof()) { + throw std::runtime_error("trailing data in KL reference"); + } +} + +} // namespace gcpp diff --git a/evals/model_comparison.h b/evals/model_comparison.h new file mode 100644 index 00000000..7e5283b3 --- /dev/null +++ b/evals/model_comparison.h @@ -0,0 +1,85 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 + +#ifndef THIRD_PARTY_GEMMA_CPP_EVALS_MODEL_COMPARISON_H_ +#define THIRD_PARTY_GEMMA_CPP_EVALS_MODEL_COMPARISON_H_ + +#include +#include + +#include +#include +#include +#include +#include + +namespace gcpp { + +// Exact, text-safe representation of serialized tokenizer bytes (no hashing). +std::string ModelComparisonHex(const std::string& bytes); + +// D_KL(root || target), in nats, over the entire vocabulary. Requires finite +// logits and matching, nonempty vocabularies. Normalizes in double precision. +double FullVocabKLDivergence(const std::vector& root_logits, + const float* target_logits, size_t size); + +struct MmluAnswer { + int token = -1; + int label = -1; + std::array logits; + std::array probabilities; + double margin = 0.0; +}; + +// Scores each label by its highest-logit token spelling. Probabilities are a +// softmax over those four scores, not sums of token probabilities. Ties select +// the lowest label. Each label needs at least one in-vocabulary token. +MmluAnswer ScoreMmluAnswer(const float* logits, size_t size, + const std::vector>& tokens); + +struct ModelComparisonMetadata { + uint32_t vocab_size = 0; + uint64_t sample_count = 0; + std::string dataset; + std::string tokenizer_hex; +}; + +struct ModelComparisonRecord { + int64_t sample_id = 0; + int32_t expected_label = 0; + std::vector prompt; + std::vector logits; +}; + +// Versioned JSON Lines: one metadata header, then one record per question. +// Float32 logits round-trip exactly. Readers hold only one question at a time. +class ModelComparisonWriter { + public: + ModelComparisonWriter(const std::string& path, + const ModelComparisonMetadata& metadata); + void Write(const ModelComparisonRecord& record); + void Finish(); + + private: + std::ofstream stream_; + ModelComparisonMetadata metadata_; + uint64_t records_written_ = 0; + bool finished_ = false; +}; + +class ModelComparisonReader { + public: + explicit ModelComparisonReader(const std::string& path); + void Validate(const ModelComparisonMetadata& expected) const; + ModelComparisonRecord Read(); + void Finish(); + + private: + std::ifstream stream_; + ModelComparisonMetadata metadata_; + uint64_t records_read_ = 0; +}; + +} // namespace gcpp + +#endif // THIRD_PARTY_GEMMA_CPP_EVALS_MODEL_COMPARISON_H_ diff --git a/evals/model_comparison_test.cc b/evals/model_comparison_test.cc new file mode 100644 index 00000000..801450ce --- /dev/null +++ b/evals/model_comparison_test.cc @@ -0,0 +1,253 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 + +#include "evals/model_comparison.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +namespace gcpp { +namespace { + +TEST(ModelComparisonMath, KnownDirectionalKL) { + const std::vector root = {std::log(0.25f), std::log(0.75f)}; + const std::vector uniform = {0.0f, 0.0f}; + EXPECT_NEAR(FullVocabKLDivergence(root, uniform.data(), 2), + 0.25 * std::log(0.5) + 0.75 * std::log(1.5), 1E-7); + EXPECT_NEAR(FullVocabKLDivergence(uniform, root.data(), 2), + 0.5 * std::log(2.0) + 0.5 * std::log(2.0 / 3.0), 1E-7); +} + +TEST(ModelComparisonMath, IdentityAndShiftInvariance) { + const std::vector root = {1.0f, -2.0f, 4.0f}; + const std::vector shifted = {1001.0f, 998.0f, 1004.0f}; + EXPECT_DOUBLE_EQ(FullVocabKLDivergence(root, root.data(), 3), 0.0); + EXPECT_DOUBLE_EQ(FullVocabKLDivergence(root, shifted.data(), 3), 0.0); + const float largest = std::numeric_limits::max(); + const std::vector extreme = {largest, -largest}; + const std::vector uniform = {largest, largest}; + EXPECT_NEAR(FullVocabKLDivergence(extreme, uniform.data(), 2), std::log(2.0), + 1E-12); + EXPECT_TRUE(std::isfinite(FullVocabKLDivergence(uniform, extreme.data(), 2))); +} + +TEST(ModelComparisonMath, RejectsInvalidLogits) { + const std::vector root = {0.0f, 1.0f}; + EXPECT_THROW(FullVocabKLDivergence({}, root.data(), 0), + std::invalid_argument); + EXPECT_THROW(FullVocabKLDivergence(root, root.data(), 1), + std::invalid_argument); + for (float value : {std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity()}) { + const std::vector bad = {0.0f, value}; + EXPECT_THROW(FullVocabKLDivergence(root, bad.data(), 2), + std::invalid_argument); + EXPECT_THROW(FullVocabKLDivergence(bad, root.data(), 2), + std::invalid_argument); + } +} + +TEST(ModelComparisonMath, ExactTokenizerEncoding) { + EXPECT_EQ(ModelComparisonHex(std::string("\0\x80\xff\n", 4)), "0080ff0a"); + EXPECT_EQ(ModelComparisonHex(""), ""); +} + +TEST(ModelComparisonMath, AnswerSpellingsAndProbabilities) { + const std::vector logits = {0, 1, 2, 3, 4, 100}; + const auto answer = ScoreMmluAnswer(logits.data(), logits.size(), + {{0, 0}, {1, 1}, {2, 2}, {3, 3}, {4, 0}}); + EXPECT_EQ(answer.label, 0); + EXPECT_EQ(answer.token, 4); // The alternate spelling of A wins. + EXPECT_EQ(answer.logits, (std::array{4, 1, 2, 3})); + EXPECT_DOUBLE_EQ(answer.margin, 1.0); + double sum = 0.0; + for (double probability : answer.probabilities) sum += probability; + EXPECT_NEAR(sum, 1.0, 1E-15); + EXPECT_NEAR(answer.probabilities[0], + 1.0 / (1 + std::exp(-1.0) + std::exp(-2.0) + std::exp(-3.0)), + 1E-15); + // The much larger non-answer token must not participate in label scoring. +} + +TEST(ModelComparisonMath, AnswerTiesAndInvalidTokens) { + const std::vector logits(4, 1.0f); + const std::vector> tokens = { + {3, 3}, {2, 2}, {1, 1}, {0, 0}}; + const auto answer = ScoreMmluAnswer(logits.data(), 4, tokens); + EXPECT_EQ(answer.label, 0); + EXPECT_DOUBLE_EQ(answer.margin, 0.0); + for (double p : answer.probabilities) EXPECT_DOUBLE_EQ(p, 0.25); + EXPECT_THROW(ScoreMmluAnswer(logits.data(), 4, {{0, 0}}), + std::invalid_argument); + EXPECT_THROW(ScoreMmluAnswer(logits.data(), 4, {{-1, 0}}), + std::invalid_argument); + EXPECT_THROW(ScoreMmluAnswer(logits.data(), 4, {{4, 0}}), + std::invalid_argument); + EXPECT_THROW(ScoreMmluAnswer(logits.data(), 4, {{0, 4}}), + std::invalid_argument); +} + +class ModelComparisonReference : public testing::Test { + protected: + void SetUp() override { + // Atomic directory creation prevents collisions between concurrent tests. + std::random_device random; + for (size_t attempt = 0; attempt < 100; ++attempt) { + directory_ = std::filesystem::temp_directory_path() / + ("gemma-comparison-" + std::to_string(random())); + if (std::filesystem::create_directory(directory_)) { + path_ = (directory_ / "reference.jsonl").string(); + return; + } + } + FAIL() << "cannot create temporary test directory"; + } + void TearDown() override { std::filesystem::remove_all(directory_); } + void WriteReference() { + ModelComparisonWriter writer(path_, metadata_); + writer.Write(record_); + writer.Finish(); + } + std::string ReadText() { + std::ifstream input(path_); + return std::string(std::istreambuf_iterator(input), {}); + } + void ReplaceText(const std::string& text) { std::ofstream(path_) << text; } + + std::filesystem::path directory_; + std::string path_; + ModelComparisonMetadata metadata_{3, 1, "dataset\n", "0080ff"}; + ModelComparisonRecord record_{7, 2, {1, 2}, {1.25f, -2.0f, 4.0f}}; +}; + +TEST_F(ModelComparisonReference, TextRoundTrip) { + record_.logits = {std::numeric_limits::max(), + std::numeric_limits::denorm_min(), -0.0f}; + WriteReference(); + EXPECT_EQ(ReadText().front(), '{'); + ModelComparisonReader reader(path_); + reader.Validate(metadata_); + const auto actual = reader.Read(); + EXPECT_EQ(actual.sample_id, record_.sample_id); + EXPECT_EQ(actual.expected_label, record_.expected_label); + EXPECT_EQ(actual.prompt, record_.prompt); + EXPECT_EQ(actual.logits, record_.logits); + EXPECT_TRUE(std::signbit(actual.logits.back())); + reader.Finish(); + EXPECT_THROW(reader.Read(), std::runtime_error); +} + +TEST_F(ModelComparisonReference, RejectsMismatchedMetadata) { + WriteReference(); + ModelComparisonReader reader(path_); + auto changed = metadata_; + ++changed.vocab_size; + EXPECT_THROW(reader.Validate(changed), std::runtime_error); + changed = metadata_; + ++changed.sample_count; + EXPECT_THROW(reader.Validate(changed), std::runtime_error); + changed = metadata_; + changed.dataset += " "; + EXPECT_THROW(reader.Validate(changed), std::runtime_error); + changed = metadata_; + changed.tokenizer_hex += "00"; + EXPECT_THROW(reader.Validate(changed), std::runtime_error); +} + +TEST_F(ModelComparisonReference, RejectsMissingAndExtraRecords) { + { + ModelComparisonWriter writer(path_, metadata_); + EXPECT_THROW(writer.Finish(), std::runtime_error); + writer.Write(record_); + EXPECT_THROW(writer.Write(record_), std::runtime_error); + writer.Finish(); + writer.Finish(); + EXPECT_THROW(writer.Write(record_), std::runtime_error); + } + ModelComparisonReader reader(path_); + EXPECT_THROW(reader.Finish(), std::runtime_error); + reader.Read(); + reader.Finish(); +} + +TEST_F(ModelComparisonReference, RejectsTruncatedAndTrailingData) { + WriteReference(); + const auto original = ReadText(); + ReplaceText(original.substr(0, original.find('\n') + 1)); + { + ModelComparisonReader reader(path_); + EXPECT_THROW(reader.Read(), std::exception); + } + ReplaceText(original.substr(0, original.size() - 5)); + { + ModelComparisonReader reader(path_); + EXPECT_THROW(reader.Read(), std::exception); + } + ReplaceText(original + "{}\n"); + { + ModelComparisonReader reader(path_); + reader.Read(); + EXPECT_THROW(reader.Finish(), std::runtime_error); + } + ReplaceText("{}"); + EXPECT_THROW(ModelComparisonReader reader(path_), std::exception); +} + +TEST_F(ModelComparisonReference, RejectsInvalidTextFields) { + WriteReference(); + const auto original = ReadText(); + for (const auto& [from, to] : + std::vector>{ + {"\"version\":1", "\"version\":2"}, + {"\"sample_count\":1", "\"sample_count\":-1"}, + {"\"vocab_size\":3", "\"vocab_size\":3.5"}}) { + std::string changed = original; + changed.replace(changed.find(from), from.size(), to); + ReplaceText(changed); + EXPECT_THROW(ModelComparisonReader reader(path_), std::exception); + } + for (const auto& [from, to] : + std::vector>{ + {"\"id\":7", "\"id\":7.5"}, + {"\"expected\":2", "\"expected\":4294967298"}, + {"\"prompt\":[1,2]", "\"prompt\":[1.5,2]"}, + {"\"logits\":[1.25,-2.0,4.0]", "\"logits\":[null,-2.0,4.0]"}}) { + std::string changed = original; + changed.replace(changed.find(from), from.size(), to); + ReplaceText(changed); + ModelComparisonReader reader(path_); + EXPECT_THROW(reader.Read(), std::exception); + } +} + +TEST_F(ModelComparisonReference, RejectsInvalidRecords) { + ModelComparisonWriter writer(path_, metadata_); + auto bad = record_; + bad.logits.pop_back(); + EXPECT_THROW(writer.Write(bad), std::invalid_argument); + bad = record_; + bad.logits[0] = std::numeric_limits::quiet_NaN(); + EXPECT_THROW(writer.Write(bad), std::invalid_argument); + bad = record_; + bad.expected_label = 4; + EXPECT_THROW(writer.Write(bad), std::invalid_argument); + bad = record_; + bad.prompt = {-1}; + EXPECT_THROW(writer.Write(bad), std::invalid_argument); + bad.prompt = {3}; + EXPECT_THROW(writer.Write(bad), std::invalid_argument); + bad.prompt.clear(); + EXPECT_THROW(writer.Write(bad), std::invalid_argument); +} + +} // namespace +} // namespace gcpp diff --git a/evals/run_mmlu.cc b/evals/run_mmlu.cc index 66044397..4d2434cf 100644 --- a/evals/run_mmlu.cc +++ b/evals/run_mmlu.cc @@ -16,17 +16,27 @@ #include #include +#include +#include +#include +#include +#include +#include +#include +#include #include +#include #include #include "evals/benchmark_helper.h" +#include "evals/model_comparison.h" #include "gemma/gemma.h" // Gemma -#include "io/io.h" // Path -#include "util/args.h" #include "hwy/base.h" #include "hwy/highway.h" #include "hwy/profiler.h" +#include "io/io.h" // Path #include "nlohmann/json.hpp" +#include "util/args.h" namespace gcpp { @@ -36,116 +46,279 @@ struct JsonArgs : public ArgsBase { } Path input; + Path reference_out; + Path reference_in; + size_t max_questions; - // Returns error string or nullptr if OK. const char* Validate() const { if (input.Empty()) return "Must specify --input"; if (!input.Exists()) return "--input file does not exist"; + if (!reference_out.Empty() && !reference_in.Empty()) { + return "Specify only one of --reference_out and --reference_in"; + } + if (!reference_in.Empty() && !reference_in.Exists()) { + return "--reference_in file does not exist"; + } return nullptr; } template void ForEach(const Visitor& visitor) { visitor(input, "input", Path(), "Full pathname of mmlu.json."); - }; + visitor(reference_out, "reference_out", Path(), + "Write root-model full-vocabulary logits to this JSON Lines file."); + visitor(reference_in, "reference_in", Path(), + "Compare this target model against a root reference file."); + visitor(max_questions, "max_questions", size_t{0}, + "Maximum questions to run; zero runs the full dataset."); + } }; -// Linear search for a few tokens is faster than std::set. -// TODO: instead of accepting for each vocab entry, filter the logits once. -class TokenSet { +// Maps both "A" and " A" tokenizer variants to answer labels 0..3. +class AnswerTokens { public: - TokenSet(const GemmaTokenizer& tokenizer, - const std::vector& strings) { - all_tokens_.reserve(strings.size()); - for (const std::string& str : strings) { - std::vector tokens; - fprintf(stderr, "%s -> ", str.c_str()); - HWY_ASSERT(tokenizer.Encode(str, &tokens)); - for (int token : tokens) { - fprintf(stderr, "%d, ", token); - all_tokens_.push_back(token); + explicit AnswerTokens(const GemmaTokenizer& tokenizer) { + for (int label = 0; label < 4; ++label) { + for (const std::string& prefix : {std::string(), std::string(" ")}) { + const std::string str = prefix + static_cast('A' + label); + std::vector tokens; + if (!tokenizer.Encode(str, &tokens)) { + throw std::runtime_error("failed to tokenize answer label"); + } + if (tokens.size() != 1) continue; + fprintf(stderr, "%s -> %d\n", str.c_str(), tokens[0]); + if (tokens[0] < 0) throw std::runtime_error("invalid answer token"); + for (const auto& [token, previous_label] : tokens_) { + if (token == tokens[0] && previous_label != label) { + throw std::runtime_error("ambiguous answer token encoding"); + } + } + tokens_.push_back({tokens[0], label}); + } + if (tokens_.empty() || tokens_.back().second != label) { + throw std::runtime_error("answer label has no single-token encoding"); } - fprintf(stderr, "\n"); } } - bool Contains(int token) const { - return std::find(all_tokens_.begin(), all_tokens_.end(), token) != - all_tokens_.end(); - } + const std::vector>& All() const { return tokens_; } private: - std::vector all_tokens_; + std::vector> tokens_; }; -void Run(GemmaEnv& env, JsonArgs& json) { +double Percentile(const std::vector& sorted, double quantile) { + if (sorted.empty()) return 0.0; + const double position = quantile * static_cast(sorted.size() - 1); + const size_t lower = static_cast(std::floor(position)); + const size_t upper = static_cast(std::ceil(position)); + const double fraction = position - static_cast(lower); + return sorted[lower] + fraction * (sorted[upper] - sorted[lower]); +} + +void Run(GemmaEnv& env, JsonArgs& args) { + using Clock = std::chrono::steady_clock; + double generate_seconds = 0.0, sample_seconds = 0.0; PROFILER_ZONE("Run.all"); - float answers = 0.0f; - float correct_answers = 0.0f; + size_t answers = 0; + size_t correct_answers = 0; + std::vector kl_values; - auto json_data = nlohmann::json::parse(ReadFileToString(json.input)); + const std::string json_text = ReadFileToString(args.input); + const auto json_data = nlohmann::json::parse(json_text); + const auto& samples = json_data.at("samples"); + if (!samples.is_array() || samples.empty()) { + throw std::invalid_argument("samples must be a nonempty array"); + } + std::set ids; + for (const auto& sample : samples) { + if (!sample.at("i").is_number_integer() || + sample.at("i") > std::numeric_limits::max() || + !sample.at("input_label").is_number_integer() || + sample.at("input_label") < 0 || sample.at("input_label") > 3 || + !sample.at("prompt").is_string() || + !ids.insert(sample.at("i").get()).second) { + throw std::invalid_argument("invalid or duplicate MMLU sample"); + } + } + const size_t sample_count = + args.max_questions == 0 + ? samples.size() + : std::min(args.max_questions, samples.size()); + + const Gemma& gemma = *env.GetGemma(); + const GemmaTokenizer& tokenizer = gemma.Tokenizer(); + const ModelComparisonMetadata metadata = { + static_cast(gemma.Config().vocab_size), sample_count, json_text, + ModelComparisonHex(tokenizer.Serialize())}; + + std::unique_ptr reference_writer; + std::unique_ptr reference_reader; + if (!args.reference_out.Empty()) { + reference_writer = std::make_unique( + args.reference_out.path, metadata); + } else if (!args.reference_in.Empty()) { + reference_reader = + std::make_unique(args.reference_in.path); + reference_reader->Validate(metadata); + kl_values.reserve(sample_count); + } + + const AnswerTokens answer_tokens(tokenizer); + + for (const auto& sample : samples) { + if (answers >= sample_count) break; + const int64_t id = sample["i"]; + fprintf(stderr, "Processing question %lld\n", static_cast(id)); + const int correct_label = sample["input_label"]; + const std::string correct_answer(1, static_cast('A' + correct_label)); - const std::vector accept_strings = { - "A", "B", "C", "D", // - " A", " B", " C", " D", // - "**", "**:", ":**", "The", "Answer", "is", ":", "."}; - const TokenSet accept_set(env.GetGemma()->Tokenizer(), accept_strings); + ModelComparisonRecord root_record; + if (reference_reader) { + root_record = reference_reader->Read(); + if (root_record.sample_id != id || + root_record.expected_label != correct_label) { + throw std::runtime_error("KL reference sample identity mismatch"); + } + } - for (auto sample : json_data["samples"]) { - const int id = sample["i"]; - fprintf(stderr, "Processing question %d\n", id); - const std::string& correct_answer = accept_strings[sample["input_label"]]; std::string prompt_string = sample["prompt"]; - // AcceptFunc restricts the output to one of these four tokens, so make an - // effort to steer the model towards that. See - // https://huggingface.co/blog/open-llm-leaderboard-mmlu prompt_string += "What is start of the line with the correct answer? " "Do not include any justifications or explanations. Reply only with a " "letter."; const std::vector prompt = env.WrapAndTokenize(prompt_string); - const size_t prompt_size = prompt.size(); - - std::vector predicted_token_ids; - predicted_token_ids.reserve(4096); - size_t generated = 0; - const StreamFunc stream_token = [&generated, prompt_size, - &predicted_token_ids](int token, - float proba) { - PROFILER_ZONE("Stream"); - ++generated; - if (generated > prompt_size) { - predicted_token_ids.push_back(token); - } - return true; - }; + if (reference_reader && root_record.prompt != prompt) { + throw std::runtime_error("KL reference tokenized prompt mismatch"); + } - // Although " A" is a token, it is difficult to associate that with the - // correct answer. Only accepting certain tokens is risky: (A) is easily - // confused with the word "A". + MmluAnswer answer; + std::vector captured_logits; + double full_vocab_kl = 0.0; + bool sampled = false; + // Exceptions must not escape a callback running on a worker thread. + std::exception_ptr sample_error; gcpp::TimingInfo timing_info; - gcpp::RuntimeConfig runtime_config = { - .max_generated_tokens = 30, - .temperature = 0.0f, - .verbosity = env.Verbosity(), - .attention_impl = env.MutableConfig().attention_impl, - .stream_token = stream_token, + RuntimeConfig runtime_config = env.MutableConfig(); + runtime_config.max_generated_tokens = 1; + runtime_config.temperature = 0.0f; + runtime_config.stream_token = [](int, float) { return true; }; + runtime_config.sample_func = + [&answer_tokens, &answer, &captured_logits, &full_vocab_kl, + &reference_writer, &reference_reader, &root_record, &sample_seconds, + &sampled, + &sample_error](size_t /*query_idx*/, size_t /*pos*/, Logits logits, + size_t /*worker*/) -> TokenAndProb { + const auto sample_start = Clock::now(); + try { + if (sampled) throw std::runtime_error("multiple answer samples"); + sampled = true; + if (reference_writer) { + captured_logits.assign(logits.data(), logits.data() + logits.size()); + } else if (reference_reader) { + full_vocab_kl = FullVocabKLDivergence(root_record.logits, + logits.data(), logits.size()); + } + + answer = + ScoreMmluAnswer(logits.data(), logits.size(), answer_tokens.All()); + sample_seconds += + std::chrono::duration(Clock::now() - sample_start).count(); + return TokenAndProb{ + .token = answer.token, + .prob = static_cast(answer.probabilities[answer.label])}; + } catch (...) { + sample_error = std::current_exception(); + return TokenAndProb{.token = 0, .prob = 0.0f}; + } }; + const auto generate_start = Clock::now(); env.GetGemma()->Generate(runtime_config, prompt, /*pos=*/0, env.MutableKVCache(), env.MutableEnv(), timing_info); - std::string output_string = env.StringFromTokens(predicted_token_ids); + generate_seconds += + std::chrono::duration(Clock::now() - generate_start).count(); + + if (sample_error) std::rethrow_exception(sample_error); + if (!sampled) throw std::runtime_error("failed to sample answer token"); + if (reference_writer) { + if (captured_logits.size() != metadata.vocab_size) { + throw std::runtime_error( + "failed to capture root full-vocabulary logits"); + } + reference_writer->Write( + {id, correct_label, prompt, std::move(captured_logits)}); + } else if (reference_reader) { + if (!std::isfinite(full_vocab_kl) || full_vocab_kl < 0.0) { + throw std::runtime_error("invalid full-vocabulary KL divergence"); + } + kl_values.push_back(full_vocab_kl); + } + + const int predicted_label = answer.label; + const std::string output_string = + predicted_label == -1 + ? std::string("?") + : std::string(1, static_cast('A' + predicted_label)); fprintf(stderr, "Correct %s, model '%s'\n", correct_answer.c_str(), output_string.c_str()); - answers += 1.0f; - if (output_string == correct_answer) { - correct_answers += 1.0f; - } - fprintf(stderr, "%.0f/%.0f = %.2f%%\n", correct_answers, answers, - correct_answers / answers); + const bool is_correct = predicted_label == correct_label; + ++answers; + correct_answers += static_cast(is_correct); + nlohmann::json result = { + {"id", id}, + {"expected", correct_answer}, + {"predicted", output_string}, + {"correct", is_correct}, + {"logits", answer.logits}, + {"probabilities", answer.probabilities}, + {"margin", answer.margin}, + }; + if (reference_reader) result["full_vocab_kl"] = full_vocab_kl; + printf("MMLU_RESULT %s\n", result.dump().c_str()); + fflush(stdout); + fprintf(stderr, "%zu/%zu = %.2f%%\n", correct_answers, answers, + 100.0 * static_cast(correct_answers) / answers); + } + + if (reference_writer) reference_writer->Finish(); + if (reference_reader) reference_reader->Finish(); + + const nlohmann::json timing = { + {"generate_seconds", generate_seconds}, + {"sample_seconds", sample_seconds}, + {"inference_seconds", generate_seconds - sample_seconds}, + {"scope", + "Generate excluding evaluation sample callback; includes MatMul " + "autotuning"}, + }; + printf("MMLU_TIMING %s\n", timing.dump().c_str()); + const nlohmann::json summary = { + {"answers", answers}, + {"correct", correct_answers}, + {"accuracy", + answers == 0 ? 0.0 : static_cast(correct_answers) / answers}, + }; + printf("MMLU_SUMMARY %s\n", summary.dump().c_str()); + + if (!kl_values.empty()) { + std::sort(kl_values.begin(), kl_values.end()); + const double mean = + std::accumulate(kl_values.begin(), kl_values.end(), 0.0) / + kl_values.size(); + const nlohmann::json kl_summary = { + {"samples", kl_values.size()}, + {"mean", mean}, + {"median", Percentile(kl_values, 0.5)}, + {"p95", Percentile(kl_values, 0.95)}, + {"max", kl_values.back()}, + {"unit", "nats"}, + {"direction", "root||target"}, + }; + printf("MMLU_KL_SUMMARY %s\n", kl_summary.dump().c_str()); } } @@ -154,17 +327,22 @@ void Run(GemmaEnv& env, JsonArgs& json) { int main(int argc, char** argv) { gcpp::InternalInit(); - { - PROFILER_ZONE("Startup.all"); - gcpp::ConsumedArgs consumed(argc, argv); - gcpp::GemmaArgs args(argc, argv, consumed); - gcpp::JsonArgs json_args(argc, argv, consumed); - gcpp::AbortIfInvalidArgs(json_args); - consumed.AbortIfUnconsumed(); + try { + { + PROFILER_ZONE("Startup.all"); + gcpp::ConsumedArgs consumed(argc, argv); + gcpp::GemmaArgs args(argc, argv, consumed); + gcpp::JsonArgs json_args(argc, argv, consumed); + gcpp::AbortIfInvalidArgs(json_args); + consumed.AbortIfUnconsumed(); - gcpp::GemmaEnv env(args); - gcpp::Run(env, json_args); + gcpp::GemmaEnv env(args); + gcpp::Run(env, json_args); + } + PROFILER_PRINT_RESULTS(); + return 0; + } catch (const std::exception& error) { + fprintf(stderr, "model comparison failed: %s\n", error.what()); + return 1; } - PROFILER_PRINT_RESULTS(); // Must call outside the zone above. - return 0; } diff --git a/examples/hello_world/run.cc b/examples/hello_world/run.cc index 9cd8b0e3..17e37254 100644 --- a/examples/hello_world/run.cc +++ b/examples/hello_world/run.cc @@ -50,7 +50,7 @@ int main(int argc, char** argv) { // Instantiate model and KV Cache gcpp::ThreadingContext ctx(args.threading); - gcpp::MatMulEnv env(ctx); + gcpp::MatMulEnv env(ctx, args.inference.MatMulSchedule()); gcpp::Gemma gemma(args, ctx); gcpp::KVCache kv_cache(gemma.Config(), args.inference, ctx.allocator); size_t generated = 0; diff --git a/gemma/api_server.cc b/gemma/api_server.cc index c713635f..f2ddfc60 100644 --- a/gemma/api_server.cc +++ b/gemma/api_server.cc @@ -427,7 +427,7 @@ void RunServer(const GemmaArgs& args) { // Initialize model ThreadingContext ctx(args.threading); - MatMulEnv env(ctx); + MatMulEnv env(ctx, args.inference.MatMulSchedule()); ServerState state; state.gemma = std::make_unique(args, ctx); state.env = &env; diff --git a/gemma/gemma_args.h b/gemma/gemma_args.h index b0746fde..db348ece 100644 --- a/gemma/gemma_args.h +++ b/gemma/gemma_args.h @@ -28,6 +28,7 @@ #include "compression/types.h" #include "gemma/configs.h" #include "io/io.h" // Path +#include "ops/matmul.h" // MMSchedule #include "util/args.h" // IWYU pragma: export #include "util/basics.h" // Tristate #include "util/mat.h" @@ -216,6 +217,7 @@ struct InferenceArgs : public ArgsBase { float temperature; size_t top_k; bool deterministic; + std::string matmul_schedule; bool multiturn; Path image_file; @@ -254,6 +256,10 @@ struct InferenceArgs : public ArgsBase { 2); visitor(deterministic, "deterministic", false, "Make top-k sampling deterministic", 2); + visitor( + matmul_schedule, "matmul_schedule", std::string("auto"), + "MatMul schedule: auto, fixed, or fixed_min_k (fewest K partitions).", + 2); visitor(multiturn, "multiturn", false, "Multiturn mode\n 0 = clear KV cache after every " "interaction\n 1 = continue KV cache after every interaction\n " @@ -302,6 +308,16 @@ struct InferenceArgs : public ArgsBase { 2); } + MMSchedule MatMulSchedule() const { + if (matmul_schedule == "auto") return MMSchedule::kAutoTune; + if (matmul_schedule == "fixed") return MMSchedule::kFixed; + if (matmul_schedule == "fixed_min_k") return MMSchedule::kFixedMinK; + HWY_ABORT( + "Invalid --matmul_schedule '%s': expected auto, fixed, or " + "fixed_min_k", + matmul_schedule.c_str()); + } + void CopyTo(RuntimeConfig& runtime_config) const { runtime_config.max_generated_tokens = max_generated_tokens; runtime_config.prefill_tbatch_size = prefill_tbatch_size; diff --git a/gemma/gemma_args_test.cc b/gemma/gemma_args_test.cc index d9ee8b71..d18e09bb 100644 --- a/gemma/gemma_args_test.cc +++ b/gemma/gemma_args_test.cc @@ -3,6 +3,7 @@ #include #include +#include #include #include "gtest/gtest.h" @@ -52,6 +53,31 @@ TEST(GemmaArgsTest, AllConsumedArgs) { CheckAllConsumed({"gemma", "--verbosity", "2", "--deterministic=true"}); } +TEST(GemmaArgsTest, MatMulSchedule) { + EXPECT_EQ(InferenceArgs().MatMulSchedule(), MMSchedule::kAutoTune); + for (const auto& entry : std::vector>{ + {"auto", MMSchedule::kAutoTune}, + {"fixed", MMSchedule::kFixed}, + {"fixed_min_k", MMSchedule::kFixedMinK}}) { + const std::vector args = {"gemma", "--matmul_schedule", + entry.first}; + std::vector ptrs; + FillPtrs(args, ptrs); + ConsumedArgs consumed(ptrs.size(), ptrs.data()); + GemmaArgs parsed(ptrs.size(), ptrs.data(), consumed); + consumed.AbortIfUnconsumed(); + EXPECT_EQ(parsed.inference.MatMulSchedule(), entry.second); + EXPECT_FALSE(parsed.inference.deterministic); + CheckAllConsumed({"gemma", "--matmul_schedule=" + entry.first}); + } +} + +TEST(GemmaArgsTest, InvalidMatMulSchedule) { + InferenceArgs args; + args.matmul_schedule = "fixd"; + EXPECT_DEATH(args.MatMulSchedule(), "Invalid --matmul_schedule"); +} + TEST(GemmaArgsTest, UnconsumedArgs) { // Single unconsumed arg CheckUnconsumed({"gemma", "--UNDEFINED"}, 1); diff --git a/gemma/run.cc b/gemma/run.cc index 65f714b5..bc100f09 100644 --- a/gemma/run.cc +++ b/gemma/run.cc @@ -260,7 +260,7 @@ void Run(const GemmaArgs& args) { PROFILER_ZONE("Run.misc"); ThreadingContext ctx(args.threading); - MatMulEnv env(ctx); + MatMulEnv env(ctx, args.inference.MatMulSchedule()); const InferenceArgs& inference = args.inference; if (inference.verbosity >= 3) env.print_best = true; const Gemma gemma(args, ctx); diff --git a/ops/brgemm.h b/ops/brgemm.h index f2c6eea4..5f66da16 100644 --- a/ops/brgemm.h +++ b/ops/brgemm.h @@ -56,8 +56,9 @@ struct BRGeMMConfig { // Generates autotuning candidates. Fixed: N_blk=32, K_blk=32 (AMX BF16). // Tunable: M_blk in {32,64}, batch_size in {16,32,64,128,256}. -inline std::vector BRGeMMCandidates(size_t M, size_t K, - size_t N) { +inline std::vector BRGeMMCandidates(size_t M, size_t K, size_t N, + bool autotune = true, + bool min_k_splits = false) { std::vector out; out.reserve(10); // At most 2 M_blk * 5 batch_size candidates. static constexpr size_t kNBlk = 32; @@ -85,6 +86,18 @@ inline std::vector BRGeMMCandidates(size_t M, size_t K, if (out.empty()) { out.push_back({std::min(M, size_t{32}), std::min(N, size_t{32}), 32, 1, 1}); } + if (!autotune) { + size_t selected = 0; + if (min_k_splits) { + for (size_t i = 1; i < out.size(); ++i) { + const size_t splits = hwy::DivCeil(k_chunks, out[i].batch_size); + const size_t best_splits = + hwy::DivCeil(k_chunks, out[selected].batch_size); + if (splits < best_splits) selected = i; + } + } + return {out[selected]}; + } return out; } diff --git a/ops/matmul-inl.h b/ops/matmul-inl.h index cabe5392..5673926f 100644 --- a/ops/matmul-inl.h +++ b/ops/matmul-inl.h @@ -372,10 +372,7 @@ class MMDecompress { // First call: generate candidates. if (HWY_UNLIKELY(!autotune.HasCandidates())) { - const MMParA other = (A.Rows() == 1) ? MMParA::kNone : MMParA::kM; - std::vector candidates = {MMParA::kK1, MMParA::kK2, MMParA::kK4, - other}; - autotune.SetCandidates(candidates); + autotune.SetCandidates(MMParACandidates(A.Rows(), env.schedule)); } const MMParA& par_a = autotune.NextConfig(); @@ -1150,10 +1147,11 @@ class MMImpl { } public: - static MMPerKey& FindOrAddPerKey(size_t M, size_t K, size_t N, size_t num_B, - size_t vector_bytes, - MatMulEnv::PerCluster& per_cluster) { - const MMKeys::Key key = MMKeys::KeyFromDims(M, K, N, num_B); + static MMPerKey& FindOrAddPerKey( + size_t M, size_t K, size_t N, size_t num_B, size_t vector_bytes, + MatMulEnv::PerCluster& per_cluster, + MMActivation activation = MMActivation::kBF16) { + const MMKeys::Key key = MMKeys::KeyFromDims(M, K, N, num_B, activation); intptr_t index = IndexOfKey(key, per_cluster.keys); // First time we see this shape/key. if (HWY_UNLIKELY(index < 0)) { @@ -1498,26 +1496,33 @@ HWY_NOINLINE MMPerKey* MatMul(const MatPtrT& A, const MatPtrT& B, // BRGeMM failed; fall through to standard matmul. } - if (HWY_UNLIKELY(!brg_tuner.HasCandidates())) { - brg_tuner.SetCandidates(BRGeMMCandidates(M, K, N)); - } + // A cached BRGeMM candidate may fail and fall back to the generic + // kernel. Do not re-enter the tuner after it has chosen a winner. + if (!brg_tuner.Best()) { + if (HWY_UNLIKELY(!brg_tuner.HasCandidates())) { + brg_tuner.SetCandidates( + BRGeMMCandidates(M, K, N, env.schedule == MMSchedule::kAutoTune, + env.schedule == MMSchedule::kFixedMinK)); + } - const BRGeMMConfig& cfg = brg_tuner.NextConfig(); - const uint64_t t0 = hwy::timer::Start(); - if (DoMatMul_BRGeMM(A, B, C_rows, M, K, N, scale, add, cfg, env.ctx, - cluster_idx)) { - const uint64_t t1 = - env.have_timer_stop ? hwy::timer::Stop() : hwy::timer::Start(); - brg_tuner.NotifyTicks(t1 - t0); - - if (HWY_UNLIKELY(env.print_best && brg_tuner.Best())) { - const BRGeMMConfig& best = *brg_tuner.Best(); - fprintf(stderr, - "BRGeMM best: %zux%zux%zu M_blk=%zu N_blk=%zu K_blk=%zu " - "batch=%zu\n", - M, K, N, best.M_blk, best.N_blk, best.K_blk, best.batch_size); + const BRGeMMConfig& cfg = brg_tuner.NextConfig(); + const uint64_t t0 = hwy::timer::Start(); + if (DoMatMul_BRGeMM(A, B, C_rows, M, K, N, scale, add, cfg, env.ctx, + cluster_idx)) { + const uint64_t t1 = + env.have_timer_stop ? hwy::timer::Stop() : hwy::timer::Start(); + brg_tuner.NotifyTicks(t1 - t0); + + if (HWY_UNLIKELY(env.print_best && brg_tuner.Best())) { + const BRGeMMConfig& best = *brg_tuner.Best(); + fprintf(stderr, + "BRGeMM best: %zux%zux%zu M_blk=%zu N_blk=%zu K_blk=%zu " + "batch=%zu\n", + M, K, N, best.M_blk, best.N_blk, best.K_blk, + best.batch_size); + } + return &per_key; } - return &per_key; } // BRGeMM failed; fall through to standard matmul. } @@ -1562,8 +1567,12 @@ HWY_NOINLINE MMPerKey* MatMul(const MatPtrT& A, const MatPtrT& B, HWY_ASSERT(K <= MMEntireA::kMaxK); HWY_ASSERT(N % kNR == 0); MMImpl::EnsureAligned(A, cache.VectorBytes()); - tuner.SetCandidates( - MMCandidates(cache, M, K, N, num_B, sizeof(TC), env.print_config)); + // Fixed scheduling must not depend on which M in this key's bucket + // happened to be seen first. + const size_t max_M = + env.schedule == MMSchedule::kAutoTune ? M : MMKeys::BucketM(M); + tuner.SetCandidates(MMCandidates(cache, max_M, K, N, num_B, sizeof(TC), + env.print_config, env.schedule)); } const MMConfig& cfg = tuner.NextConfig(); @@ -1625,7 +1634,7 @@ HWY_NOINLINE MMPerKey* TwoMatMul(const MatPtrT& A, const MatPtrT& B1, MMImpl::EnsureAligned(A, cache.VectorBytes()); const size_t max_M = MMKeys::BucketM(M); tuner.SetCandidates(MMCandidates(cache, max_M, K, N, num_B, sizeof(BF16), - env.print_config)); + env.print_config, env.schedule)); } const MMConfig& cfg = tuner.NextConfig(); diff --git a/ops/matmul.cc b/ops/matmul.cc index 2167dc80..ba49acda 100644 --- a/ops/matmul.cc +++ b/ops/matmul.cc @@ -392,13 +392,32 @@ class GenerateCandidates { // Facade to avoid exposing `GenerateCandidates` in the header. std::vector MMCandidates(const CacheInfo& cache, size_t max_M, size_t K, size_t N, size_t num_B, - size_t sizeof_TC, bool print_config) { - return GenerateCandidates(cache, max_M, K, N, num_B, sizeof_TC, - print_config)(); + size_t sizeof_TC, bool print_config, + MMSchedule schedule) { + auto candidates = + GenerateCandidates(cache, max_M, K, N, num_B, sizeof_TC, print_config)(); + if (schedule == MMSchedule::kAutoTune) return candidates; + + size_t selected = 0; + if (schedule == MMSchedule::kFixedMinK) { + for (size_t i = 1; i < candidates.size(); ++i) { + if (candidates[i].RangesOfKC(K).NumTasks() < + candidates[selected].RangesOfKC(K).NumTasks()) { + selected = i; + } + } + } + return {candidates[selected]}; +} + +std::vector MMParACandidates(size_t M, MMSchedule schedule) { + if (schedule != MMSchedule::kAutoTune) return {MMParA::kK1}; + const MMParA other = (M == 1) ? MMParA::kNone : MMParA::kM; + return {MMParA::kK1, MMParA::kK2, MMParA::kK4, other}; } -MatMulEnv::MatMulEnv(ThreadingContext& ctx) - : ctx(ctx), A_BF(ctx.allocator), C_tiles(ctx) { +MatMulEnv::MatMulEnv(ThreadingContext& ctx, MMSchedule schedule) + : ctx(ctx), schedule(schedule), A_BF(ctx.allocator), C_tiles(ctx) { const size_t num_clusters = ctx.pools.NumClusters(); per_cluster.resize(num_clusters); for (size_t cluster_idx = 0; cluster_idx < num_clusters; ++cluster_idx) { diff --git a/ops/matmul.h b/ops/matmul.h index af0a45a6..f235bb4f 100644 --- a/ops/matmul.h +++ b/ops/matmul.h @@ -573,9 +573,17 @@ class MMConfig { static_assert(sizeof(MMConfig) == 32); // for faster indexing #pragma pack(pop) +// Fixed modes select one legal candidate independently of timing measurements. +// kFixedMinK prefers fewer intermediate stores/roundings along K, retaining +// candidate order to break ties. Neither fixed mode promises optimal speed. +enum class MMSchedule : uint8_t { kAutoTune, kFixed, kFixedMinK }; + +std::vector MMParACandidates(size_t M, MMSchedule schedule); + std::vector MMCandidates(const CacheInfo& cache, size_t M, size_t K, size_t N, size_t num_B, size_t sizeof_TC, - bool print_config); + bool print_config, + MMSchedule schedule = MMSchedule::kAutoTune); // State machine for choosing the best `TConfig`, which is `MMConfig` for the // main MatMul autotuner. @@ -637,8 +645,8 @@ class MMAutoTune { skipped_.Set(my_idx); } - // After sufficient rounds, choose the winner. - if (rounds_complete_ == 4) { + // A single candidate needs no comparison with other configurations. + if (candidates_.size() == 1 || rounds_complete_ == 4) { for (size_t i = 0; i < candidates_.size(); ++i) { worst_min_ticks_ = HWY_MAX(worst_min_ticks_, min_ticks_[i]); if (min_ticks_[i] == best_ticks_) { @@ -677,6 +685,10 @@ class MMAutoTune { //------------------------------------------------------------------------------ +// The representation consumed by the kernel (F32 inputs converted to BF16 +// also use kBF16). Reserve separate tuning state for the forthcoming A8 paths. +enum class MMActivation : uint8_t { kBF16, kI8, kI8Block }; + // Map of previously seen dimensions to index via linear search. class MMKeys { public: @@ -699,12 +711,19 @@ class MMKeys { } // Compresses the dimensions into a single Key for faster comparison. - static Key KeyFromDims(size_t M, size_t K, size_t N, size_t num_B) { + static Key KeyFromDims(size_t M, size_t K, size_t N, size_t num_B, + MMActivation activation = MMActivation::kBF16) { HWY_DASSERT(M < (Key{1} << 16)); // batch sizes are smaller HWY_DASSERT(K < (Key{1} << 20)); HWY_DASSERT(N < (Key{1} << 20)); HWY_DASSERT(num_B == 1 || num_B == 2); + HWY_DASSERT(activation == MMActivation::kBF16 || + activation == MMActivation::kI8 || + activation == MMActivation::kI8Block); + // M: bits 0..15, K: 16..35, activation: 36..37, N: 40..59, + // num_B: 60..61. BF16 retains the existing key encoding. const Key key = static_cast(BucketM(M)) | (static_cast(K) << 16) | + (static_cast(activation) << 36) | (static_cast(N) << 40) | (static_cast(num_B) << 60); HWY_DASSERT(key != kPadding); @@ -733,7 +752,9 @@ class MMKeys { HWY_DASSERT(capacity_ >= num_unique_ + 1); hwy::AlignedFreeUniquePtr new_keys = hwy::AllocateAligned(capacity_); - hwy::CopyBytes(keys_.get(), new_keys.get(), num_unique_ * sizeof(Key)); + if (num_unique_ != 0) { + hwy::CopyBytes(keys_.get(), new_keys.get(), num_unique_ * sizeof(Key)); + } // Pad for SIMD. for (size_t i = num_unique_; i < hwy::RoundUpTo(num_unique_, NU64); ++i) { new_keys[i] = kPadding; @@ -766,9 +787,12 @@ struct MMPerKey { // Stores state shared across MatMul calls. Non-copyable. `ctx` must outlive // `MatMulEnv`. struct MatMulEnv { - explicit MatMulEnv(ThreadingContext& ctx); + explicit MatMulEnv(ThreadingContext& ctx, + MMSchedule schedule = MMSchedule::kAutoTune); ThreadingContext& ctx; + // Immutable because cached candidates depend on the scheduling policy. + const MMSchedule schedule; bool have_timer_stop = false; // Whether `MMCandidates()` should print the set of parameters. diff --git a/ops/matmul_test.cc b/ops/matmul_test.cc index 88988333..598913a3 100644 --- a/ops/matmul_test.cc +++ b/ops/matmul_test.cc @@ -30,6 +30,9 @@ #include #include +#include +#include +#include #include "ops/matmul.h" #include "util/basics.h" @@ -131,11 +134,36 @@ void TestMatMul(size_t rows_ac, size_t cols_a_rows_b, size_t cols_bc, bool add, MatMulSlow(A, BT, add_row, env, C_slow); // A few reps to get coverage of the various autotuned code paths. MMOptions options; + std::vector first_output; + std::vector first_fused_output; for (size_t rep = 0; rep < 16; ++rep) { MMPerKey* per_key = MatMulStatic(A, BT, add_row, env, C, options); AssertClose(A, BT, C_slow, C, env.ctx.allocator, env.row_ptrs, line); // Check before TwoMatMulStatic(), which can invalidate per_key. const bool autotune_done = !!per_key->autotune.Best(); + if (env.schedule != MMSchedule::kAutoTune) { + bool selected = autotune_done; +#if GEMMA_ONEDNN_BRGEMM + selected |= per_key->brgemm_autotune.Best() != nullptr; +#endif +#if GEMMA_ONEDNN_MATMUL + selected |= per_key->onednn_built; +#endif + HWY_ASSERT(selected); + if constexpr (!IsBF16()) { + HWY_ASSERT(per_key->autotune_par_a.Best() != nullptr); + } + if (rep == 0) first_output.resize(C.Extents().Area()); + for (size_t row = 0; row < C.Rows(); ++row) { + TC* first = first_output.data() + row * C.Cols(); + if (rep == 0) { + hwy::CopyBytes(C.Row(row), first, C.Cols() * sizeof(TC)); + } else { + HWY_ASSERT(std::memcmp(first, C.Row(row), C.Cols() * sizeof(TC)) == + 0); + } + } + } // Ensure the tiled view returns the same result as C. if constexpr (IsBF16() && IsBF16()) { @@ -167,6 +195,18 @@ void TestMatMul(size_t rows_ac, size_t cols_a_rows_b, size_t cols_bc, bool add, options.SetFunc(fused); TwoMatMulStatic(A, BT, BT, env, C2, options); HWY_ASSERT_EQ(C.Extents().Area(), total_view_area.load()); + if (env.schedule != MMSchedule::kAutoTune) { + if (rep == 0) first_fused_output.resize(C2.Extents().Area()); + for (size_t row = 0; row < C2.Rows(); ++row) { + TC* first = first_fused_output.data() + row * C2.Cols(); + if (rep == 0) { + hwy::CopyBytes(C2.Row(row), first, C2.Cols() * sizeof(TC)); + } else { + HWY_ASSERT(std::memcmp(first, C2.Row(row), + C2.Cols() * sizeof(TC)) == 0); + } + } + } options.func = nullptr; // reset for next call // TwoMatMulStatic() does not support adding a bias vector. @@ -175,13 +215,115 @@ void TestMatMul(size_t rows_ac, size_t cols_a_rows_b, size_t cols_bc, bool add, } } - if (autotune_done) break; + if (autotune_done && env.schedule == MMSchedule::kAutoTune) break; } } using F32 = float; using SFP = SfpStream; +// Fixed schedules exercise both activation conversion and the fused path. +// TestMatMul compares against the scalar reference and checks repeated results +// bit-for-bit, including the transition from the first call to the cached path. +void TestFixedSchedules() { + ThreadingArgs threading_args; + threading_args.max_threads = 2; + threading_args.bind = Tristate::kFalse; + ThreadingContext ctx(threading_args); + for (MMSchedule schedule : {MMSchedule::kFixed, MMSchedule::kFixedMinK}) { + MatMulEnv env(ctx, schedule); + for (size_t M : {size_t{1}, size_t{3}, size_t{4}, size_t{17}, size_t{65}}) { + TestMatMul(M, 258, 32, true, env, __LINE__); + TestMatMul(M, 256, 32, false, env, __LINE__); + TestMatMul(M, 4096, 32, false, env, __LINE__); + } + // K above the one-block limit must still yield a legal fixed schedule. + TestMatMul(3, 2 * kMaxKC, 32, true, env, __LINE__); + } +} + +void AssertSameConfig(const MMConfig& a, const MMConfig& b) { + HWY_ASSERT_EQ(a.MR(), b.MR()); + HWY_ASSERT_EQ(a.MC(), b.MC()); + HWY_ASSERT_EQ(a.KC(), b.KC()); + HWY_ASSERT_EQ(a.NC(), b.NC()); + HWY_ASSERT_EQ(static_cast(a.Order()), static_cast(b.Order())); + HWY_ASSERT_EQ(a.InnerTasks(), b.InnerTasks()); +} + +// The first M encountered in a shared bucket must not affect a fixed config. +void TestFixedBucketOrder() { + ThreadingArgs threading_args; + threading_args.max_threads = 2; + threading_args.bind = Tristate::kFalse; + ThreadingContext ctx(threading_args); + auto A = GenerateMat(Extents2D(7, 4096), MatPadding::kOdd, ctx); + auto B = GenerateTransposedMat(Extents2D(32, 4096), + MatPadding::kPacked, ctx); + MatStorageT C1("C1", Extents2D(7, 32), ctx.allocator, MatPadding::kOdd); + MatStorageT C2("C2", Extents2D(7, 32), ctx.allocator, MatPadding::kOdd); + for (MMSchedule schedule : {MMSchedule::kFixed, MMSchedule::kFixedMinK}) { + MatMulEnv small_first(ctx, schedule); + MatMulEnv large_first(ctx, schedule); + A.OverrideRows(4); + C1.OverrideRows(4); + const MMConfig small_config = *MatMulStatic( + A, B, nullptr, small_first, C1, MMOptions())->autotune.Best(); + A.OverrideRows(7); + C2.OverrideRows(7); + const MMConfig large_config = *MatMulStatic( + A, B, nullptr, large_first, C2, MMOptions())->autotune.Best(); + AssertSameConfig(small_config, large_config); + C1.OverrideRows(7); + MatMulStatic(A, B, nullptr, small_first, C1, MMOptions()); + for (size_t row = 0; row < C1.Rows(); ++row) { + HWY_ASSERT(std::memcmp(C1.Row(row), C2.Row(row), + C1.Cols() * sizeof(BF16)) == 0); + } + } +} + +void TestScheduleCandidates() { + ThreadingArgs threading_args; + threading_args.max_threads = 2; + threading_args.bind = Tristate::kFalse; + ThreadingContext ctx(threading_args); + for (size_t M : {size_t{1}, size_t{3}, size_t{17}, size_t{127}}) { + for (size_t K : {size_t{1}, size_t{258}, size_t{4096}, 2 * kMaxKC}) { + for (size_t num_B : {size_t{1}, size_t{2}}) { + for (size_t sizeof_TC : {sizeof(BF16), sizeof(float)}) { + const auto all = + MMCandidates(ctx.cache_info, M, K, 128, num_B, sizeof_TC, false); + for (MMSchedule schedule : + {MMSchedule::kFixed, MMSchedule::kFixedMinK}) { + const auto fixed = MMCandidates(ctx.cache_info, M, K, 128, num_B, + sizeof_TC, false, schedule); + HWY_ASSERT_EQ(size_t{1}, fixed.size()); + const auto again = MMCandidates(ctx.cache_info, M, K, 128, num_B, + sizeof_TC, false, schedule); + AssertSameConfig(fixed.front(), again.front()); + if (schedule == MMSchedule::kFixed) { + AssertSameConfig(all.front(), fixed.front()); + } else { + const size_t splits = fixed.front().RangesOfKC(K).NumTasks(); + for (const MMConfig& candidate : all) { + HWY_ASSERT(splits <= candidate.RangesOfKC(K).NumTasks()); + } + // First candidate with this split count wins ties. + for (const MMConfig& candidate : all) { + if (candidate.RangesOfKC(K).NumTasks() == splits) { + AssertSameConfig(candidate, fixed.front()); + break; + } + } + } + } + } + } + } + } +} + // Sweep all dimensions for a single input type and Highway target, to verify // the remainder handling. void TestTiny() { @@ -339,9 +481,96 @@ namespace gcpp { int64_t first_target = 0; // none run yet HWY_BEFORE_TEST(MatMulTest); HWY_EXPORT_AND_TEST_P(MatMulTest, TestTiny); +HWY_EXPORT_AND_TEST_P(MatMulTest, TestFixedSchedules); +HWY_EXPORT_AND_TEST_P(MatMulTest, TestFixedBucketOrder); +HWY_EXPORT_AND_TEST_P(MatMulTest, TestScheduleCandidates); HWY_EXPORT_AND_TEST_P(MatMulTest, TestAllMatMul); HWY_AFTER_TEST(); +TEST(MatMulSchedulingTest, SingleCandidate) { + MMAutoTune tuner; + tuner.SetCandidates({42}); + ASSERT_EQ(tuner.NextConfig(), 42); + tuner.NotifyTicks(123); + ASSERT_NE(tuner.Best(), nullptr); + EXPECT_EQ(*tuner.Best(), 42); + EXPECT_EQ(tuner.BestTicks(), 123); + EXPECT_EQ(tuner.FirstConfigTicks(), 123); + EXPECT_EQ(tuner.WorstMinTicks(), 123); +} + +TEST(MatMulSchedulingTest, MultipleCandidatesStillTune) { + MMAutoTune tuner; + tuner.SetCandidates({0, 1}); + for (size_t round = 0; round < 4; ++round) { + ASSERT_EQ(tuner.Best(), nullptr); + ASSERT_EQ(tuner.NextConfig(), 0); + tuner.NotifyTicks(110); + ASSERT_EQ(tuner.Best(), nullptr); + ASSERT_EQ(tuner.NextConfig(), 1); + tuner.NotifyTicks(100); + } + ASSERT_NE(tuner.Best(), nullptr); + EXPECT_EQ(*tuner.Best(), 1); +} + +TEST(MatMulSchedulingTest, ActivationConversionCandidates) { + for (size_t M : {size_t{1}, size_t{17}}) { + EXPECT_EQ(MMParACandidates(M, MMSchedule::kAutoTune).size(), 4); + for (MMSchedule schedule : {MMSchedule::kFixed, MMSchedule::kFixedMinK}) { + const auto candidates = MMParACandidates(M, schedule); + ASSERT_EQ(candidates.size(), 1); + EXPECT_EQ(candidates.front(), MMParA::kK1); + } + } +} + +#if GEMMA_ONEDNN_BRGEMM +TEST(MatMulSchedulingTest, BRGeMMCandidates) { + for (size_t M : {size_t{32}, size_t{64}, size_t{127}}) { + for (size_t K : {size_t{32}, size_t{1024}, size_t{16384}}) { + const auto all = BRGeMMCandidates(M, K, 128); + const auto fixed = BRGeMMCandidates(M, K, 128, false); + ASSERT_EQ(fixed.size(), 1); + EXPECT_EQ(fixed.front().M_blk, all.front().M_blk); + EXPECT_EQ(fixed.front().batch_size, all.front().batch_size); + const auto min_k = BRGeMMCandidates(M, K, 128, false, true); + ASSERT_EQ(min_k.size(), 1); + const size_t splits = hwy::DivCeil(K / 32, min_k.front().batch_size); + for (const auto& candidate : all) { + EXPECT_LE(splits, hwy::DivCeil(K / 32, candidate.batch_size)); + } + } + } +} +#endif + +TEST(MatMulSchedulingTest, ActivationKeys) { + std::set keys; + for (size_t M : {size_t{1}, size_t{4}, size_t{64}, size_t{65535}}) { + for (size_t K : {size_t{1}, size_t{1} << 19, (size_t{1} << 20) - 1}) { + for (size_t N : {size_t{4}, size_t{1} << 19, (size_t{1} << 20) - 1}) { + for (size_t num_B : {size_t{1}, size_t{2}}) { + const auto original = static_cast(MMKeys::BucketM(M)) | + (static_cast(K) << 16) | + (static_cast(N) << 40) | + (static_cast(num_B) << 60); + EXPECT_EQ(original, MMKeys::KeyFromDims(M, K, N, num_B)); + for (MMActivation activation : + {MMActivation::kBF16, MMActivation::kI8, + MMActivation::kI8Block}) { + const auto key = MMKeys::KeyFromDims(M, K, N, num_B, activation); + EXPECT_NE(key, MMKeys::kPadding); + EXPECT_TRUE(keys.insert(key).second); + EXPECT_EQ(key, MMKeys::KeyFromDims(MMKeys::BucketM(M), K, N, num_B, + activation)); + } + } + } + } + } +} + } // namespace gcpp #endif