Skip to content

feat: add inference speed benchmark#743

Draft
lwalew wants to merge 1 commit into
ddmms:mainfrom
lwalew:feat/add-inference-speed-benchmark
Draft

feat: add inference speed benchmark#743
lwalew wants to merge 1 commit into
ddmms:mainfrom
lwalew:feat/add-inference-speed-benchmark

Conversation

@lwalew

@lwalew lwalew commented Jul 24, 2026

Copy link
Copy Markdown

Pre-review checklist for PR author

  • I've confirmed the contribution guidelines.
  • I have reviewed and understand all AI-generated code in this PR.
  • I have added human-written tests for the new logic.
  • I have properly cited any upstream algorithms or libraries the AI utilized.
  • I have disclosed significant AI tool usage in the PR description.

Summary

Migrates the inference speed benchmark from the MLIP Audit suite into ml-peg under the molecular_dynamics category.

For each structure in a size-stratified protein dataset (elements N, H, O, S, C), the underlying InferenceSpeedBenchmark times 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:

  • Table metric Forward Time / Atom (µs/atom, lower is better): mean over structures of average_forward_time / num_atoms, skipping structures whose forward pass failed.
  • Scaling plot (line plot): forward-pass time (s) vs. number of atoms, one line per model, sorted by system size and restricted to successful forward passes.

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.

⚠️ Blocked

This PR is blocked and must remain a draft until the inference_speed benchmark lands on the mlipaudit mlpeg-migration branch.

  • ml-peg pins mlipaudit to the mlpeg-migration branch (it carries the ASE BaseCalculator fix and other changes the rest of ml-peg needs). That pin is intentionally left unchanged here.
  • InferenceSpeedBenchmark currently exists only on the mlipaudit branch feat/scaling-benchmark-revamp, at src/mlipaudit/benchmarks/inference_speed/inference_speed.py. It is not on mlpeg-migration.
  • Until inference_speed is merged into mlpeg-migration, the import from mlipaudit.benchmarks.inference_speed.inference_speed import ... will not resolve, so the benchmark cannot run.

Additional notes:

  • The shared mlipaudit wiring (the mlipaudit optional extra + conflicts entry in pyproject.toml, the [tool.uv.sources] pin, and the mlip_audit entry in frameworks.yml) is already present on main and overlaps with other in-flight mlipaudit PRs, so this PR does not touch those files.
  • Data packaging requirement: the S3 object inputs/molecular_dynamics/inference_speed/inference_speed.zip must contain an inference_speed/ directory of size-prefixed .xyz files (e.g. 121_1ay3.xyz) — the benchmark reads data_dir / "{structure}.xyz" from data_input_dir / "inference_speed".

Linked issue

Resolves #742

Progress

  • Calculations
  • Analysis
  • Application
  • Documentation

Testing

Static checks only so far: python -m py_compile on the new modules and all pre-commit hooks (ruff-check, ruff-format, numpydoc-validation, whitespace/EOF) pass.

End-to-end testing is pending two prerequisites:

  1. The mlipaudit dependency landing on mlpeg-migration (see Blocked above).
  2. The S3 data upload of inputs/molecular_dynamics/inference_speed/inference_speed.zip.

The good/bad thresholds in metrics.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, and plot_from_table_column / read_plot patterns.

@ElliottKasoar ElliottKasoar added new benchmark Proposals and suggestions for new benchmarks external Contribution from external framework labels Jul 25, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

external Contribution from external framework new benchmark Proposals and suggestions for new benchmarks

Projects

None yet

Development

Successfully merging this pull request may close these issues.

New benchmark: inference speed

2 participants