[evaluation] Add Marin-native OLMoBaseEval Easy Table 9 BPB evaluator#6726
[evaluation] Add Marin-native OLMoBaseEval Easy Table 9 BPB evaluator#6726Calvin-Xu wants to merge 11 commits into
Conversation
Port the OLMoBaseEval Easy Table 9 bits-per-byte suite (51 components) into Marin so data-mixing checkpoints can be evaluated region-locally on Iris/TPU, replacing the convoluted external Stanford-SC eval path. The runtime consumes a frozen, model-independent request set (the exact (context, continuation) gold pairs, exported once from official OLMo-Eval) and scores it natively in Levanter: continuation-masked fp32 log-probs over a UTF-8-byte denominator, the unweighted per-instance mean per task, then the 51 Table 9 components (MMLU collapsed to four size-weighted category buckets) and the unweighted-mean macro. No runtime dependency on SC or OLMo-Eval. Validated against the SC oracle on the baseline_proportional 300m checkpoint: 6 of 7 subset components match to <=1.7e-7 on a us-east5 v6e-8 (highest matmul precision; codex_humaneval is byte-exact on CPU fp32), aggregation reproduces the SC macro to <=2.2e-16, and tokenization matches official OLMo-Eval exactly (no BOS). The spec and the canary parity report are under .agents/projects/.
…roject constant Factor the SC-compatible metric-key emission into metrics.sc_compat_metrics so the partial-request-set path in run.py and build_wandb_metrics share one implementation, and default the CLI --wandb-project to run.DEFAULT_WANDB_PROJECT instead of a raw literal.
|
Claude finished @Calvin-Xu's task in 7m 49s —— View job Code reviewReviewed all 18 added files with parallel correctness + CLAUDE.md/TESTING.md compliance passes (two independent bug passes, one guidelines audit, one test-quality audit), then validated each finding against the source. The parity-critical math is solid — both bug passes independently confirmed: the BPB sign/denominator, the unweighted per-instance mean ( Three findings, posted inline:
PR description: good — leads with what the change does and why; the parity-evidence paragraph is substantive justification, not boilerplate. No issues. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 90c9998579
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| JsonFileTrackerConfig(output_path=this_output_path()), | ||
| ), | ||
| per_device_eval_parallelism=per_device_batch_size, | ||
| mp=jmp.get_policy("c=bf16"), |
There was a problem hiding this comment.
When this evaluator is launched through experiments/evals/olmo_base_eval_table9.py, olmo_base_eval_step builds the trainer with c=bf16. The scorer then loads the checkpoint with compute_dtype=mp.compute_dtype and casts the model to that compute policy before computing loss, so the Iris parity/canary runs use bf16 model compute rather than the fp32 evaluator described by the commit and used by the CLI default. That can silently drift BPB/W&B results for the main launcher; default this path to f32 or expose/pass the dtype explicitly.
Useful? React with 👍 / 👎.
| from marin.execution.executor import executor_main | ||
|
|
||
| # Frozen Table 9 request set (51 components / 104 scored tasks), region-local to us-east5. | ||
| REQUEST_SET_DIR = "gs://marin-us-east5/raw/eval-datasets/olmo_base_eval_table9/v1" |
There was a problem hiding this comment.
This file is under experiments/, whose scoped AGENTS.md says literal artifact paths must use InputName.hardcoded(...) rather than raw gs:// URLs. The request set and checkpoint constants in this launcher are absolute gs://marin-us-east5/... strings, so the experiment bypasses Marin prefix/mirror path handling and bakes a fixed bucket into the DAG; wrap these artifact paths (and allow request_set_dir to accept the resolved InputName) instead.
Useful? React with 👍 / 👎.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| JsonFileTrackerConfig(output_path=this_output_path()), | ||
| ), | ||
| per_device_eval_parallelism=per_device_batch_size, | ||
| mp=jmp.get_policy("c=bf16"), |
There was a problem hiding this comment.
🤖 The executor path forces a bf16 forward, which defeats the parity goal.
olmo_base_eval_step hardcodes mp=jmp.get_policy("c=bf16") with no way to override it. But the parity step launched through this path — baseline_proportional_300m_parity in the experiment launcher — exists specifically to reproduce the SC oracle, and cli.py deliberately defaults to fp32 for exactly that reason ("fp32 reproduces the SC oracle exactly (the SC HF provider scored in fp32)").
loss_dtype=jnp.float32 in score_summed_logprobs only makes the cross-entropy fp32; the upstream logits are still bf16, so the per-token log-probs (and therefore BPB) diverge from the fp32 oracle by bf16 rounding. As written, anyone running the launcher in parity/all mode cannot reach the <=1.7e-7 parity the report cites — and there is no dtype/mp parameter on olmo_base_eval_step to fix it.
Suggest threading a dtype/mp parameter through olmo_base_eval_step (defaulting to f32 for the parity step, matching the CLI), so the two entry points don't silently disagree on precision. Fix this →
| paths = sorted(glob.glob(f"{requests_dir.rstrip('/')}/**/*-requests.jsonl", recursive=True)) | ||
| if not paths: | ||
| raise FileNotFoundError(f"no *-requests.jsonl under {requests_dir}") | ||
| for path in paths: | ||
| with open(path) as handle: |
There was a problem hiding this comment.
🤖 Minor: use fsspec here for consistency with the rest of the module.
convert_olmo_requests discovers paths with glob.glob and reads them with the builtin open(), so it only works on a local filesystem. lib/marin/AGENTS.md says "Use fsspec.open for filesystem access," and the sibling request_set.py already does (and accepts gs://).
This is mitigated by the module's offline-build framing (the docstring says it runs from a local OLMo-Eval checkout), so it's low priority — but switching to fsspec.open + fsspec glob would let the bridge read an exported request dump from GCS without a copy, and keeps the package internally consistent.
| def test_scored_tasks_are_47_leaves_plus_57_subjects(): | ||
| assert len(scored_tasks()) == 104 | ||
| assert set(scored_tasks()) == set(leaf_components()) | set(mmlu_subjects()) |
There was a problem hiding this comment.
🤖 Minor (test quality): this test compares scored_tasks() to its own definition.
scored_tasks() is defined as (*leaf_components(), *mmlu_subjects()), so set(scored_tasks()) == set(leaf_components()) | set(mmlu_subjects()) holds by construction and can't regress independently — TESTING.md flags reimplementation/tautological assertions. The len == 104 is the only non-derived bit, and it's already implied by the 47 and 57 counts asserted in the two tests above.
Same applies to lines 88–89 of test_leaf_components_exclude_the_four_mmlu_buckets (the set-relation holds by construction; the len(leaves) == 47 on line 87 is the meaningful check). Consider dropping the circular assertions and keeping the counts. Not blocking.
…aths Address PR review: the executor launcher built the trainer with c=bf16, so Iris parity/canary runs would have used bf16 compute and drifted from the fp32 results the CLI and parity report describe. Default olmo_base_eval_step to f32 (configurable via dtype) and set jax_default_matmul_precision="highest" in the eval body so every launch path reproduces the ~1e-7 SC-oracle parity regardless of how it is started. Also reference the launcher's request set and checkpoints as prefix-relative InputName.hardcoded paths (per experiments/AGENTS.md) instead of raw gs:// URLs, and pin only the region so Iris picks the available us-east5 zone.
|
🤖 Thanks — both addressed in 2092737:
|
…ological test asserts Use fsspec.open + fs.glob in convert_olmo_requests so the OLMo-Eval request bridge can read an exported dump from gs:// (and stays consistent with the rest of the package). Replace the by-construction set-relation assertions in test_components with a check of leaf_components against the independent EXPECTED_TABLE9 oracle and a uniqueness check on scored_tasks.
|
🤖 Thanks @claude — all three addressed:
|
|
Claude finished @Calvin-Xu's task in 55s —— View job Verifying the three fixes — all confirmed ✅Reviewed
All three review findings are properly addressed. Nice work — the parity fix in particular closes a real silent-drift gap between the two entry points. |
…ns[gold_idx]) convert_olmo_requests scored the singular request.continuation, which is only continuations[0] — correct for single-continuation PPL tasks and for MC tasks whose gold is choice 0 (arc, csqa, mmlu), but wrong for MC tasks whose gold is a later choice (piqa, coqa, socialiqa, winogrande, hellaswag, basic_skills), where it scored the first choice instead of the gold. Select continuations[gold_idx] for multi-choice records (and continuation_prompts[gold_idx] as the context for per-choice-prompt tasks like winogrande), keeping continuations[0] for the single-continuation case (the BPB metric's single-output path). Bump the request set to v2.
…rter.from_hf from_hf probes every registered Levanter config and constructs each one's converter, which for the gemma configs tries to load the gated google/gemma-2b reference repo. Llama checkpoints match before gemma in the registry so they load, but a Qwen3 checkpoint (the delphi 3e18 canaries) is registered after gemma and fails with a 403 on the gated repo. Select the config class directly from config.json's model_type via LmConfig.get_known_choices(), which uses each architecture's own (non-gated) default reference. Verified on both the Llama parity checkpoint and the Qwen3 canaries.
…canaries) Record the full result: macro abs diff 1.9e-8 and all 51/51 components within 1e-3 vs the SC oracle on baseline_proportional, the gold-selection fix that got there (correcting the earlier piqa-as-SC-data hypothesis), the model_type config fix for the Qwen3 canaries, and the two canary macros (the Table-9-optimized DSP mixture is 0.055 BPB below proportional).
1 participant
Make Marin/Iris able to evaluate the OLMoBaseEval Easy Table 9 bits-per-byte
suite (51 components) directly and region-locally, replacing the convoluted
external Stanford-SC (SC) eval path used for data-mixing checkpoints.
The runtime consumes a frozen, model-independent request set — the exact gold
(context, continuation)pairs, exported once from official OLMo-Eval — andscores it natively in Levanter on TPU: continuation-masked fp32 log-probs over a
UTF-8-byte denominator, the unweighted per-instance mean per task, then the 51
Table 9 components (MMLU collapsed to four size-weighted category buckets) and the
unweighted-mean macro. The eval path depends on neither SC nor OLMo-Eval, and the
model config is derived from the checkpoint's
model_type, so both Llama andQwen3 checkpoints load region-locally.
Layout:
lib/marin/src/marin/evaluation/olmo_base_eval/: components, bpb(byte/token/BOS kernel), aggregate, request_set, metrics, run (Levanter/Iris
runtime), generate_requests (offline OLMo-Eval bridge), cli.
experiments/evals/olmo_base_eval_table9.py: Iris launcher (parity + 2 canaries)..agents/projects/: design spec and the canary parity report.On a us-east5 v6e-8 (fp32, highest matmul precision) the evaluator reproduces the
SC oracle on
baseline_proportional300m over the full 104-task / 88,592-instancesuite: macro abs diff
1.9e-8and all51/51components within1e-3(max2.4e-6). The two delphi 3e18 canaries (Qwen3) run end-to-end with W&B writeback;the Table-9-optimized DSP mixture is
0.055BPB below proportional, the expecteddirection.