feat: add inference speed benchmark#743
Draft
lwalew wants to merge 1 commit into
Draft
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Pre-review checklist for PR author
Summary
Migrates the inference speed benchmark from the MLIP Audit suite into ml-peg under the
molecular_dynamicscategory.For each structure in a size-stratified protein dataset (elements N, H, O, S, C), the underlying
InferenceSpeedBenchmarktimes the model forward pass (energy + forces) and runs a short MD simulation per backend. This PR wires that up for ml-peg's ASE calculators and exposes:Forward Time / Atom(µs/atom, lower is better): mean over structures ofaverage_forward_time / num_atoms, skipping structures whose forward pass failed.This is a wall-clock, hardware-dependent measurement (H100 reference), not a level of theory, so results are only comparable across models run on the same hardware.
This PR is blocked and must remain a draft until the
inference_speedbenchmark lands on the mlipauditmlpeg-migrationbranch.mlpeg-migrationbranch (it carries the ASEBaseCalculatorfix and other changes the rest of ml-peg needs). That pin is intentionally left unchanged here.InferenceSpeedBenchmarkcurrently exists only on the mlipaudit branchfeat/scaling-benchmark-revamp, atsrc/mlipaudit/benchmarks/inference_speed/inference_speed.py. It is not onmlpeg-migration.inference_speedis merged intomlpeg-migration, the importfrom mlipaudit.benchmarks.inference_speed.inference_speed import ...will not resolve, so the benchmark cannot run.Additional notes:
mlipauditoptional extra + conflicts entry inpyproject.toml, the[tool.uv.sources]pin, and themlip_auditentry inframeworks.yml) is already present onmainand overlaps with other in-flight mlipaudit PRs, so this PR does not touch those files.inputs/molecular_dynamics/inference_speed/inference_speed.zipmust contain aninference_speed/directory of size-prefixed.xyzfiles (e.g.121_1ay3.xyz) — the benchmark readsdata_dir / "{structure}.xyz"fromdata_input_dir / "inference_speed".Linked issue
Resolves #742
Progress
Testing
Static checks only so far:
python -m py_compileon the new modules and allpre-commithooks (ruff-check, ruff-format, numpydoc-validation, whitespace/EOF) pass.End-to-end testing is pending two prerequisites:
mlpeg-migration(see Blocked above).inputs/molecular_dynamics/inference_speed/inference_speed.zip.The
good/badthresholds inmetrics.yml(0.5 / 5.0 µs/atom) and the score midpoint are estimates and need tuning against real H100 runs.New decorators/callbacks
None. Reuses the existing
@plot_scatter(line plot),@build_table, andplot_from_table_column/read_plotpatterns.