HydraGNN

Scalable PyTorch Implementation of Multi-Headed Graph Neural Networks

Capabilities

• Multi-headed Prediction for graph and node-level properties
• Distributed Training via DDP, FSDP (v1/v2), and DeepSpeed at supercomputing scale
• Convolutional Layers as a hyperparameter
• Geometric Equivariance in convolution and prediction (EGNN, PaiNN, PNAEq, MACE, DimeNet)
• Global Attention (GPS)
• Multiple Precision Training (FP32, BF16, FP64)
• Machine-Learned Interatomic Potentials with energy-conserving force prediction
• Gradient Checkpointing for memory-efficient training

Optional graph-level conditioning

• Enable with NeuralNetwork.Architecture.use_graph_attr_conditioning (off by default) and choose mode via graph_attr_conditioning_mode ("concat_node" default, "film", or "fuse_pool").
• concat_node (default) appends graph_attr to node embeddings and projects back to hidden dimension; FiLM scales/shifts invariant channels per graph; fuse_pool fuses graph_attr with the pooled graph embedding before the heads.
• Conditioning consumes data.graph_attr and requires those global attributes to be rotation/translation invariant; providing orientation-dependent values will break equivariance by design.

Dependencies

To install required packages with only basic capability (torch, torch_geometric, and related packages) and to serialize+store the processed data for later sessions (pickle5):

Python version: Installation via install_dependencies.sh is currently tested and supported for Python 3.10, 3.11, and 3.12 only.

Recommended approach - standard installation:

# Install all core dependencies (base + PyTorch + PyTorch Geometric)
pip install -r requirements.txt

# Or use the installation script
./install_dependencies.sh all

Alternative approach for reproducible installation:

# Use the provided installation script
./install_dependencies.sh

# Or install manually with consistent settings:
pip install --no-build-isolation -v -r requirements.txt

Modular installation (choose what you need):

# Base only (scientific computing, materials science, visualization)
pip install -r requirements-base.txt

# Add PyTorch
pip install -r requirements-torch.txt

# Add PyTorch Geometric  
pip install -r requirements-pyg.txt

# Add optional features (HPO, FAIRChem, etc.)
pip install -r requirements-optional.txt

If you plan to modify the code, include packages for formatting (black) and testing (pytest) the code:

pip install -r requirements-dev.txt
# Or with the script:
./install_dependencies.sh all dev

Detailed dependency installation instructions are available on the Wiki

Installation

After checking out HydgraGNN, we recommend to install HydraGNN in a developer mode so that you can use the files in your current location and update them if needed:

python -m pip install -e .

Or, simply type the following in the HydraGNN directory:

export PYTHONPATH=$PWD:$PYTHONPATH

Alternatively, if you have no plan to update, you can install HydraGNN in your python tree as a static package:

python -m pip install .

Quick Start

For detailed instructions, see the Comprehensive User Manual.

Below are the four main functionalities for running the code.

1. Training a model, including continuing from a previously trained model using configuration options:

import hydragnn
hydragnn.run_training("examples/configuration.json")

2. Saving a model state:

import hydragnn
model_name = model_checkpoint.pk
hydragnn.save_model(model, optimizer, model_name, path="./logs/")

3. Loading a model state:

import hydragnn
model_name = model_checkpoint.pk
hydragnn.load_existing_model(model, model_name, path="./logs/")

4. Making predictions from a previously trained model:

import hydragnn
hydragnn.run_prediction("examples/configuration.json", model)

The run_training and run_predictions functions are convenient routines that encapsulate all the steps of the training process (data generation, data pre-processing, training of HydraGNN models, and use of trained HydraGNN models for inference) on toy problems, which are included in the CI test workflows. Both run_training and run_predictions require a JSON input file for configurable options. The save_model and load_model functions store and retrieve model checkpoints for continued training and subsequent inference. Ad-hoc example scripts where data pre-processing, training, and inference are done for specific datasets are provided in the examples folder.

Datasets

Built in examples are provided for testing purposes only. One source of data to create HydraGNN surrogate predictions is DFT output on the OLCF Constellation: https://doi.ccs.ornl.gov/

Detailed instructions are available on the Wiki

Configurable settings

HydraGNN uses a JSON configuration file (examples in examples/):

There are many options for HydraGNN; the dataset and model type are particularly important:

• ["Verbosity"]["level"]: 0, 1, 2, 3, 4 (int)
• ["Dataset"]["name"]: CuAu_32atoms, FePt_32atoms, FeSi_1024atoms (str)

Additionally, many important arguments fall within the ["NeuralNetwork"] section:

• ["NeuralNetwork"]

  • ["Architecture"]

    • ["mpnn_type"]
      Accepted types: CGCNN, DimeNet, EGNN, GAT, GIN, MACE, MFC, PAINN, PNAEq, PNAPlus, PNA, SAGE, SchNet (str)
    • ["num_conv_layers"]
      Examples: 1, 2, 3, 4 ... (int)
    • ["output_heads"]
      Task types: node, graph (int)
    • ["global_attn_engine"] Accepted types: GPS, None
    • ["global_attn_type"] Accepted types: multihead
    • ["pe_dim"] Dimension of positional encodings (int)
    • ["global_attn_heads"] Examples: 1, 2, 3, 4 ... (int)
    • ["hidden_dim"]
      Dimension of node embeddings during convolution (int) - must be a multiple of "global_attn_heads" if "global_attn_engine" is not "None"
    • ["enable_interatomic_potential"]
      Enable MLIP mode with dynamic graph construction and energy-conserving force prediction (bool, default false)

  • ["Variables of Interest"]

    • ["input_node_features"]
      Indices from nodal data used as inputs (int)
    • ["output_index"]
      Indices from data used as targets (int)
    • ["type"]
      Either node or graph (string)
    • ["output_dim"]
      Dimensions of prediction tasks (list)

  • ["Training"]

    • ["num_epoch"]
      Examples: 75, 100, 250 (int)
    • ["batch_size"]
      Examples: 16, 32, 64 (int)
    • ["Optimizer"]["learning_rate"]
      Examples: 2e-3, 0.005 (float)
    • ["compute_grad_energy"]
      Use the gradient of energy to predict forces (bool)
    • ["precision"]
      Training precision: "fp32", "bf16", "fp64" (str, default "fp32")
    • ["conv_checkpointing"]
      Enable gradient checkpointing to reduce memory usage (bool, default false)

Citations

If you use this software, please cite both releases:

Original release: "HydraGNN: Distributed PyTorch implementation of multi-headed graph convolutional neural networks", Copyright ID#: 81929619 https://doi.org/10.11578/dc.20211019.2

Newest release: Lupo Pasini, Massimiliano, Choi, Jong Youl, Mehta, Kshitij, Zhang, Pei, Weaver, Rylie, Messerly, Richard, Chowdhury, Arindam, Raman, Adithya, & Aji, Ashwin M. (2026). HydraGNN v5.0. https://doi.org/10.11578/dc.20260512.1

Contributing

We encourage you to contribute to HydraGNN! Please check the guidelines on how to do so.

Documentation

• Quick Start: This README provides basic usage examples
• Comprehensive User Manual: Detailed guide covering data pre-processing, model construction, scalable data management, and training
• Wiki: Additional technical documentation and datasets