Parallelization in RuNNer 2¶

RuNNer features a hybrid MPI/OpenMP parallelization to make ideal use of the available computational resources. The way the resources are split among MPI tasks and OpenMP threads depends on the specific application. This guide gives an overview about the implementation and practical tips for getting the best performance out of RuNNer.

How to run parallelized calculations with MPI and OpenMP¶

OpenMP¶

OpenMP is based on the shared-memory model where threads that are working in parallel can access the same memory space. Here, repeating sections of the code that are independent of each other can be executed by multiple threads to accelerate the runtime. Whereever applicable, RuNNer 2 tries to make use of OpenMP parallelization (e.g. feature calculation, forward/backward passes, QEq). The parallelization is obtained by setting the OMP_NUM_THREADS environment variable to the number of cores that should be available to RuNNer 2. Since in OpenMP all threads need access to the same memory space, a parallelization across nodes by pure OpenMP is not possible.

Library dependencies¶

Various sections of the RuNNer 2 code rely on BLAS for linear algebraic operations (e.g. forward and backward passes, Kalman filter updates) and LAPACK for solving the linear system of equations in the direct charge equilibration (qeq_method direct). Depending on the respective BLAS and LAPACK implementation, the OpenMP parallelization is already provided (MKL, BLIS, AOCL) or must be built with the respective library (OpenBLAS + LAPACK).

Similarly, the Fast-Fourier-Transformation (FFT) which is used in the electrostatics and iterative charge equilibration is, depending on the respective FFT implementation, already OpenMP parallelized (MKL, AOCL) or must be built with OpenMP support (FFTW3).

MPI¶

MPI: MPI follows the distributed-memory model where each MPI task owns its assigned batch of data. Here, each MPI task executes large parts of the code independently of the other tasks, and most communication only happens with the root task at selected points. In RuNNer 2, the structures from input.data are evenly distributed among the MPI tasks (round-robin). Each task then performs the feature calculation and neural network predictions for its assigned block.

In order to run MPI-parallelized RuNNer 2 calculations, the RuNNer executable must be built with MPI support, see Installation. The calculation can be then launched via

mpirun -np X /path/to/your/RuNNer_mpi.x

with X specifying the number of MPI tasks.

Precompute mode¶

OpenMP¶

The calculation of the atomic features and feature derivatives are directly parallelized by OpenMP. The size of each parallel loop can be increased with the batchsize_structures keyword, increasing efficiency.

MPI¶

The structures in the dataset are distributed across MPI task, each task handling the calculation of features and feature derivatives separately. The results are communicated back to the root task for writing to the output files. Thus, the precompute mode scales ideally with the number of MPI tasks until the writing of the feature and feature derivatives to file, which is serial, becomes the bottleneck.

Predict mode¶

OpenMP¶

Most notably, the calculation of atomic features and feature derivatives, the forward and backward passes of the atomic neural networks, and the electrostatics (for 3G and 4G) and charge equilibration (for 4G) are parallelized via OpenMP. This covers most of the computationally intensive parts of the prediction, including force and stress evaluation. The parallel efficiency is controlled primarily by the batchsize_structures keyword with additional but smaller gains from the batchsize_elements keyword.

MPI¶

The structures in the data set are distributed across MPI tasks, and predictions are performed independently for their assigned structures. The predict mode scales perfectly with the number of MPI tasks as long as the number of tasks is smaller than the total number of structures in the data set, ensuring that each task owns at least one structure.

Train mode¶

OpenMP¶

Similar to predict mode, most of the prediction of the target properties is parallelized via OpenMP. In addition to that, the calculation of the gradient of the loss with respect to the neural network parameters and, if selected by optimizer (..) kalman[_std], the Kalman Filter updates are parallelized by OpenMP as well.

The parallel efficiency is controlled primarily by the batchsize_structures keyword with additional but smaller gains from the batchsize_elements keyword.

During training with the Kalman filter optimizer, the parameter updates often become the primary bottleneck. Since Kalman matrices are typically small, the OpenMP parallelization often only provides a decent speedup up to a small number of threads. In these cases, it is advisable to switch to MPI training instead.

MPI¶

Structures are distributed across the MPI tasks. Each task independently predicts the target property and computes the loss and its gradients. For the parameter update, all losses and gradients are sent to the root task, where they are averaged and used to update the model. This yields an effective batch size equal to the number of MPI tasks. The updated neural network parameters are then broadcast to all tasks.

This approach is termed "multistream training".

Interfaces¶

OpenMP¶

Same as in predict mode. When using RuNNer 2 as a library inside an MD executable, ensure that both the MD code and the RuNNer 2 library are linked against the same OpenMP runtime to avoid oversubscription or runtime errors.

MPI¶

Supported in combination with LAMMPS. LAMMPS decomposes the system into spatial subdomains processed independently by MPI tasks with the RuNNer 2 library evaluating the local energy/force/stress contributions on each task, which LAMMPS then communicates across the tasks. Parallel efficiency generally improves with larger systems (more atoms per subdomain). Exception: for 4G-HDNNP models the required global charge equilibration is performed on the root task, introducing a bottleneck that severely limits the parallelization by MPI.