Add support for caching and re-using dynamics contexts - #191
Merged
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.
This PR extends the
DynamicsCacheused for replica exchange with proper caching, allowing contexts to be re-used for simulations at lambda values other than which they were created. This allows us to run replica exchange on more modest hardware, where there isn't sufficient memory to create all contexts up front, as well as offering the ability to run simulations with a very large number of replicas, e.g. large ABFE simulations.The user can now specify a
max_contextsoption which limits the total number of contexts that are created. Whenmax_contextsisNone, then the existing behaviour is preserved. Settingmax_contextsto less than the number of replicas means that contexts will be re-used. In this situation, we have added book-keeping to store the energy trajectory for each replica so that we can swap it on re-use. We then callset_lambda(),delete_all_frames(), and re-set the integrator clock to the start of the repex cycle. In situations wheremax_contextsis equal to thenum_gpus * oversubscription_factor, there is no performance loss for using the new approach, i.e. you only pay a hit when there are GPU resources that are under utilised. (We warn when this is the case.) When using all replicas, we now suggest an appropriatemax_contextsif GPU memory is exhausted on setup. (We don't automatically fall back to themax_contextspathway due to the two caveats below.)Running with a restricted number of contexts comes with two caveats:
frame_frequencyequal to thecheckpoint_frequency, which is a common production use case anyway.update_constraintsoption that specifies whether constraints should be updated whenmax_contextsis used. (By default they are.) IfFalse, then the user can additionally supply aconstraint_lambda_indexoption that specifies the replica that should be used for the constraint value. I intend to see if there is a good default for this, e.g. should we use the lambda value where the constraint is at its mid-point, or one of the end-states? In previous work I found that, when running simulations in forward/reverse, you'd get the correct result when pinning the constraint to one end-state, i.e. lambda=0 for forwards and lambda=1 for reverse, then averaging the ddG values.The PR also makes the GCMC equilibration pathway consistent for the
RunnerandRepexRunner, as well as removing minimisation on restart, which was inconsistent with the other restart guards.