PaperlayerView exact revision on arXiv

Source-linked AI summary

High-performance training and inference for deep equivariant interatomic potentials

Chuin Wei Tan, Marc L. Descoteaux, Mit Kotak, Gabriel de Miranda Nascimento, Seán R. Kavanagh, Laura Zichi, Menghang Wang, Aadit Saluja, Yizhong R. Hu, Tess Smidt, Anders Johansson, William C. Witt, Boris Kozinsky, Albert Musaelian

arXiv:2504.16068v1physics.comp-phcs.LGphysics.chem-ph

TL;DR

Growing datasets and downstream demands call for scalable MLIP software, while equivariant models pose computational challenges because their operations lack optimized library implementations. This work overhauls NequIP with compiled and distributed training, AOTI inference, and custom Allegro kernels; training and inference both show speedups over TorchScript.

Abstract

from arXiv · show

Machine learning interatomic potentials, particularly those based on deep equivariant neural networks, have demonstrated state-of-the-art accuracy and computational efficiency in atomistic modeling tasks like molecular dynamics and high-throughput screening. The size of datasets and demands of downstream workflows are growing rapidly, making robust and scalable software essential. This work presents a major overhaul of the NequIP framework focusing on multi-node parallelism, computational performance, and extensibility. The redesigned framework supports distributed training on large datasets and removes barriers preventing full utilization of the PyTorch 2.0 compiler at train time. We demonstrate this acceleration in a case study by training Allegro models on the SPICE 2 dataset of organic molecular systems. For inference, we introduce the first end-to-end infrastructure that uses the PyTorch Ahead-of-Time Inductor compiler for machine learning interatomic potentials. Additionally, we implement a custom kernel for the Allegro model's most expensive operation, the tensor product. Together, these advancements speed up molecular dynamics calculations on system sizes of practical relevance by up to a factor of 18.

I. INTRODUCTION

Large, diverse MLIP datasets and growing downstream demands make scalable, efficient, adaptable software infrastructure essential. This work revamps NequIP to improve performance and versatility for training and inference of deep equivariant models.

I. INTRODUCTION

  • Large and diverse datasets support more generalizable MLIPs, while growing computational demands make scalable, efficient, adaptable infrastructure essential.Frameworks must also support novel training paradigms, hyperparameter optimization, and diverse hardware.Source: §I · p. 1
  • The NequIP framework was revamped to improve computational performance and versatility across training and inference for deep equivariant networks.The framework was developed to support the NequIP and Allegro architectures.Source: §I · p. 1
  • Equivariant MLIPs incorporate physical symmetries into model structure, but their mathematical operations can lack optimized library implementations.More powerful compilers and custom optimized implementations motivate this work.Source: §I · p. 1

A. TorchInductor Compilation

TorchInductor offers more complex kernel fusion than TorchScript, but applying it to NequIP and Allegro requires compiling their energy derivatives. Its dynamic-shape support also suits atomistic simulations, where input sizes vary.

A. TorchInductor Compilation

  • TorchInductor transforms PyTorch operation graphs into device-specific code and can perform more complex kernel fusion than TorchScript.It uses Triton on GPUs and produces C++ and OpenMP code on CPUs.Source: A. TorchInductor Compilation · pp. 1–2
  • Tracing the full model into a torch.fx graph captures both energy prediction and backward-mode derivatives using primitive PyTorch operations, enabling train-time and ahead-of-time compilation.This addresses the challenge that derivative computation is not directly supported as a compiler primitive.Source: A. TorchInductor Compilation · p. 2
  • TorchInductor can generate one set of optimized kernels for dynamic tensor shapes, accommodating changes in atom and neighbor-pair counts during molecular dynamics.Many other neural network compilers enforce static tensor shapes.Source: A. TorchInductor Compilation · p. 2

B. Compiled and Distributed Training

The framework compiles the full model for training and adds distributed multi-GPU training through a custom DDP approach. Deferring communication until after backpropagation lets TorchInductor compile the entire backward pass.

B. Compiled and Distributed Training

  • Applying torch.compile to the primitive-operation graph generates optimized kernels for model predictions and train-time parameter-gradient computation.The procedure compiles both the forward pass and the backward pass for loss gradients.Source: B. Compiled and Distributed Training · p. 2
  • The framework adds multi-GPU training with a custom DDP approach that communicates after the full backward pass, allowing TorchInductor to compile that pass as one graph.Unlike standard DDP, this approach does not use gradient bucketing to overlap computation and communication.Source: B. Compiled and Distributed Training · p. 2
  • Omitting gradient bucketing has not imposed meaningful limitations on MLIP training, though bucketing could be reintroduced for larger models.The authors identify larger models as a possible reason to revisit the strategy.Source: B. Compiled and Distributed Training · p. 2

C. Ahead-of-Time Compilation for Inference

The framework introduces end-to-end AOTI compilation for MLIPs as a replacement for TorchScript export. Compiled models can run in non-Python codes such as LAMMPS and in Python-based codes such as ASE.

C. Ahead-of-Time Compilation for Inference

  • The prior LAMMPS integration used TorchScript exports and dedicated plugins to load and call compiled NequIP and Allegro models.The Allegro plugin also uses Kokkos to reduce interface overhead and enable GPU-aware MPI communication.Source: C. Ahead-of-Time Compilation for… · p. 2
  • The framework introduces the first end-to-end AOTI compilation for MLIPs as a replacement for TorchScript export.AOTI exports compiled PyTorch models as self-contained native code for non-Python environments.Source: C. Ahead-of-Time Compilation for… · p. 2
  • AOTI-compiled models can be used in high-performance codes such as LAMMPS without embedded Python interpreters, and remain usable in Python-based codes such as ASE.The approach avoids complicated, high-overhead embedded-Python solutions.Source: C. Ahead-of-Time Compilation for… · p. 2

D. Optimized Tensor Product Kernel

The Allegro tensor product remains computationally expensive, and TorchInductor cannot fuse its PyTorch implementation into one GPU kernel. Fused Triton kernels address this bottleneck for inference and integrate with AOTI.

D. Optimized Tensor Product Kernel

  • TorchInductor splits Allegro’s expensive tensor-product implementation into two kernels, requiring costly intermediate memory rather than fusing it into one GPU kernel.The tensor product is Allegro’s most computationally expensive operation.Source: D. Optimized Tensor Product Kernel · p. 2
  • Fused Triton kernels for the tensor product and its first derivative accelerate inference of energy and forces and integrate with AOTI; training kernels remain future work.AOTI allows these kernels to be exported for LAMMPS and other inference software.Source: D. Optimized Tensor Product Kernel · p. 3
  • The custom kernel uses a compressed sparse representation of Wigner 3-j contraction coefficients and avoids intermediate outer products, saving memory.Source: D. Optimized Tensor Product Kernel · p. 3

III. CASE STUDY: TRAINING AND INFERENCE OF SPICE 2 MODELS

The case study trains three Allegro models on SPICE 2, a large dataset of organic molecular systems. Training is restricted to neutral-total-charge systems, and the models vary in cost-accuracy trade-offs.

III. CASE STUDY: TRAINING AND INFERENCE OF SPICE 2 MODELS

  • SPICE 2 contains more than two million organic-system configurations, with energies and forces calculated using DFT and the ωB97M-D3(BJ) functional.Systems include drug-like molecules, amino acids, ions, water clusters, and their interactions.Source: §III · p. 3
  • To avoid multiple possible DFT labels for a configuration, the models were trained on SPICE 2 systems with neutral total charge.This is less restrictive than training on neutral per-atom formal-charge systems and more restrictive than using the full dataset.Source: §III · p. 3
  • The small, medium, and large Allegro models offer different cost-accuracy trade-offs through their rotation order, tensor-feature count, and layer count.All three use a radial cutoff rmax of 6.0 Å.Source: §III · p. 3, Table 1 · p. 3

A. Model Accuracy

On SPICE 2 test subsets, accuracy improves from the small to medium to large Allegro model, while energy errors are larger on the unrestricted set containing charge states absent from training. On TorsionNet 500, their barrier-height errors are comparable to those reported by prior studies.

A. Model Accuracy

  • The evaluation uses both the official SPICE 2 test set and TorsionNet 500, which assesses relative energy differences between molecular conformers.Source: A. Model Accuracy · p. 3
  • Accuracy improves from small to medium to large across system types and charge schemes; unrestricted-set energy errors are large, while force errors are only slightly larger than on filtered subsets.The unrestricted test set includes charge states absent from training; the test set contains systems distinct from training and spans four categories.Source: A. Model Accuracy · p. 3, Figure 1 · p. 4
  • Allegro’s TorsionNet 500 barrier-height errors are comparable to those of Kovács et al. [18] and Eastman et al..The table reports mean absolute error and root mean square error for barrier-height predictions.Source: Table 2 · p. 4, A. Model Accuracy · p. 4

B. Train-time Acceleration

Distributed Allegro training with torch.compile speeds up epochs relative to TorchScript and scales across multiple ranks. At 256 ranks, parallel efficiency remains 40% on AMD and 24% on NVIDIA systems.

B. Train-time Acceleration

  • The comparison trains the medium Allegro model on SPICE 2 on Frontier MI250X and Perlmutter A100 systems, including typical production-run overheads.Each MPI rank uses a local batch of eight atomic configurations.Source: B. Train-time Acceleration · p. 4, Figure 2 · p. 4
  • Distributed training scales well up to 128 ranks; at 256 ranks, torch.compile achieves parallel efficiencies of 40% on AMD and 24% on NVIDIA, relative to two-rank performance.Source: B. Train-time Acceleration · p. 5
  • torch.compile is 2.4–5.0 times faster than TorchScript across the tested MPI-rank range on MI250X and A100 (80GB) GPUs.Source: Figure 2 · p. 4, B. Train-time Acceleration · p. 5

C. Inference Acceleration

AOTI and an optimized tensor-product kernel improve Allegro inference throughput and memory efficiency across single-GPU and multi-node molecular dynamics benchmarks. The largest reported gains include 4–18× speedups over TorchScript and successful runs of systems that exceed TorchScript’s memory capacity.

Single-GPU inference

  • 4–18× speedups over TorchScript were observed for small-molecule and water simulations across tested hardware and models with AOTI plus the optimized tensor-product kernel.The advantage tends to be greater for larger Allegro models, whose tensor-product operations are more expensive.Source: C. Inference Acceleration · p. 5
  • AOTI plus the custom kernel raised the large model’s single-rank water limit from about 100 atoms with TorchScript to 4,320–5,184 atoms, depending on GPU.The reported maxima were 4,320 atoms on MI250X and A100 80GB, and 5,184 on H100.Source: C. Inference Acceleration · p. 5, Figure 4 · p. 6

Multi-GPU scaling

  • Strong scaling was reported up to 256 nodes on both biomolecular systems and platforms; AOTI throughput on Perlmutter began saturating between 256 and 512 nodes, while Frontier continued improving through 512.The tests used 23,558-atom DHFR and 408,609-atom cellulose systems.Source: C. Inference Acceleration · p. 5
  • 6.6× on Frontier and 6.9× on Perlmutter were the average speedups over TorchScript where the baseline fit in memory.The largest reported speedup was 10.9× for DHFR on 64 Perlmutter nodes; the maximum throughput was 205 timesteps per second for DHFR on 256 Perlmutter nodes.Source: C. Inference Acceleration · p. 6
  • On the 408,609-atom cellulose system, AOTI plus the custom kernel ran on 16 nodes, while TorchScript exceeded A100 40GB memory even at 512 nodes.The 16-node run corresponds to approximately 6,384 atoms per GPU.Source: C. Inference Acceleration · p. 6

IV. CONCLUSION

The redesigned NequIP framework combines distributed training, compiler technologies, and custom kernels to improve MLIP training and inference performance. The conclusion presents it as an extensible foundation, while identifying further evaluation of acceleration techniques as future work.

IV. CONCLUSION

  • The redesigned framework integrates PyTorch 2.0 compiler technologies, distributed training, and custom Triton kernels, with substantial training and inference performance improvements demonstrated using Allegro models trained on SPICE 2.The work trained three Allegro models and demonstrated support for training on large datasets.Source: §IV · p. 6
  • An in-depth study of the strengths and weaknesses of additional acceleration techniques remains future work.The framework supports per-edge-type cutoff radii, whose tuning can significantly affect speed because Allegro cost scales linearly with edge count; faster, less accurate TensorFloat32 was not enabled.Source: §IV · pp. 6–7
  • The updated framework’s modular structure supports extension packages such as Allegro and lets them leverage its acceleration techniques, tools, integrations, and infrastructure.The conclusion presents NequIP as an extensible foundation for MLIP architectures and training strategies.Source: §IV · p. 7

V. APPENDIX

The appendix describes Allegro’s mixed-precision sequence and records the software versions used in the work.

A. Mixed-Precision Architecture

  • Allegro embeds edge vectors largely in double precision, then uses single precision through its layers and readout before casting energy predictions back to double precision.The backward pass for forces and other derivatives follows the same precision transitions in reverse.Source: A. Mixed-Precision Architecture · p. 7

B. Code

  • The work used PyTorch 2.6.0, nequip 0.7.0, allegro 0.4.0, and pair nequip 0.7.0.The passage identifies the corresponding project repositories.Source: B. Code · p. 7
Loading 2504.16068v1…