[aoti_torch: CUDA out of memory] Pathological memory allocation issue while attempting to use Allegro with LAMMPS

[jingyetan]

Hi, I’m running into what looks like a severe AOTI memory issue when using an Allegro checkpoint with the latest pair_nequip_allegro-patched LAMMPS on NERSC Perlmutter, and I wanted to ask whether this is expected, a known issue, or a likely compile/runtime incompatibility.

Setup

• Machine: NERSC Perlmutter
• GPUs tested:
  • 4 × A100 40 GB
  • 4 × A100 80 GB (-C "gpu&hbm80g")
• LAMMPS: custom build patched with the latest pair_nequip_allegro
• Pair style available in LAMMPS: allegro, allegro/kk, nequip, etc.
• Model: nequip.net:mir-group/Allegro-OAM-L:0.1
• Compiled artifact: AOTI only, since I need the fast runtime path
• Compilation command used:

nequip-compile nequip.net:mir-group/Allegro-OAM-L:0.1
./mir-group__Allegro-OAM-L__0.1.nequip.pt2
--mode aotinductor
--modifiers enable_TritonContracter
--device cuda
--target pair_allegro

Then I load that .nequip.pt2 inside LAMMPS through pair_style allegro/kk.

What works

• LAMMPS loads the model successfully.
• Type mapping is recognized correctly (Cu -> type 1).
• Kokkos multi-GPU setup is active and domain decomposition is happening correctly.
• EAM runs fine in the same general workflow, so basic Kokkos/MPI/GPU launching seems okay.

The problem

The run fails during VerletKokkos::setup() / first model execution with extremely large CUDA allocations from AOTI.

For example:

Case 1: 4 × A100 40 GB, 4 MPI ranks, full-ish test

• ~768k atoms total: OOM around 51.5 GiB requested per GPU
• ~560k atoms total: OOM around 37.6 GiB requested per GPU
• ~426k atoms total: still OOM

Case 2: 4 × A100 80 GB, 4 MPI ranks

Even on 80 GB GPUs, with only ~96k atoms total (~24k atoms/rank), each rank still drives memory usage to around 78 GB allocated by PyTorch and then fails trying to allocate another 12.9 GiB.

Representative error:

Exception in aoti_torch: CUDA out of memory. Tried to allocate 12.88 GiB.
GPU 0 has a total capacity of 79.25 GiB of which 552.62 MiB is free.
Including non-PyTorch memory, this process has 78.70 GiB memory in use.
Of the allocated memory 78.02 GiB is allocated by PyTorch ...

The stack trace consistently goes through:

• torch::inductor::AOTIModelContainerRunnerCuda::run_impl
• LAMMPS_NS::PairNequIPAllegro::call(...)
• LAMMPS_NS::PairAllegroKokkos::compute(int, int)
• LAMMPS_NS::VerletKokkos::setup(int)

and the failure happens inside generated AOTI wrapper code like:

... wrapper.cpp line 6842 / 6899
aoti_torch_empty_strided(...)

Why I think something is wrong beyond “the system is too big”

What seems odd to me is:

1. The memory blow-up happens already during setup / first call.
2. It persists even when the system is reduced substantially.
3. On 4 × 80 GB GPUs, a ~96k atom total system still reaches ~78 GB per GPU, which feels far beyond what I would expect for a sane domain-decomposed Allegro run.
4. The behavior seems more like a huge temporary/workspace allocation in the AOTI wrapper than normal steady-state MD memory.

Additional notes

• I need to stay with AOTI, not eager mode.
• Earlier, with NequIP, I had some OOM issues that could be worked around by avoiding compiled mode / using eager-style execution, but here I specifically want the AOTI fast path.
• I have also tested 4-rank / 4-GPU launching carefully; the multi-GPU mapping appears correct.

Questions

1. Is this a known issue with AOTI-compiled Allegro models in LAMMPS?
2. Is enable_TritonContracter possibly causing pathological memory planning here?
3. Is there a recommended nequip-compile recipe for pair_nequip_allegro that is known to be stable in LAMMPS?
4. Could this indicate a mismatch between:
   • checkpoint version,
   • nequip-compile version,
   • libtorch version,
   • and/or the patched pair_nequip_allegro interface?
5. Are there any recommended runtime environment variables or compile flags for reducing these AOTI temporary allocations while still staying in AOTI mode?

[cw-tan]

If you're using the pair_allegro pair style, which is required for multi-rank simulations with the Allegro model, the relevant --target is pair_allegro. See https://nequip.readthedocs.io/projects/allegro/en/latest/guide/lammps.html
e.g.

nequip-compile \
    path/to/ckpt_file/or/package_file \
    path/to/compiled_model.nequip.pt2 \
    --device [cpu|cuda] \
    --mode aotinductor \
    --target pair_allegro

Under the hood using --target pair_nequip and --target pair_allegro compiles for either pair style, so swapping them will probably just lead to nonsense behavior. The handling of ghost atoms is very different.

Note that pair_nequip does not support multi-rank simulations. In general, look at this summary table: https://nequip.readthedocs.io/en/latest/integrations/lammps/index.html#summary

The other option is ML-IAP, which will allow you to use cuEquivariance (https://nequip.readthedocs.io/projects/allegro/en/latest/guide/cuequivariance.html), which should be faster than the Triton TP. For NequIP models, ML-IAP will allow for multi-rank simulations, with OpenEquivariance accelerations.

[jingyetan]

Thanks Chuin Wei for getting back to me so quickly!

My current lammps build is patched with "pair_nequip_allegro" (https://github.com/mir-group/pair_nequip_allegro ), I made a mistake on my previous post, I just updated with this:

nequip-compile nequip.net:mir-group/Allegro-OAM-L:0.1
./mir-group__Allegro-OAM-L__0.1.nequip.pt2
--mode aotinductor
--modifiers enable_TritonContracter
--device cuda
--target pair_allegro

I had previously attempted to use nequip through ML-IAP for multi-rank MD simulation but I had no luck on that, do you have any pointer on this direction?

[cw-tan]

What was the problem with ML-IAP? This should generally work: https://nequip.readthedocs.io/en/latest/integrations/lammps/mliap.html

There might be cluster specific things to be aware of, but I just compiled everything fine recently.

[jingyetan]

Quick clarification, from what you said earlier, "For NequIP models, ML-IAP will allow for multi-rank simulations, with OpenEquivariance accelerations", were you referring to LAMMPS compiled with ML-IAP, and using nequip model (not allegro) with OpenEquivariance for multiple GPUs ranks? If so I could go about re-compiling and reproducing my previous errors to discuss more with you on that direction as well. However, this pathological OOM issue appeared in both the ML-IAP builds I had attempted previously, as well as the current build with pair_nequip_allegro. Do you think this error comes from compiling the allegro model artifacts? or from the actual software stack I compiled?

[cw-tan]

from what you said earlier, "For NequIP models, ML-IAP will allow for multi-rank simulations, with OpenEquivariance accelerations", were you referring to LAMMPS compiled with ML-IAP, and using nequip model (not allegro) with OpenEquivariance for multiple GPUs ranks?

yes

However, this pathological OOM issue appeared in both the ML-IAP builds I had attempted previously, as well as the current build with pair_nequip_allegro. Do you think this error comes from compiling the allegro model artifacts? or from the actual software stack I compiled?

I imagine that it could be that the large Allegro model is just heavy and OOMs more quickly. We've managed to pack 100k atoms with the NequIP large model on a single GPU while the Allegro large model OOMs.

[kavanase]

Just to add to @cw-tan's points;

• The Allegro-OAM-L model is still fairly large and memory intensive, so I think this is likely a case of the system being too big for the large model. In some (ASE) inference tests we ran, we had OOMs with this large Allegro model at >100k atoms (though I'm not sure what GPU setup was used for that, was single GPU anyway).
• Using cuEquivariance (as suggested by @cw-tan, with MLIAP) should significantly reduce the memory demand. Moreover, using TF32 if acceptable (often the case for MD simulations, but should probably be tested) will also significantly reduce the memory demand.
• Generally speaking, our NequIP-OAM models are faster and significantly lower memory demand for the same accuracy, as our Allegro-OAM models.
• We will also be releasing small and medium Allegro-OAM models shortly.

[jingyetan]

Hi Sean, thank you for joining in the discussion. Yes Allegro-OAM-L is definitely a bit large, and I am definitely very looking forward to the smaller model releases. However, the OOM issue I face here might not be from model size itself. for the same simulation, the more GPUs/memory resources I use (from 140GB to 440GB to 4*80GB etc), the OOM errors seem to just pick up all the available memories and still keep asking for more; on the other hand, for the same GPU resources, even as I scale the MD system size done as much as possible, the OOM issue persists and asking for the same if not more memory. I am not exactly sure how to properly describe this weird OOM situation, but I think something is not working properly in memory allocation here.

[AoyanLiang]

I am also seeing similar large GPU memory footprint with pair_style allegro.

System:

• ~1900 atoms
• A5000 GPU (24 GB)
• LAMMPS 30 Mar 2026
• Kokkos enabled
• Single GPU run (-k on g 1)

Model compilation:

nequip-compile \
  Allegro-OAM-L-0.1.nequip.zip \
  mir-group__Allegro-OAM-L__0.1.nequip.pt2 \
  --mode aotinductor \
  --modifiers enable_TritonContracter \
  --device cuda \
  --target pair_allegro

Observed behavior: The simulation OOMs during the initial minimization with only ~1900 atoms. The error reports ~23 GB already allocated by PyTorch and then fails on an additional ~1.2 GB allocation. To avoid OOM I have to decompose the system across multiple GPUs, but performance becomes worse due to communication overhead.

The same physical system runs successfully with a NequIP-OAM-L model through pair_style mliap unified on a single GPU and uses much less memory (only 1.2 GB).

[cw-tan]

Sorry for the late response @AoyanLiang , can you try using ML-IAP with Allegro and the same settings -- it is possible to do so. That will be one way to figure out if there's something intrinsic about the pair_allegro pair style, or if it's something intrinsic about the model. With ML-IAP, you'll also be able to use cuEquivariance with Allegro. Did you use the OpenEquivariance kernels with NequIP? That might explain part of it since those kernels would avoid materializing the largest (PyTorch) tensor --- the spherical tensor features on edges --- while Allegro has to materialize such objects even if it has a kernel.

[AoyanLiang]

Thanks for the suggestion. I tested Allegro-OAM-L-0.1 through ML-IAP on the same ~1900 atom system, and it still OOMs on a single A5000. So the large memory footprint does not seem specific to pair_style allegro. I also compiled/tested an Allegro ML-IAP model with cuEquivariance enabled, but it still OOMs on one A5000.

For comparison, NequIP-OAM-L-0.1 through ML-IAP runs on the same system on one A5000 and uses only about 1.2 GB. That NequIP run is using OpenEquivariance kernels.

This seems to suggest that the large memory footprint is more related to the Allegro-OAM-L model itself than to only the LAMMPS interface.

[kavanase]

@AoyanLiang this preprint might give a useful reference (though there are many issues with their speed benchmarking), which reports the raw memory demand of Allegro-OAM-L (without CuEquivariance or graph compilation it would seem):
https://arxiv.org/pdf/2607.07647