CUDA Out of Memory with GCMC

[AlghamdiNada]

Hi, I am trying to run Grand Canonical Monte Carlo (GCMC) for water adsorption in silica using MACE. The issue is that I have a CUDA out of memory error after few thousands steps (~4000 with one gpu). I also tried using more than one gpu, but I still have cuda out of memory. I am using  A100 GPUs 64GB.

Input

This is the input file that I am using. Note that I specify bond and angle styles as harmonic but then the coefficient is 0 in order to be able to run the MC with the water molecule, so at the end the energies that I get are only mace predicted energies.

units metal
boundary p p p
atom_style full
neighbor 1.0 bin
neigh_modify delay 1
pair_style mace no_domain_decomposition

atom_modify   map yes
newton        on

bond_style harmonic
angle_style harmonic

read_data ../part1/1_SiOwithwater.data
molecule h2omol ../H2O.mol

lattice sc 3
create_atoms 0 box mol h2omol 45585

lattice none 1
group SiO type 1 2
group H2O type 3 4

pair_coeff * * ./MACE_MPtrj_2022.9.model-lammps.pt Si O H O
bond_coeff * 0.0 0.0
angle_coeff * 0.0 0.0

delete_atoms overlap 2 H2O SiO mol yes

# Next 4 lines to count the number of water molecules
variable oxygen atom "type==3"
group oxygen dynamic all var oxygen
variable nO equal count(oxygen)
fix myat1 all ave/time 100 10 1000 v_nO file numbermolecule.dat

## the GCMC step
variable tfac equal 5.0/3.0
variable xlo equal xlo+0.1
variable xhi equal xhi-0.1
variable ylo equal ylo+0.1
variable yhi equal yhi-0.1
variable zlo equal zlo+0.1
variable zhi equal zhi-0.1
region system block ${xlo} ${xhi} ${ylo} ${yhi} ${zlo} ${zhi}
fix fgcmc H2O gcmc 100 100 0 0 65899 300 -0.5 0.1 &
    mol h2omol tfac_insert ${tfac} group H2O &
    full_energy pressure 10000 region system

run 45000
write_data SiOwithwater.data
write_dump all atom dump.lammpstrj

Running environment

I used the following specifications to build lammps-mace:

cmake \
    -D CMAKE_BUILD_TYPE=Release \
    -D CMAKE_INSTALL_PREFIX=$(pwd) \
    -D CMAKE_CXX_STANDARD=17 \
    -D CMAKE_CXX_STANDARD_REQUIRED=ON \
    -D BUILD_MPI=ON \
    -D BUILD_OMP=ON \
    -D BUILD_SHARED_LIBS=ON \
    -D PKG_KOKKOS=ON \
    -D Kokkos_ENABLE_CUDA=ON \
    -D CMAKE_CXX_COMPILER=$(pwd)/../lib/kokkos/bin/nvcc_wrapper \
    -D Kokkos_ARCH_AMDAVX=ON \
    -D Kokkos_ARCH_AMPERE100=ON \
    -D CMAKE_PREFIX_PATH=$(pwd)/../../libtorch-gpu \
    -D PKG_ML-MACE=ON \
    -D CAFFE2_USE_CUDNN=True \
    -D PKG_RIGID=ON \
    -D PKG_MC=ON \
    -D PKG_MOLECULE=ON \
    ../cmake

When I run, I load these modules:

module load gcc/12.2.0
module load gsl/2.7.1--gcc--12.2.0
module load openmpi/4.1.6--gcc--12.2.0
module load fftw/3.3.10--openmpi--4.1.6--gcc--12.2.0
module load openblas/0.3.24--gcc--12.2.0
module load cuda/12.1
module load intel-oneapi-mkl/2023.2.0

Error message

This is the error message that I get

RuntimeError: CUDA out of memory. Tried to allocate 6.46 GiB. GPU 0 has a total capacity of 63.42 GiB of which 6.42 GiB is free. Including non-PyTorch memory, this process has 57.00 GiB memory in use. Of the allocated memory 55.10 GiB is allocated by PyTorch, and 302.98 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

I am not sure why I have cuda out of memory error although the size of my system is small, I would appreciate any insights, thanks!

[ilyes319]

Hey, what is your model size, and your cutoff? This greatly influence the number of atoms that can fit on a GPU.

[gabor1]

Could you please repeat your run but without changing the number of atoms in the system (i.e. always reject the GCMC steps which change the number of atoms). Let's see if the problem is with the changing atom number.

[jmargraf]

I've encountered a similar problem with a atom-swapping MD/MC code in ase. When the order of elements gets modified between MD steps, the GPU memory starts to slowly but steadily fill up, until it runs out. If I just do sequential short MDs without swapping atoms this doesn't happen.

[ilyes319]

That is interesting, I cant really explain it. Can you estimate the rate of increase in Mb/steps/atoms?

[jmargraf]

Here's the GPU memory usage over time for a MD simulation of an oxide with 448 atoms in the super cell, with and without swaps:

MD time fs	swaps MiB	no swaps MiB
0	4472	4414
100	4536	4414
200	4536	4414
300	5256	4414
400	5256	4414
500	5322	4414
600	5344	4414
700	5344	4414
800	6052	4414
900	6074	4414
1000	6074	4414
1100	6124	4414
1200	6822	4414
1300	6822	4414

So the increase is not linear, I'm assuming it depends on the number of successful swap attempts?

Some further info: I'm using the mace_agnesi_medium.model with float32 precision. When I do a swap, I make a copy of the atoms object, then modify the atom order and assign a copy of the mace calculator. I've also tried other variants, including reloading the MACE calculator from scratch in regular intervals, but the memory keeps building up. So it looks to me like there is some issue with memory that is no longer needed not being deallocated properly?

[ilyes319]

I see, thank you for the additional infos. I am wondering why reloading from scratch does not erase this problem. There must be a weird caching happening indeed. Can you try adding the following line in between swaps: torch.cuda.empty_cache().

[jmargraf]

I've tried cuda.empty_cache() to no avail, but I will investigate some more.

[wcwitt]

If @jmargraf is having the same problem, it's (hopefully) not a LAMMPS issue.

[jmargraf]

Yes, perhaps something more general about simulations where the number and order of atoms is not constant. Or my implementation is dumb, and it's an unrelated issue ;-)

[jmargraf]

Some more insight: The problem seems to come from making copies of the mace calculator (or deepcopies of atoms objects, which amounts to the same thing). I now modified the code to swap atom positions instead of modifying the element order, and changed everything to operate on a single atoms object in place. This seems to solve the issue.

I'm not sure how to deal with this more generally though, when I do want to make copies of atoms objects.

[bernstei]

The ASE atoms has a reference to the calculator, and also vice versa, so I'm not so surprised it can go wrong. I believe that the default behavior of ASE is to wipe the calculator reference when making a copy of an atoms object. You then need to reattach a calculator to the new object. But, especially if you abandon the original object, I don't see why you'd need to copy the calculator anyway. What's wrong with

atoms.calc = calc
.
.
.
atoms_new = atoms.copy() # atoms_new has no calculator
modify(atoms_new)
atoms_new.calc = calc # or maybe = atoms.calc
.
.
.

[jmargraf]

Yes, that's true. The issue might have been that I've been using deepcopy(atoms) instead of atoms.copy(). The problem with the latter is that it loses information about energies, forces, and (I think) momenta, which are nice to have in the trajectory.

[bernstei]

Yet another reason why the ASE atoms/calculator model for how results are stored is broken. Copy results from the calculator into info/arrays properties (ase.io.extxyz.save_calc_results) , use a file format that stores them, and everything will be better.

[jmargraf]

I didn't know about this function, thanks!

[bernstei]

It's new-ish, and comes out of our conflict with Ask about the whole extxyz/ASE per-config/atom information storage issue. It's on our list of ASE design decisions to finally hash out, perhaps at the proposed upcoming ASE CECAM workshop.