Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add a BatchRelaxer to speed up geometry optimisation of many small structures #854

Draft
wants to merge 9 commits into
base: main
Choose a base branch
from

Conversation

zhubonan
Copy link
Contributor

Add a batch relaxer which combines multiple atoms as a single batch to the underlying model. This can give speed-ups when lots of smaller structures needs to be relaxed using GPU (typical while using the foundation models for high-throughput studies).

The current code is working, but a few useful functionalities can still be added:

  • Support for ensemble models
  • Support other types of MACE like EnergyDipoleMACE?
  • Since we already run the optimiser step-by-step, it does no harm to add some detection for pathological structures (exploded/collapsed) which can make the main worker process to stuck.
  • Automatically determine the best batch size by collecting the timings of the compute method.

@zhubonan zhubonan marked this pull request as draft March 10, 2025 15:40
@ilyes319
Copy link
Contributor

Hey looks great, ping me when you think it is ready to merge.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants