Large Memory usage - any tips to bring this down?

[ChrisDavi3s]

Hi all,

Benchmarking chgnet to see if it fits some of our use cases. Without diving too deep, we get out of memory errors (this was on a 16gb card, we’ve tried up to 40Gb A100’s) for certain systems and not for others:

  File “/###/lib/python3.9/site-packages/chgnet/model/layers.py", line 356, in forward
    new_angle_feas = self.twoBody_bond(total_fea)
  File “/###/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File “/###/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File “/###/lib/python3.9/site-packages/chgnet/model/functions.py", line 181, in forward
    gate = self.sigmoid(self.bn2(self.mlp_gate(x)))
  File “/###/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File “/###/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File “/###/lib/python3.9/site-packages/torch/nn/modules/activation.py", line 301, in forward
    return torch.sigmoid(input)
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 142.00 MiB. GPU 0 has a total capacity of 15.77 GiB of which 39.38 MiB is free. 

Trial systems used:

A Li argyrodite:

• A supercell
• 3328 atoms
• 16Gb ram usage

LiF + Iron nanoparticle:

• Created by inserting an iron nanoparticle in a matrix of iron
• 2500 atoms
• 45gb+ ram usage (crashes on the largest system we have access to!)

For comparison we’ve tried relaxing both systems in universal MACE (with ram usage barely going beyond 10gb for the model) and then rerunning chgnet to similar results.

I understand that chgnet has been successfully used for 10k atoms on an A100 - is there anything obvious we’re doing wrong?

Happy to provide more details on everything, and apologies if the information required / question is unhelpful