Batchwise optimizer performance improvements - #795
Draft
jnsLs wants to merge 7 commits into
Draft
Conversation
Rewrite the batchwise optimizer for speed and remove its dependency on ASE. The optimizer now operates directly on torch tensors and can run on either CPU or GPU. - batchwise_optimization: LBFGS update rewritten (~2x faster), fixed-atom masking, device handling fixes, no more ASE Atoms round-trips - neighborlist: torch-based FilterNeighbors, skin-based update tracking that avoids rebuilding unchanged neighborhoods - batchwise_trajectory: new trajectory writer for batchwise relaxations - atomistic/prior_bonds: covalent bond prior can now be specified - tests: unit, integration and benchmark coverage for the optimizer and trajectory writer - howto_batchwise_relaxations: batch size sweep added Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The batch-wise optimizer needs an on-device check for whether the neighbor lists went stale, so that a relaxation step that barely moves the atoms never ships the batch to the cpu. That gate stays. Two things that came along with it did not have to. SkinNeighborList.forward had dropped _remove_neighbors_in_skin, so pairs out to cutoff + cutoff_skin reached the model -- wrong for a model without a cutoff function, extra pair work for one with, and a silent change for the plain SpkCalculator and ase MD users of the transform. Pruning is back, and the converter does the same pruning for the whole batch on its own device on the steps that reuse a cached list. AtomsConverter._transform_inputs was a second conversion path that split a batch, re-ran the transforms and re-collated. It duplicated __call__, and it raised KeyError on any batch without properties.idx -- which is every batch out of BatchwiseTrajectoryReader.frame, so a relaxation could not be resumed from a trajectory frame. update_inputs now rebuilds via batch_to_atoms and __call__ instead; that is within noise of the hand-rolled path, since a rebuild is dominated by the neighbor list construction. Also folds the two copies of the cache-store block into _store_reference, holds the skin neighbor list explicitly rather than indexing transforms[0], and drops two unused imports. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
auxiliary_output_modules let BatchwiseCalculator splice extra output modules into the model it had just loaded, so that a prior contributed to the energy and hence to the forces. Attaching a prior is a property of the model rather than of the calculator that evaluates it, and schnetpack already has the pattern: a prior writes its own output_key the way ZBLRepulsionEnergy does, and an Aggregation module sums the terms into the total energy. What was there also did not work properly. SpkCalculator accepted the argument and ignored it outright. The two _initialize_model overrides inserted at different positions, len(output_modules) - 1 against 1, and neither re-ran collect_derivatives or collect_outputs, so a prior declaring model_outputs or required_derivatives went unnoticed. CovalentBond assigned to output_key instead of contributing a term, which overwrites the model's prediction outright when given output_key="energy"; it applied its unit factors in the wrong direction; it required equal atom counts; and its defaults were specific to one project. It was referenced nowhere, had no tests and was never listed in the api docs. The howto gains a section building the same thing out of the existing pieces: a HarmonicBond prior shaped after ZBLRepulsionEnergy, and a with_prior helper that splices it plus an Aggregation in ahead of the response module and refreshes the two caches. It is demonstrated on ethanol's C-O bond, restrained to 1.8 Ang and relaxing to 1.77 against 1.44 unrestrained, which is what shows the prior reaching the forces and not just the reported energy. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…xed skin bug in md nbh list, fixed triplets bug
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
1 participant
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.
Rewrite the batchwise optimizer for speed and remove its dependency on ASE. The optimizer now operates directly on torch tensors and can run on either CPU or GPU.