DOCSEARCHbyAlgolia
Submit Docs

NequIP Documentation

repository·main·Indexed 21 days ago

https://github.com/mir-group/nequip

An open-source framework for building E(3)-equivariant interatomic potentials, optimized for high-performance training and inference. It features integrations with ASE and LAMMPS, support for multi-GPU training via DDP, and GPU kernel accelerations using OpenEquivariance and CuEquivariance. The framework allows for custom architectures via extension packages, such as allegro, and provides tools for LMDB dataset conversion and adaptive loss weighting.

Tokens
32.4K
Snippets
101
Records
179
Agent score
76%

What's inside NequIP

  1. Use NequIP DataModules for training

    main

    NequIP provides several specialized DataModule classes to handle different types of datasets. These classes are responsible for loading, preprocessing, and providing data to the model during training or evaluation.

    For general configuration guidance on how to set up your data, refer to the ../guide/configuration/data guide in the documentation.

    Available DataModule implementations include:

    • NequIPDataModule: General purpose data module.
    • ASEDataModule: For datasets using the Atomic Simulation Environment (ASE) format.
    • sGDML_CCSD_DataModule: Specialized for sGDML CCSD datasets.
    • rMD17DataModule: For MD17 datasets.
    • MD22DataModule: For MD22 datasets.
    • NequIP3BPADataModule: Specialized for 3BPA datasets.
    • COLLDataModule: For COLL datasets.
    • TM23DataModule: For TM23 datasets.
    • SAMD23DataModule: For SAMD23 datasets.
    • WaterDataModule: Specialized for water datasets.

  3. Understand NequIP file types

    main

    NequIP uses several specific file formats for configuration, training, distribution, and inference:

    • Config Files (.yaml): Describe training jobs for use with nequip-train.
    • Checkpoint Files (.ckpt): Produced during training. Used for restarting interrupted runs, fine-tuning models via nequip.model.ModelFromCheckpoint, or as source for compilation.
    • Package Files (.nequip.zip): Produced by nequip-package. These are archival formats containing both model weights and the necessary code, making them largely version-independent. They can be used for fine-tuning via nequip.model.ModelFromPackage or for compilation.
    • Compiled Model Files (.nequip.pth or .nequip.pt2): Produced by nequip-compile. These are used exclusively for inference with integrations like ASE or LAMMPS. The extension depends on the compilation mode (torchscript uses .pth, aotinductor uses .pt2).

  4. Navigate NequIP documentation

    main

    The NequIP documentation is organized into four primary sections to support different user needs:

    • User Guide: Covers getting started, configuration, training techniques, accelerations, and reference materials.
    • Python API: Detailed documentation of the underlying Python classes and functions. Note that most configuration file options correspond directly to these API elements.
    • Integrations: Instructions on how to use NequIP models with other software packages.
    • Developer Guide: Information for those looking to extend the framework or contribute to the codebase.

  5. Avoid checkpoint chaining and path issues when loading models

    main

    When using ModelFromCheckpoint or ModelFromPackage, the loaders save the paths provided to them into the resulting checkpoint file of any new training run. This creates a dependency chain.

    Critical Requirements

    • Use absolute paths: Always provide absolute paths instead of relative paths to avoid broken links if the working directory changes.
    • Do not move files: Moving or modifying the original checkpoint/package files will cause subsequent loading of the new checkpoint to fail.
    • Maintain directory structure: Do not change the directory structure of your model files.

    Breaking the Chain

    Iterated nested use of ModelFromCheckpoint leads to "checkpoint chaining," where loading the final checkpoint requires every intermediate file in the chain to be present and accessible.

    To break this chain, use nequip-package to convert a checkpoint file into a packaged model (.nequip.zip), and then load that package using nequip.model.ModelFromPackage.

  6. Handle validation and test metrics in DDP

    main

    When using DDP, torch.utils.data.distributed.DistributedSampler may duplicate data samples across devices if the dataset size is not evenly divisible by the number of ranks. This duplication can lead to incorrect validation or test metrics.

    To ensure accuracy:

    1. Ensure your dataset can be evenly distributed across all ranks.
    2. Alternatively, perform validation and testing on a single rank only.

  7. Scale batch size and learning rate for DDP

    main

    In NequIP, the batch_size defined in the data section of the config is the per-rank batch size.

    Effective Batch Size = per-rank batch size × number of ranks.

    When increasing the number of ranks, you must adjust hyperparameters (like learning rate) to account for the increased effective batch size. You can use omegaconf interpolation and NequIP's int_div resolver to dynamically calculate the per-rank batch size based on the global batch size and the number of tasks (e.g., from SLURM).

    batch_size: ${int_div:${effective_global_batch_size},${oc.env:SLURM_NTASKS}}

  8. Configure Loss and Metrics in NequIP

    main

    Loss functions and metrics are configured by specifying a data field (e.g., total_energy, forces) and an error quantity (e.g., MeanSquaredError, MeanAbsoluteError).

    • Loss functions determine what the model optimizes during training.
    • Metrics are used for monitoring progress and conditioning training behavior like early stopping or learning rate scheduling.

    Configuration is handled via MetricsManager objects within the training_module section of your configuration file.

  9. Implement custom data transforms

    main
    Data transforms are used to preprocess data during the loading process. To implement a custom transform, create a class that implements a __call__ method. This method must accept and return an AtomicDataDict object, allowing you to modify the data dictionary in place or return a modified version.

  10. Handle units in NequIP models

    main

    NequIP does not enforce a specific unit system; it uses the units provided in your dataset. You are responsible for ensuring consistency across all inputs and outputs.

    Example: If your length unit is Å and energy labels are in eV, the model's force predictions will be in eV/Å. Ensure your provided force labels are also in eV/Å.

  12. Handle stress sign conventions in NequIP

    main

    NequIP adopts the convention where stress = (-1/volume) * virial. This matches the ASE convention but differs from VASP. In this convention, positive diagonal stress entries imply the system is under tensile strain (wants to compress).

    If your dataset uses a different sign convention, you can:

    1. Preprocess the dataset manually.
    2. Use the nequip.data.transforms.StressSignFlipTransform data transform during training.
    from nequip.data.transforms import StressSignFlipTransform

# Use this transform to flip the sign of stress labels to match NequIP convention