Skip to content

Latest commit

 

History

30 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

phosbench

Production MLIP molecular dynamics on consumer GPUs — an end-to-end deployment study on black phosphorus (RTX 3080 Ti, 12 GB).

Machine-learned interatomic potentials (MACE foundation models) promise DFT-quality MD at classical cost, but the headline acceleration numbers (NVIDIA's cuEquivariance-in-LAMMPS results, NVIDIA Technical Blog, Oct 2025) are water benchmarks on datacenter A100/H100s, throughput-first, and the closest independent study (arXiv:2510.23621) is likewise water-throughput, covering e3nn-vs-cuEq precision arms on an RTX 2080 Ti. A consumer GPU tells a different story: a GeForce RTX card runs fp64 at 1/64 of fp32 throughput and has only 12 GB of VRAM. This study measures, on monolayer phosphorene (a maximally anisotropic 2D semiconductor), what it takes to run trustworthy production MD on that hardware: when cuEquivariance pays off and when it does not, what fp32 does to physical observables (phonon dispersions, anisotropic elastic constants, energy drift, not force RMSE), and where the time actually goes (Nsight). Every figure caption states the deployment decision it informs.

A 4-minute walkthrough (video) → youtu.be/44fmO8QmyGQ.

The 2-page decision memo for a lab deciding today → docs/engagement-memo.md.

Scope discipline: this extends the published datacenter throughput results to the consumer break-even boundaries and solid-state observables that drive real lab deployment decisions. It is not the first consumer MACE benchmark (arXiv:2510.23621 includes an RTX 2080 Ti). It is, to our knowledge, the first break-even/OOM characterization and the first precision-vs-solid-state-observable study for the MACE + cuEquivariance stack.

TL;DR — the recommendation matrix

In plain terms: for small simulations (below roughly 400 atoms) the default settings already win, and turning on the "GPU accelerator" flag actually makes them slower. Above that size, that one flag cuts the time per step by 1.4–5× and lets a single 12 GB workstation card handle ~3.9× bigger simulations. Using the faster (single-precision) math is physically safe here — the real risk was never the numbers, it was the AI model itself, which silently got a key property wrong until we caught it and fixed it.

For a lab running MACE-class potentials on a 12 GB consumer GPU:

system size backend & precision why (measured)
below break-even (~313–346 at. large, ~373–454 medium, ~947–982 small) e3nn / fp32 (cuEq off) step time is host-bound (GPU kernels ≤ 8% of step); cuEq kernels make it slower (×0.74–0.96 for the MP-0 family, down to ×0.40 for OMAT-medium at 64 atoms)
any size, if you wire CUDA-graph capture cuEq / fp32 + CUDA graph (fixed-topology or padded) the row above is an eager-dispatch artifact: graph capture removes the launch overhead and the crossover disappears — ×4–9 per-step below ~512 atoms, numerically exact (Part 2)
break-even – ~3,000 atoms cuEq / fp32 ×1.4 (small model) to ×5.0 (large model) per-step speedup at 2,944 atoms
~3,000 – 23,000 atoms cuEq / fp32 (above 5,760 atoms, the only option) e3nn OOMs at 2,944 (medium/large) – 5,760 (small) atoms; cuEq's ×3.4–6.1 smaller activation memory reaches 11,520 (medium) – 23,040 (small) atoms
any size, fp64 don't — use e3nn/fp64 sparingly for reference data fp64 costs ×3.2 (64 at.) to ×10 (1,408 at.) on GA102; last working rung 1,408 atoms (OOM at 2,944)
property workflows (phonons, elastic) hybrid: displaced-force evaluations on e3nn/fp64, production MD on cuEq/fp32 fp32 finite-difference noise pollutes small displacements: at the standard 0.01 Å amplitude it adds ~0.03–0.04 cm⁻¹ RMSE and spurious imaginary acoustic points — harmless under the standard 0.3 THz 2D-flexural tolerance, but enough to flip a naive ω² > 0 stability check; fp32 forces an explicit displacement-amplitude choice
CPU (8-core Ryzen 5800X) never for this workload GPU wins ×13 even at 64 atoms, ×64 at 1,408

error budget The precision chart: per observable, fp32 and backend errors are 1–3 orders of magnitude below the model-vs-literature error.

The error budget that matters: across every observable we measured, |fp32 − fp64| and |cuEq − e3nn| are ≤ 3 % (typically ≤ 1 %) of |model − DFT-literature|. Worst case is C22, where 0.07 N/m of precision spread stands against 2.4 N/m of model error. Numerical precision is not the error you should pay to reduce — model choice is (see Zero-shot validation failure, below).

Headline figures

break-even vram
Below the model-dependent crossover (~313–982 atoms), leave cuEq off — the host, not the kernel, limits the step. On 12 GB, cuEq buys a ×3.9 reachable-size increase before it buys speed; size production cells from oom_boundary.csv.
time share phonons
Faster kernels cannot fix host-bound small systems; after cuEq, the ASE loop is the next bottleneck. cuEq/fp32 (dashed) reproduces the e3nn/fp64 dispersion to ≤ 0.08 cm⁻¹ per branch — fp32 production MD is physically safe here.
The raw throughput and the two crossovers behind the map (click to expand)

throughput vs size The un-normalized view behind the speedup map: absolute MD throughput (ns/day) per backend/precision as system size grows; × marks the last size that fits in 12 GB. cuEq/fp32 owns the large-system frontier, e3nn dies early to OOM, and fp64 is off the bottom. This is the data the recommendation matrix compresses.

break-even modes Force-call (kernel) crossover vs MD-loop (wall-clock) crossover, per model: the kernel speedup usually leads the wall-clock one because ASE host overhead delays the MD crossover, so enabling cuEq the moment the kernels break even is premature; wait for the MD line. Medium is the measured exception (MD 373 vs force-call 454 atoms).

Key findings

  1. Break-even is real and model-size-dependent. cuEquivariance crossover sits at ~313–346 atoms (MACE-MP-0 large), ~373–454 (medium), ~947–982 (small). For three of four models the wall-clock (MD) crossover lands later than the bare force-call crossover — ASE host overhead delays it; medium is the measured exception (373 MD vs 454 force-call). Below break-even cuEq costs up to 26 % (and 60 % for OMAT-medium at 64 atoms). The crossover positions carry roughly ±1 ladder-rung uncertainty from the unlocked consumer boost clocks (SM clock/temp logged per data point; medians of per-step laps reported). Part 2 (below) shows this crossover is a launch-overhead artifact — it disappears under CUDA-graph capture.

  2. cuEquivariance's bigger gift on 12 GB is memory, not speed. At 2,944 atoms (medium model) cuEq peaks at 1.4 GiB where e3nn needs 7.4 GiB (×5.3; ×3.4–6.1 across models); the reachable system size grows ×3.9 (2,944 → 11,520 atoms). The OOM boundary per config is tabulated in results/figures/oom_boundary.csv.

  3. After acceleration, the bottleneck is the host. Nsight: at 140 atoms GPU kernels occupy 1 % of step wall-time (cuEq) — faster kernels cannot help; at 2,944 atoms e3nn kernels occupy 60 % of step wall-time vs cuEq's 10 % (kernel share = summed kernel durations / NVTX force_eval span on the nsys timeline, a wall-time ratio, not SM occupancy; hardware counters were permission-blocked). cuEq shifts the limiter from kernels to the Python/ASE loop: the production path beyond this study is LAMMPS ML-IAP (Kokkos) or CUDA-graph-style batching, per NVIDIA's datacenter results. Part 2 (below) runs the CUDA-graph half of that sentence: ×9.1 → ×1.12 (140 → 2,944 atoms) — the decay curve this finding predicts.

  4. fp64 is effectively unavailable on consumer Ampere. ×3.2–×10 measured cost (the fp64 GEMMs dominate the timeline: the three largest cutlass...d884gemm kernels take ~74 % of kernel time, and all d884gemm kernels together ≈ 82 %), ×2 memory, last working rung 1,408 atoms (OOM at 2,944). Consumer deployment forces the fp32 question — which is exactly why the accuracy gates below matter. (On A100/H100 fp64 is 1:2, so the datacenter column of the matrix differs.)

  5. fp32/cuEq preserve the physics — measured, not assumed.

    • Parity gate (cuEq vs e3nn @ fp32, 140 atoms): ΔE below fp32 representation resolution (both backends round to the bitwise-identical fp32 energy), max |ΔF| = 5×10⁻⁴ meV/Å.
    • Phonon dispersion: per-branch RMSE ≤ 0.08 cm⁻¹ (fp32 vs fp64), invisible against the ~20 cm⁻¹ scale separating the model's top optical mode from the bulk Raman lines (context, not a fitted comparison).
    • Anisotropic elastic constants: C11/C22 = 33/105 N/m, identical to 0.22 % across e3nn-fp64 / e3nn-fp32 / cuEq-fp32, from relaxed-ion energy curvature (R² > 0.9997). C12 is not reliably determinable on this pbc=[T,T,F] geometry with mace-torch 0.3.16: the ≈30 N/m obtained from the stress cross-slope rescaled by the measured ×17.8 factor (finding 8) rests on a uniform-scaling assumption, so it is indicative only and should not be weighed against the ~18 N/m PBE literature value.
    • NVE drift (512 atoms, 25 ps, 300 K): |slope| ≤ 0.012 µeV/atom/ps in all three cells — fp32/cuEq conserve energy as well as fp64 at this horizon. Bonus deployment fact: free-running NVE (no per-step sync) runs cuEq at 31 ms/step vs the 78 ms/step the per-step-synced harness measures at this size — cuEq pipelines launches ahead, so the sweep's break-even numbers are conservative for production MD.
    • NPT lattice: invalidated by the upstream stress inconsistency (see finding 8); the failure itself is precision/backend-independent.
    The three gate figures behind finding 5 (click to expand)

    parity gates Swapping e3nn→cuEq at fp32 is numerically free (both bars far under the 1 meV/atom, 1 meV/Å gate line); the fp32↔fp64 gap is larger but still ~2 orders below the gate — precision, not backend, is the knob that moves, and even it stays inside spec.

    elastic anisotropy C11/C22 are identical to 0.22 % across all three cells and reproduce the ~3–4× DFT anisotropy (dashed lines: DFT C11 ≈ 24, C22 ≈ 103 N/m, reference, not a fitted target). Accept a cell for strained MD only if it keeps C22 ≫ C11; C12 is the one number off the broken stress path, shown for completeness, not trusted.

    NVE drift NVE energy drift over 25 ps at 300 K: |slope| ≤ 0.012 µeV/atom/ps in all three cells — fp32/cuEq conserve energy as well as fp64 at this horizon, so accept a cell for long production MD on this criterion.

  6. …but fp32 changes how you must measure. Finite-difference phonons at the standard 0.01 Å displacement pick up fp32 force noise: spurious imaginary acoustic artifacts at −0.007 THz (e3nn) to −0.012 THz (cuEq) that shrink (cuEq: −0.003 THz) or change sign (e3nn) at 0.05 Å, while fp64 shows the opposite trend, clean at 0.01–0.03 Å and picking up its own −0.002 THz anharmonic artifacts at 0.05 Å. No verdict changes under the standard 0.3 THz 2D-flexural tolerance; the deployment point is that fp32 makes the displacement amplitude an explicit engineering choice instead of a default. Hence the hybrid policy in the matrix.

    phonon noise vs displacement fp32 finite-difference force noise falls as the displacement grows, so the 0.01 Å textbook default is the worst case for fp32: keep displaced-force phonon/elastic workflows on e3nn/fp64 and production MD on cuEq/fp32.

  7. Zero-shot validation failure, caught, diagnosed, and fixed. Caught: all three foundation models tested (MACE-MP-0, MACE-MPA-0, MACE-OMAT-0) reproduce the zigzag lattice constant within 0.9–2.6 % but compress the soft armchair axis by 7–10 % (4.17–4.30 Å vs 4.62 Å DFT) — exactly the direction whose DFT stiffness is ~4.3× lower (C11 ≈ 24 vs C22 ≈ 103 N/m); armchair stiffness comes out ~38 % high. Diagnosed: a first fine-tune on the open GAP-20 dataset (forces-weighted loss, 40 epochs) repaired nothing: validation energy RMSE sat at 106 meV/atom; the soft-axis lattice is an energy-landscape property, so underfit energies leave it broken. Fixed: an energy-weighted stage-two (SWA, E:F = 1000:100) fine-tune (≈2 h of training wall-time on this same RTX 3080 Ti, both stages, from the training logs; budget ~a GPU-day end-to-end with data prep and validation iterations) lands every observable within ~5 % of literature: a −2.0 %, b −0.1 %, C11 23.0 N/m (−4 %), C22 100.8 N/m (−2 %), anisotropy 4.38 (DFT 4.29), top optical mode 463 cm⁻¹ vs 467 cm⁻¹ Raman. Precision/backend invariance persists after the fix (C11: e3nn-fp64 23.01 vs cuEq-fp32 22.94).

    model fix Validate, then fine-tune: the gate that failed three foundation models defines the acceptance test their replacement must pass.

  8. Measured here: MACE analytic stress is ×17.8 off on this slab. Hellmann-Feynman check (scripts/90_diag_stress_hf.py): analytic get_stress() is 17.81× smaller than dE/dε of the same energy surface (e3nn/fp64/CPU — independent of cuEquivariance; a fully-periodic control on the same script isolates the slab geometry as the trigger). Stress-slope elastic constants are unusable on pbc=[T,T,F] geometries with mace-torch 0.3.16; this repo's C11/C22 use energy-curvature fits (R² > 0.9997). The bug also poisons NPT: the barostat balances an unscaled kinetic pressure against a ×17.8-undersized virial, so the cell inflates monotonically (~+20 % in-plane over 50 ps, identical across e3nn-fp64 / e3nn-fp32 / cuEq-fp32, three-way agreement that exonerates precision/backend and indicts the stress path). Deployment rule until the upstream fix: no Berendsen/Parrinello-style barostats on partially periodic MACE systems. Upstream issue with this minimal repro: ACEsuit/mace#1509.

    NPT lattice inflation This is the bug manifestation, not a physical result: with the ×17.8-undersized virial the barostat cannot find equilibrium, so both in-plane lattice constants inflate monotonically (~+20 % over 50 ps) and walk straight out of the fp64 reference band, identically across e3nn-fp64 / e3nn-fp32 / cuEq-fp32. Read it as "what the stress bug does to a barostat," not as thermal expansion. The three-way agreement is the evidence that exonerates precision/backend and indicts the stress path.

Known-issues table (what we worked around, honestly)

issue consequence here
cuEq + fp64 reported broken upstream (MACE #1203, #1298) did not reproduce on this stack: the cuEq/fp64 probe ran, matched the fp64 reference, and nsys shows real cuEq kernels (segmented_polynomial_* + fp64 d884gemm), no silent fallback. The cell stays out of the headline matrix anyway: GA102 fp64 is hardware-inappropriate (finding 4)
SM86 is not a tuned cuEq target (kernels are SM80/90/100+) speedups here are a lower bound; verified real cuEq kernels run via Nsight (segmented_polynomial_*), no silent fallback
torch.compile × cuEq zero-gradient bug (cuEq #77) torch.compile disabled everywhere
MACE slab stress ×17.8 (this work, 90_diag_stress_hf.py) elastic constants via energy curvature; in-plane NPT/barostat results are invalid
consumer boost clocks drift no root to lock clocks → SM clock/temp/power logged per measurement, medians of per-step laps reported
ncu hardware counters need NVreg_RestrictProfilingToAdminUsers=0 kernel analysis via nsys timelines only (sufficient for time-share)

Reproduce

# environment (Python 3.11; torch needs the CUDA 12.8 wheel index)
conda create -n scicomp python=3.11 && conda activate scicomp
pip install -r requirements.txt \
    --extra-index-url https://download.pytorch.org/whl/cu128
# full transitive lock: docs/requirements-freeze.txt
# one-time: structure + gates. In the published run the physics gate returns
# nonzero because every zero-shot foundation model fails; that negative gate is
# part of the result, not a stack failure.
bash scripts/stage_a.sh || true
# sweeps / accuracy arms / profiling (hours; queue-friendly, resumable)
python scripts/10_sweep_throughput.py --backends e3nn,cueq --dtypes float32 \
    --models small,medium,large,medium-omat-0 --modes md,force_call
python scripts/20_phonons.py --displacements 0.01,0.03,0.05
python scripts/21_elastic.py && python scripts/23_elastic_recompute.py
python scripts/22_md_stability.py
bash scripts/30_profile_nsys.sh
python scripts/40_make_plots.py && python scripts/41_error_budget.py
# repair the armchair axis by fine-tuning on GAP-20 (Zenodo 4003703)
python scripts/50_finetune_prep.py
bash scripts/51_finetune.sh          # first pass, force-weighted; recorded underfit
bash scripts/53_finetune_swa.sh      # stage-two energy-weighted SWA fix
python scripts/55_ft_validation.py
python scripts/42_model_fix_figure.py

Pinned stack (verified): mace-torch 0.3.16 · cuequivariance(-torch/-ops) 0.10.0 · torch 2.11.0+cu128 · e3nn 0.4.4 · ase 3.28.0 · phonopy 4.1.0 · driver 610.43 (CUDA 13.3) · Ubuntu 24.04 · RTX 3080 Ti 12 GB · Ryzen 7 5800X.

Protocol design, gate thresholds and the schedule that produced this in ~2 days: PROTOCOL.md. The 2-page consulting-style summary for a lab deciding today: docs/engagement-memo.md.

MPI communication share — a scale-out profiling demo

Everything above is single-rank, single-GPU ASE. The scale-out question an SA must answer first is when does adding ranks stop paying? Here is the methodology on a workload that runs on the CPU head node today.

What was measured. A fixed-size Lennard-Jones melt (500,000 atoms, 800 MD steps, LAMMPS 22 Jul 2025 + OpenMPI on a Ryzen 5800X) run as strong scaling at np ∈ {1, 2, 4, 8} — one rank per physical core, SMT deliberately unused. Each run's LAMMPS loop-timing breakdown (Pair / Neigh / Comm / …) was parsed straight from the log.

MPI comm share Comm share rises 0.7 %→9.0 % from 1→8 ranks (LJ melt, 500k atoms): profile a customer this way to find the rank count where communication eats the pair-time win and scale-out stops paying.

ranks Pair % Comm % loop time
1 84.6 0.7 134.8 s
2 82.5 1.5 71.7 s
4 79.0 2.5 39.5 s
8 70.1 9.0 24.9 s

As ranks rise on the fixed problem the communication share climbs (0.7 %→9.0 %) while the pair (force-compute) share falls (84.6 %→70.1 %); wall time still drops (5.4× at 8 ranks) but the marginal return is shrinking. The rank count where the rising Comm tax overtakes the marginal Pair speedup is where scale-out stops paying: the single number to establish before committing a customer to more nodes/GPUs.

Scope (honest). Classical LJ, single node, one run per rank count, a methodology demonstration, not MACE-in-LAMMPS. The MACE path at scale is LAMMPS ML-IAP (pair_style mace), where LAMMPS prints the same Pair/Comm breakdown, so this reading transfers directly; only the absolute Pair cost (a neural potential is far heavier than LJ) and thus the comm-vs-compute crossover rank change.

Part 2 — closing the host gap: CUDA Graphs

Added 2026-07-06, after the frozen Part 1 campaign. Full study with methods and per-arm data: docs/cudagraph-study.md; code: scripts/70–73.

Finding 3 named the fix but did not run it. Part 2 captures the MACE force evaluation into a CUDA graph (fixed-topology path through the calculator's own converter, static I/O buffers, hand capture with torch.cuda.graph; no torch.compile, cuEq #77) and measures what comes back. Parity gate before any timing: the frozen topology first matches the normal ASE calculator path, then graph-replay forces match frozen eager to max |ΔF| ≤ 8×10⁻⁷ eV/Å (fp32 roundoff) at every size.

atoms eager (ms/step) graph (ms/step) speedup peak VRAM eager → graph
140 17.14 1.88 ×9.1 103 → 67 MiB
512 17.23 4.18 ×4.1 327 → 103 MiB
2,944 23.26 20.85 ×1.12 1,487 → 143 MiB

cudagraph speedup The decay curve is finding 3 validated by intervention: at 140 atoms the step was ~94 % host launch overhead — a graph collapses it; at 2,944 atoms kernels already fill the step, so there is nothing left for a graph to reclaim.

The headline consequence: finding 1's break-even collapses at the force-eval level. Side by side, cuEq + graph beats e3nn + graph at every size: ×7.1 at 140 atoms (1.88 vs 13.30 ms), the very size where eager cuEq loses; at 2,944 atoms e3nn cannot even capture (OOM: graph pool + e3nn's 7.4 GiB working set exceeds 12 GB) where cuEq + graph fits easily. The eager crossover was never about kernel work: cuEq's many tiny kernels simply pay more launch dispatch than e3nn's few heavy ones. Remove the dispatch and cuEq wins everywhere; hence the matrix's second row.

break-even transformation The same decision chart before and after: under eager dispatch the cuEq advantage dips below 1 (finding 1's break-even, as seen by this harness); under graph capture it never crosses, holding ×4.6–11.9 across 64–1,408 atoms. Left panel: cuEq eager sits on a flat ~17 ms launch-overhead floor all the way through 1,408 atoms; the graph removes that floor.

Batch the below-break-even regime: 8 replicas of the 140-atom cell packed into one captured graph run at 1.04 ms/cell — ×2.1 vs eager, ×1.7 vs replaying a single-cell graph 8 times (parity 5.8×10⁻⁷ eV/Å). For ensembles of small independent cells, don't just capture — batch the capture.

Honest boundary: all of this is the per-step ceiling at frozen topology, the regime between neighbour-list rebuilds. Production MD needs periodic recapture or padded capture (MACE ships padding_tools) and end-to-end numbers including that cost; this section must not be cited as an end-to-end MD speedup. An unexpected free observation: in a bare force-eval loop, free-running ≈ synced eager (unlike the MD loop's 31 vs 78 ms/step) — without integrator Python between calls there is nothing to pipeline, isolating per-call dispatch as the cost only a graph removes.

Where this goes next

  • Datacenter column: the sweep harness is config-driven and re-runs unmodified on A100/H100 (fp64 at 1:2 changes the precision economics; SM80-tuned kernels should raise the cuEq column). Published water numbers suggest ×3–5 at scale; our large model reaches that bracket at 2,944 atoms (×5.0); medium sits at ×2.4, small at ×1.4–1.7.
  • Scaling out: LAMMPS ML-IAP (Kokkos) is the supported multi-GPU path, the configuration behind NVIDIA's Oct 2025 LAMMPS numbers. Note the deployment subtlety: the ASE-side enable_cueq=True runtime conversion is not what LAMMPS consumes; LAMMPS takes the exported e3nn model through its own cuEquivariance-accelerated ML-IAP route. An ASE-vs-LAMMPS single-GPU step-time comparison on this card is the next measurement on the list.
  • Model fix: done; see finding 7 (scripts/50_finetune_prep.py, scripts/53_finetune_swa.sh, scripts/55_ft_validation.py). The benchmark numbers needed no re-run: they are model-fidelity-independent by construction.
  • Host gap, CUDA-graph half: done — see Part 2 (scripts/70–73, docs/cudagraph-study.md). The remaining honest step on this branch is padded capture wired into a real MD driver, measured end-to-end including recapture cost.

About

Production MLIP molecular dynamics on consumer GPUs — a measured MACE + cuEquivariance deployment study on phosphorene (break-even maps, precision error budget, Nsight, RTX 3080 Ti).

Topics

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages