Integrating with Colmena for concurrent training, sampling, reference computations

[WardLT]

I was really excited to learn about Psiflow from the Parsl team a few days ago because I've been exploring concepts for simultaneously managing training, simulation, and reference energy calculation. To date, I've been working on these ideas as part of the "FFF" library but I'd really prefer to work on a project with others (i.e., yours) than create another.

I'm thinking of re-implementing some of your "Learning" classes using Colmena - a tool built on top of Parsl for managing dynamic workflows. Before getting started I wanted to confirm this was the part of Psiflow that contains logic for submitting/processing Parsl tasks, and confirm make sure that I'd not be creating changes which duplicate work you're already doing.

Any advice on where/how to get started?

[svandenhaute]

Excited to hear you're interested in building on top of psiflow!
The learning classes define the learning algorithm, i.e. the specific sequence of (train, sample, reference_evaluate) operations. They mainly use BaseWalker, BaseModel, and BaseReference instances to submit Parsl tasks. Simple (mostly static) approaches such as SequentialLearning and IncrementalLearning are very easy to implement in this way. Some of it is dynamic with a bunch of join_apps (the best example being the sample + reference evaluation loop, which allows for a certain number of retries based on the exit status of each). I've found that to work, but also to be the main source of subtle bugs in the past.
Concurrent learning, in which (train, sample, reference_evaluate) all run concurrently at any given moment, is somewhat more difficult. I've made a first attempt at doing so with bare Parsl in the ConcurrentLearning class. I'm not entirely sure what the benefit would be of using Colmena for this, but if there's good arguments for it I'm very happy to help you with this. While psiflow's high-level API still makes a lot of sense, there's definitely room for improvement in the way in which the Parsl apps are defined and executed. That, together with better workflow monitoring was going to be my main focus over the next months.

As you can see in the code, the actual Parsl apps are defined in the various walker, model, and reference subclasses using the create_apps method. It uses the loaded psiflow configuration in order to essentially create partial functions with execution parameters automatically built into it in order to not clutter other parts of the code with that. All the apps are stored in the context instance, so they're accessible from everywhere.

May I ask what type of application your thinking of, and on what scale?

[WardLT]

Hi, Sander! Thanks for the quick reply.

There are two challenges I'm looking at: training reactive forcefields for dehydrogenation using CCSD(T) and dynamic stability of MOFs. In terms of scale, I'm looking to deploy of workflows on exascale systems (the work is part of the US's Exascale Computing Project). I'm also working on forcefields for water trained at the MP2 level, but that project's late enough stages that I probably won't redo it with Psiflow before the paper.

Colmena will fill a similar role to join apps, in that it provides a way to express what happens when a task completes. Instead of having the logic embedded in stateless workflow function, you have it as methods of a stateful "steering class." An example for concurrent learning: https://github.com/exalearn/fast-finetuned-forcefields/tree/main/colmena/initial/

I'll take a look at ConcurrentLearning to get an idea of what the application looks like. It'll probably be early July before I get coding, but wanted to at least chat now so I can think on if/how to approach using Psiflow.

[svandenhaute]

Very interesting! Which QC software would you use for the CCSD(T) calculations? ORCA support is something I've been planning to add for a while now. It would also require some small changes throughout the code as now it is assumed that all Atoms instances are 3D periodic. This is minor, though, and I can try and get that done before you start.

What about file storage? Now, it's kind of assumed everything resides on a shared file system because this has always been the case for us so far. Do you plan to store intermediate results in some (compressed?) database in order to not overload the file system in terms of inodes? Psiflow is rather inode-hungry...

[WardLT]

Very interesting! Which QC software would you use for the CCSD(T) calculations?

I use TAMM, which is part of the eventual NWChemEx code and can run on Aurora's Intel GPUs.

Good point about the periodicity. All of my workflows have assumed nonperiodic, but I'm hopeful we can make it easy to run either.

What about file storage?

I keep the training data in an ASE db file (sqlite backend, which I'll eventually have to replace), and keep other details (e.g., lists of sampled structures, model(s)) in memory or node-local storage.

Where do you see the most disk usage?

[svandenhaute]

Most of it is consumed by extended XYZ files of MD trajectories and/or datasets. The XYZ header sometimes contains additional information related to e.g. the location of stdout/stderr of a QM singlepoint of the corresponding confguration, or the value of some collective variable. It's probably not the best way to do it. Do you think it makes sense to switch to sqlite?

If you plan to use TAMM/NWChemEx, you'd have to write a TAMMReference class similar to the existing CP2KReference, with pointers to basis set files and other input options. If you have a MWE, I can take a look as well.