NequIP
Abstract
[0.19.1] Added docs page for the NequIP foundation potentials, with links to fine-tuning, compilation and the benchmarks reported in the paper, and citation instructions (CITATION.bib, README, docs) for using them TorchProfilerCallback for profiling training with torch.profiler NEQUIP_NODE_PARALLEL_COMPILE environment variable (1/true/yes/y) to serialize the first train-time compile within each node. Intended for large multi-node jobs, where every rank generating the same kernels at once against a shared cache directory on a parallel filesystem can crash or hang. Off by default. Fixed multi-rank checkpoint resume restoring tensors onto the GPU index recorded in the checkpoint rather than each rank's own device: SimpleDDPStrategy now moves the model to its device before restoring state (which also avoids openequivariance modules wedging when state is restored onto them beforehand), and loads model, optimizer, and EMA state directly onto that device excessive host memory use when resuming from a checkpoint: nequip-train held two full copies of the checkpoint in CPU memory at once (one to read hyperparameters and the run index, one from load_from_checkpoint); the first is now freed before the second load, and the second is loaded with map_location="cpu" https://github.com/mir-group/nequip/issues/612: Expected all tensors to be on the same device when loading a CUDA-built package on a CUDA-visible machine, seen as nequip-package update failing its own verification: e3nn's CodeGenMixin serializes generated fx.GraphModules as opaque nested pickles whose storages carry an absolute device tag and are unreachable by map_location, so the Wigner-3j buffers stayed on cuda:0 while the rest of the model loaded on CPU; the CPU-remapping context manager added for CPU-only hosts is now applied unconditionally silently NaN parameter gradients (and hence all-NaN metrics from the first optimization step) under train-time compile (compile_mode: compile) for frames with a zero-volume cell, e.g. non-periodic data, for which ASE supplies an all-zero cell: stress is virial / volume, and although it is unused by the loss, AOTAutograd bakes every output's backward into the joint graph and hands unused outputs a materialized zero tangent, so 0/0 produced a NaN that merged into the displacement gradient shared with forces and poisoned every parameter gradient (eager never schedules that branch, so it was unaffected); zero-volume frames now get NaN stress from a division that never sees a zero denominator nequip-compile failing on PyTorch 2.13 with RuntimeError: PyTorch is checking whether allow_tf32 is enabled for cuDNN ... for both --device cuda and --device cpu on machines with a CUDA-enabled build: 2.13 propagates torch.backends.fp32_precision to cuDNN, whose legacy allow_tf32 flag cannot encode "ieee" and raises when torch.export snapshots it; the legacy flag is now kept in sync, leaving all fp32_precision values (and hence numerics) untouched train-time compile (compile_mode: compile) crashing on PyTorch 2.12 with derivative for aten::silu_backward is not implemented: 2.12 keeps make_fx's decomposition table in a thread-local contextvar not seen by the autograd backward worker thread, so the silu_backward patch was skipped; tracing is now single-threaded on 2.12 to keep the table visible (no-op on other versions) silent wrong forces in the torch-sim integration for batched evaluation: NequIPTorchSimCalc used a cached system count, so a batch reusing the previous call's atomic numbers with a different number of systems (e.g. monatomic cells) used a stale layout; the system count now comes from the input state's cell batch dimension
Publication details
- Journal / source
- Zenodo (CERN European Organization for Nuclear Research)
- Publication date
- 2 September 2026