AI-coupled HPC Workflows

Workflows in high energy physics are multi-scale, comprising quantum field theoretic calculations, detector simulations, and classical reconstruction of physical objects. Each scale has considerable computational requirements and, using only traditional methods, it is impossible to produce sufficient numbers of Monte Carlo events to maintain statistical adequacy with recorded data. As such, ML techniques, including generative adversarial networks, are becoming an increasingly attractive alternative to standard frameworks implementing step-by-step model predictions.

For example, Fig. 2 shows an illustration of the ATLAS Collaboration’s “fast” simulation framework, Atlfast3 [atlfast3], where ML-based techniques are used in place of intensive Geant4 [geant4] simulations. In this configuration, surrogates are employed in lieu of full Geant4 simulations for specific particle types, energies and subdetectors to reduce overall simulation times up to several orders of magnitude. At the same time, those surrogates allow to maintain accurate detector modeling for new physics searches, statistically-limited analyses, background processes, and detector upgrades.

Another significant example of ML-in-HPC can be found in the traditional ab initio MD (AIMD) methods for modeling atomistic phenomena. Due to demanding computational requirements (cubic scaling in the number of electronic degrees of freedom), most AIMD applications are limited to $O\left(1000\right)$ atoms. However, AIMD plays a major role in addressing many issues related to, e.g., drug discovery, complex chemical processes and nanotechnology. As such, tremendous efforts have been afforded to more efficient methods, including ML.

Jia et al. [9355242] offer a powerful example of ML-in-HPC applied to AIMD. Jia’s approach employs an ML-based simulation protocol which uses surrogates (Deep Potential MD) in conjunction with a highly-optimized code (a GPU-accelerated DeePMD-kit [WANG2018178]) to simulate $O\left({10}^8\right)$ nanosecond-long trajectories in 24 hours. This record-setting accomplishment efficiently scaled to the whole 4,560 nodes of the Summit supercomputer, reaching double/mixed-single/mixed-half precision performance of 91/162/275 PFLOPS. Compared to other state-of-the-art, Jia et al. showed more than $O\left({10}^{-3}\right)$ reduction in time-to-solution (TTS) [s/step/atom].

Ref. [9355242] is a prime example of how ML-in-HPC, which is on the cusp of a paradigmatic change as learning approaches influence the way both ODE and PDEs are being solved. For example, Karniadakis et al [raissi2019physics] are investigating how to solve and discover new PDEs via deep learning. For partial differential equations (PDEs), neural operators directly learn the mapping from any functional parametric dependence to the solution. Thus, they learn an entire family of PDEs, in contrast to classical methods which solve one instance of the equation [li2020fourier].

Replacing computationally intensive computations with surrogate approximators aims to reduce TTS often by sacrificing accuracy with respect to more complete models. Optimally balancing TTS and accuracy is a non-trivial task, and the former is often neglected in order to reach the desired accuracy. Proxima [proxima], provides real-time feedback from executing simulations and it has been utilized to develop systematic and automated methods for dynamically tuning surrogate configurations. An iterative simulation workflow, representative of the ML-about-HPC mode, evaluates uncertainties associated with the use of surrogates, concurrently updating configurations based on a distance metric to learn features and an accuracy metric to evaluate prediction. Between iterations, surrogates evolve and replace less optimized ones, providing the coupling between concurrent surrogate tuning and the main HPC campaign.

The Proxima framework has been demonstrated in a Mone Carlo sampling application, where the first-principles Hartree-Fock [osti_5156729] prediction target is replaced with a Proxima-managed surrogate. Mean absolute error (MAE) and TTS comparisons are made between Proxima and a surrogate strategy with a fixed distance threshold (based on scientific trial-and-error). In certain scenarios, the fixed strategy outperforms Proxima in terms of TTS. However, the utility of user-defined error bounds ensures more robust results with Proxima. By determining values for surrogate configurations automatically during workflow execution, Proxima is able to satisfy error bounds while achieving as much as 5.5x speedup in TTS.

Estimating properties of large collections of molecules is often necessary to find candidates for medical therapeutics, next-generation batteries, etc. However, the number of possible candidates for a single application, and therefore the number of different experimental configurations required to test them all, is often intractable. Large-scale workflows have thus come to adopt methodologies to provide an ML model training and retraining runtime to decide which computations to perform based on previous outputs.

The Colmena framework [ward2021colmena] facilitates a user-defined steering for workflow execution. Using an example application involving an ML-guided search of ${10}^5$ molecules with high resistance to oxidation electrolyte design, the Colmena workflow provides components to actively train and learn as simulations are executed. Colmena performs a concurrent execution mode, having ML and traditional simulations running side-by-side throughout the workflow. Candidacy of molecules is evaluated using ML models which are scored based on selection criteria and then are ordered based on their score. Molecules appearing at the top of the ordered list reflect most suitable candidates, and subsequently additional simulations are executed. Colmena is reported to find candidate molecules at rates 100 times that of traditional computational solutions, and scaling up to 1024 nodes (65,636 cores) on Theta supercomputer.

A naturally more complicated scenario is the development of therapeutics. This type of R&D can take years or decades to come to fruition due to the complicated computational modeling involved in searches for candidate drugs and strict FDA approval procedure. In the context of cancer treatments, for example, it is suggested that Ras proteins are involved in nearly a third of all human cancers in the US [SIMANSHU201717]; however, many physiochemical properties of Ras-Raf-membrane dynamics are not fully understood. The inherent multi-scale nature of such processes makes computational modeling challenging, and each scale is traditionally simulated separately.

The massively parallel Multiscale Machine-Learned Modeling Infrastructure (MuMMI) [mummi1-short, mummi2-short] couples three resolution scales with ML-based selection, effectively promoting important configurations from coarse-grained to all atomistic (highest resolution). The autonomy and full power of MuMMI is realized through dynamic co-scheduling of tasks which is achieved by tying together application and coordination layers of the workflow. MuMMI achieves a 98% GPU occupancy for more than 83% of 600,000 node hours, coordinating 24,000 jobs and managing several terabytes of data daily. Furthermore, the split architecture, separating the workflow application from coordination, permits generalizability, making the infrastructure attractive beyond drug design.

Increasing computational power helps to produce more rapidly predictive models, larger volumes of collected data enables higher fidelity predictions. However, improving models typically implies introducing additional complexity such as substantially increasing trainable parameters. As such, building ML models for complex diseases—such as cancer—involves a significant amount of trial-and-error, and intervention from both epidemiologists and ML experts, making diagnosis, detection, prognosis and prediction extremely time-consuming tasks.

Work from Balaprakash et al. [balaprakash2019scalable] introduces a reinforcement, learning-based neural architecture search for autonomous deep learning development. By targeting specific class of cancer data, the automated approach finds neural architectures requiring fewer trainable parameters—thus reducing training time—which produce equivalent or better accuracy to manually finely-tuned architectures. Scalability is demonstrated using 1024 nodes of the Theta supercomputer, with the best neural architecture outperforming the manually designed network in terms of scientific results, and having $11.5\times$ fewer trainable parameters and $2.5\times$ faster training time. These results suggest ML-driven neural architecture search has the potential to accelerate cancer research, allowing researchers in the field to automate neural architecture discovery using HPC.

To accommodate needs of scientists and non-ML experts, a recent effort providing flexible user tools for distributed and scalable reinforcement learning (RL) is the EXARL [exarl-github] from the Co-Design Center for Exascale Machine Learning Technologies (ExaLearn). EXARL enables interfacing with existing exascale applications—e.g., LAMMPS [LAMMPS] and NWChem [nwchem]—which domain scientists can guide using RL algorithms and associated neural network architectures. In addition to the miniGAN proxy application [exarl], the ExaLearn team has demonstrated its usefulness and exercise its scalable RL to a block copolymer application [doi:10.1177/10943420211029302]. This is generally a complicated problem as materials may evolve toward generic states or become trapped in a metastable state, and thus requiring hundreds of experimental trials to reach a target state. By mapping this problem to RL, wherein a NN is trained to update annealing temperatures for subseqeunt block copolymer simulations, EXARL was able to show learning convergence for guiding the annealing process to both equilibrium and non-equilibrium states. Ongoing and future work includes expanding to new scientific domain use-cases, and enablement of further scaling and execution of multi-process applications.

While HPC drives much of scientific research and discovery, the platforms themselves require continuous performance analysis and optimizations to reach their full potential. This is particularly difficult for I/O systems which are commonly bottlenecks in computing systems, trailing in performance with respect to computational capabilities by several orders of magnitude. This is due to complexity of I/O systems, requiring intimate knowledge of the underlying components and potentially thousands of parameters need to mutual optimization.

The SIOX Project [kunkel2014siox] monitors, diagnoses and optimizes I/O system parameters of HPC platforms. The modular design of SIOX provides plug-and-play capability, allowing to use diverse monitoring tools for data production. Plugins use offline ML to predict the performance gains or losses for different optimization actions and online ML to perform anomaly detection. SIOX implements actuator tasks to apply the selected optimizations and evaluator tasks to measure achieved performance.

The three modes of coupling ML with HPC are not mutually exclusive. In fact, the most ambitious multi-scale or multi-stage campaigns involve all three modes. For example, considering the universe of about 10${}^{68}$ possible drug compounds, efficient and high throughput frameworks for early stage drug discovery [Bohacek_et_al:2010] are needed. In silico methodologies need to be improved to better select lead compounds that can proceed to later stages of the drug discovery protocol accelerating the entire process [antunes2015understanding, zhou2020artificial, smith2018transforming]. Innovations that integrate AI and simulation at multiple levels are demonstrating promise in overcoming fundamental limitations.

We discuss IMPECCABLE as a representative campaign that is comprised of ML-in-, ML-out-, and ML-about-HPC workflows, supplanting traditional HPC with learning everywhere. Although, IMPECCABLE was developed for COVID19 therapeutics, the multi-stage and AI-HPC integrated campaign is representative of a range of campaigns in material and molecule design, climate science, inter alia.

The campaign consists of an iterative loop initiated with ML predictions (ML1), followed by data processing stages S1, S2, S3. ML/AI techniques (ML1 and S2) interfaced with physics-based methods estimate docking poses of compounds that are promising leads for a given protein target (S1) and binding free-energy computations (S3). Put together, the campaign glues together learning methods with innovative physics-based methods, with iterative algorithms allowing both upstream and downstream feedback to overcome fundamental limitations of classical in silico drug design [lee2020scalable-short]. It includes high-throughput structure-based protein-ligand docking simulations, followed by iterative refinements to these virtual screening results to filter out compounds that “show promise” in biochemical or whole-cell assays, safety and toxicology tests.

ML techniques overcome the limitations of S1 and S3 by predicting the likelihood of binding between small molecules and a protein target (ML1), and accelerating the sampling of conformational landscapes to bound the binding free-energy values for a given protein-ligand complex (S2). Interfacing ML approaches with physics-based models (docking and MD simulations), we achieve at least three orders of magnitude improvement in the size of compound libraries that can be screened with traditional approaches, while simultaneously providing access to binding free-energy calculations that can impose better confidence intervals in the ligands selected for further (experimental or computational) optimization.

S1 is an example of ML-in-HPC mode(i.e., training and using a surrogate in lieu of computations), while S2 [brace2021achieving] represents a common instance of the ML-about-HPC mode. Although ML-out-HPC was not implemented on HPC platforms at the time of publication, prototypes were used to determine optimal allocation of computational resources across S1-S3 [woo2021optimal].

The impact of the algorithmic, methodological and infrastructural innovations resulted by measuring both raw throughput—defined as ligands per unit time, scientific performance—defined as effective ligands sampled per unit time, as well as the quality of ligands selected [clyde2021high-short, saadi2020impeccable-short, lee2020scalable-short, babuji2020targeting]. Thus, the IMPECCABLE [clyde2021high-short, saadi2020impeccable-short] drug discovery pipeline is the quintessential example of the learning everywhere paradigm.

Traditionally, scientific workflows were defined as either High Throughput Computing (HTC) or High Performance Computing (HPC). The former came to define the distributed workflows of the grid era; the later epitomized by complex and large DAGs of processing. The increasing importance and popularity of ensembles of HPC simulations, resulted in a convergence of these two primary modes – high-throughput of high-performance computing (HT-HPC). And ultimately workflows involving dependencies between large number of (parallel) tasks, and represented by DAG task-graph. The current workflow middleware reflects these dominant paradigms and trends. Moving forward, they will be supplanted by middleware systems which support ML coupling to HT-HPC workflows at multiple levels of the application.

Integrating ML methods with HPC simulations, results in three primary classes of workflows: (1) Hybrid HPC-HTC workflows; (2) ML-coupled workflows, discussed in §III; and (3) Edge-to-center workflows, which typically involve integrating distributed ML with HPC workflows (e.g., with ML on the edge). This is rapidly becoming an increasingly important type of workflow with distributed data production and ML execution, and their need to couple to large data-centers.

Unsurprisingly, coupling ML to HPC simulations also introduces many challenges — at application, middleware and resource levels. In this chapter we focus on three main middleware development challenges related to resource and task execution management to realize the full potential of ML for scientific workflow applications: (1) task heterogeneity; (2) adaptive execution; and (3) application performance.

ML introduces multiple levels of task heterogeneity. Alongside traditional CPU, GPU and, possibly, multi-node MPI tasks, ML usually requires the execution of high-throughput function calls, often implemented in an interpreted language as Python, and that may depend on datasets distributed across repositories managed by diverse organization and platforms. As a consequence, the middleware that manages the execution of the workflow application, has to be able to concurrently schedule, place and execute MPI executables alongside Python functions with wildly varying execution lifetimes— the former for hours, the latter for as little as fractions of seconds.

One of the main scientific reasons to use ML in workflow applications is to improve the analysis that can be done on the data produced by part of the tasks of the workflow application. While some tasks progress, ML models can be used to learn relevant features and better drive the progress of the workflow at runtime. In order to leverage the potential of ML-based analysis, the workflow application has to become adaptive, i.e., being able to integrate the results of ML inferences and alter the workflow graph accordingly, and define the amount of learning to perform at runtime, especially when that amount cannot be known in advance, before execution [brace2021achieving]: simulations must be paused and restarted with new starting points, and/or a diverse number, type or size of simulations must be started to account for changed requirements, based on ML inferences. Further, ML training can vary at runtime, both in amount per model and across multiple models, when used. That has consequences for the capabilities of the workflow execution middleware. Alongside the capability of traversing an acyclic direct graph (DAG) to produce a concrete execution plan, workflow middleware has to update that DAG, pausing/restarting the execution of some of its nodes, adding/removing some nodes, and/or dynamically changing the amount of resources allocated to those nodes.

Finally, for ML to be useful it must enable improvements in both scientific and execution performance. On one side, the use of ML modeling and inference needs to improve the scientific computation that it drives, e.g., the accuracy and/or physically simulated duration. On the other side, ML has to effectively and efficiently use available resources when integrated within a workflow application. Resource efficiency depends on both the amount of time those resources are used in order to achieve the planned goal of the workflow application, and the percentage of available resources utilized to achieve that goal. This means that the workflow execution middleware has to manage the concurrent execution of heterogeneous tasks in a way that maximizes resource utilization while minimizing the workflow application total time to completion.

The ML-enabled workflow frameworks described earlier address some or all the challenges of task heterogeneity, adaptive execution and framework’s performance (as opposed to scientific performance), at different levels of the middleware software stack.

Proxima [zamora2021proxima] is implemented as a Python library used to wrap a Python function. Based on its inputs, Proxima calculates when to infer via a surrogate model or running the wrapped function. Inferring via a surrogate model is often faster than executing the wrapped function, resulting in an overall speedup. Proxima continually monitors the function execution, dynamically adapting the surrogate configuration parameters and determining when to retrain the surrogate model at runtime. While Proxima executes different types of functions (inference, monitoring, evaluation, configuration and retraining), it is not optimized for HPC and does not concurrently execute those different functions at scale. Proxima implements adaptivity, by retraining at runtime and parametrizing the model. Finally, Proxima performance as Python library is evaluated in terms of Proxima logic, model (re)training, surrogate usage, and inference.

Colmena [ward2021colmena] a general-purpose Python library for steering ensembles of experiments on HPC computing systems. Colmena is designed to execute different types of tasks, including: simulation, inference (via surrogate models) model training, and candidate generation. Different from Proxima, Colmena is designed to scale on HPC platforms, addressing the heterogeneity challenge by coordinating the (possibly concurrent) execution of different types of tasks. Similar to Proxima, Colmena enables adaptivity via surrogate parameterization and (re)training. Colmena uses Parsl as its runtime, avoiding a reimplementation of a ML-specific and general-purpose runtime capabilities. Colmena’s performance is measured in terms of communication overheads (e.g., requests or result object, and data input or output), and scaling performance with different task duration, result size, and number of workers.

EXARL [exarl] is a Python framework build on OpenAI Gym to enable the implementation of arbitrary reinforcement learning (RL) algorithms and their execution at scale. EXARL implements agents, each based on a learner/actors architecture in which each agent concurrently uses a scalable number of learners. Leaners can be implemented as multi-process or MPI executables; multiple agents can be executed concurrently. EXARL does not offer specific capabilities for mapping and launching its agents, relying on third party tools like, for example, batch system and an MPI infrastructure. As such, EXARL does not support task heterogeneity and implementing adaptivity requires coding capabilities on top of its agents. Performance is currently under evaluation, in terms of scalability of the size of each learners, and number of concurrent learners and agents.

MUMMI [mummi1-short, mummi2-short] is a Python workflow manager that coordinated the execution of massively parallel multiscale simulations. MUMMI allows to coordinate the concurrent execution of macro- and micro-scale simulation tasks, coupling them via ML methods to decide what space of the macro-scale simulations should be explored by the micro-scale ones. MUMMI uses the Flux job scheduler to coordinate the scheduling and execution of heterogeneous tasks on both CPUs and GPUs, and the Maestro workflow plugin to interface its workflow manager component to Flux. MUMMI enables adaptivity, allowing (re)training of ML models at runtime and using them for steering the simulations. MUMMI’s performance is evaluated in terms of resource utilization and number of concurrent simulations executed.

IMPECCABLE [saadi2020impeccable-short, lee2020scalable-short] is a drug discovery pipeline that executes heterogeneous tasks (i.e., MD simulations, ML training and inference) on both CPUs and GPUs at scale. Implemented using RADICAL-Cybertools as workflow middleware and runtime systems, it also uses DeepDriveMD. IMPECCABLE enables adaptivity by clustering MD trajectories to steer the ensemble of MD simulations. This may include either starting new simulations (i.e., expanding the pool of initial MD simulations), or killing unproductive MD simulations (i.e., simulations stuck in meta-stable states). IMPECCABLE also supports runtime evaluation of training of docking surrogate(s). IMPECCABLE’s performance is evaluated in terms of resource utilization, framework’s overheads, and total time to completion of the pipeline and each of its stages.

Importantly, the capabilities offered by DeepDriveMD and RADICAL-Cybertools are portable across use cases and computational campaigns. DeepDriveMD and RADICAL-Cybertools capabilities which are utilized for IMPECCABLE, also allowed for coordinating the diverse simulations coupled to ML models, and automate their execution at scale for the #COVIDIsAirborne [dommer2021covidisairborne] campaign. Work is underway to use RADICAL-Cybertools to support workflow orchestration, heterogeneous task execution and adaptivity at scale.