CUDA out of memory

[adam-norris]

I trained a model using l_max=3 and PBC=True on a relatively small dataset of 100 training samples using the example.yaml config file defaults. When I go to evaluate it on a test set of ~3000 training samples (molecule with 96 atoms), I'm getting the following "CUDA out of memory" error:

RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "<eval_with_key>.161", line 544, in forward
    einsum_119 = torch.functional.einsum('dbca,dbc->dba', tensordot_47, mul_53);  tensordot_47 = mul_53 = None
    reshape_166 = einsum_119.reshape(getitem_4, 224);  einsum_119 = getitem_4 = None
    cat = torch.cat([reshape_31, reshape_45, reshape_97, reshape_133, reshape_18, reshape_70, reshape_108, reshape_160, reshape_21, reshape_47, reshape_64, reshape_78, reshape_99, reshape_106, reshape_114, reshape_135, reshape_154, reshape_34, reshape_43, reshape_51, reshape_72, reshape_91, reshape_110, reshape_127, reshape_141, reshape_162, reshape_37, reshape_49, reshape_57, reshape_66, reshape_80, reshape_87, reshape_101, reshape_116, reshape_125, reshape_137, reshape_156, reshape_24, reshape_53, reshape_62, reshape_74, reshape_93, reshape_112, reshape_120, reshape_129, reshape_143, reshape_150, reshape_164, reshape_27, reshape_59, reshape_68, reshape_82, reshape_89, reshape_103, reshape_118, reshape_139, reshape_148, reshape_158, reshape_40, reshape_55, reshape_76, reshape_85, reshape_95, reshape_122, reshape_131, reshape_145, reshape_152, reshape_166], dim = 1);  reshape_31 = reshape_45 = reshape_97 = reshape_133 = reshape_18 = reshape_70 = reshape_108 = reshape_160 = reshape_21 = reshape_47 = reshape_64 = reshape_78 = reshape_99 = reshape_106 = reshape_114 = reshape_135 = reshape_154 = reshape_34 = reshape_43 = reshape_51 = reshape_72 = reshape_91 = reshape_110 = reshape_127 = reshape_141 = reshape_162 = reshape_37 = reshape_49 = reshape_57 = reshape_66 = reshape_80 = reshape_87 = reshape_101 = reshape_116 = reshape_125 = reshape_137 = reshape_156 = reshape_24 = reshape_53 = reshape_62 = reshape_74 = reshape_93 = reshape_112 = reshape_120 = reshape_129 = reshape_143 = reshape_150 = reshape_164 = reshape_27 = reshape_59 = reshape_68 = reshape_82 = reshape_89 = reshape_103 = reshape_118 = reshape_139 = reshape_148 = reshape_158 = reshape_40 = reshape_55 = reshape_76 = reshape_85 = reshape_95 = reshape_122 = reshape_131 = reshape_145 = reshape_152 = reshape_166 = None
          ~~~~~~~~~ <--- HERE
    reshape_167 = cat.reshape(add_3);  cat = add_3 = None
    return reshape_167
RuntimeError: CUDA out of memory. Tried to allocate 7.85 GiB (GPU 0; 39.59 GiB total capacity; 33.84 GiB already allocated; 2.50 GiB free; 35.22 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

I didn't have any trouble running out of memory using l_max=2 , however, it seems l_max=3 requires quite a bit more memory at inference time. I tried using a lower batch size but ran into the same time. For context, this was using a 1X NVIDIA A100-SXM4-40GB. Is this normal behavior to expect when using a high l_max value?? And are there any ways to get around the high memory requirements?

Thanks in advance!

[simonbatzner]

Hi @adam-norris, yes, higher l's are expected to have poor scaling. The scaling goes as m=2l+1 so for l_max=3, you'd have 1 + 3 + 5 + 7 = 16 elements per feature. So that means your device will have to store 16 floats in memory per atom per feature.

If you're sure you want/need l_max=3 features, then I'd recommend using a smaller number of features for them. So typically what works well is to use large number of features for low l's and small number of features for high l's. You'll have to use the advanced option for specifying that, given in full.yaml. Here's an example for an SE(3)-equivariant (turning parity off, usually doesn't add much) configuration with decaying numbers of features for higher l's.

# alternatively, the irreps of the features in various parts of the network can be specified directly:
# the following options use e3nn irreps notation
# either these four options, or the above three options, should be provided--- they cannot be mixed.
chemical_embedding_irreps_out: 128x0e                                              
feature_irreps_hidden: 64x0e + 32x1e + 16x2e + 8x3e
irreps_edge_sh: 0e + 1e + 2e + 3e                                                
conv_to_output_hidden_irreps_out: 16x0e                                         

The other thing is obviously to reduce the batch size (often 1 actually work really well if you're training on forces). Let us know if you have other questions.

[adam-norris]

Thanks for the quick reply and detailed explanation! Reducing the batch size did the trick : )

[Linux-cpp-lisp]

When I go to evaluate it

Does this mean you are using nequip-evaluate? nequip-evaluate uses its own batch size, taken from a command-line option (see nequip-evaluate --help), and the default is pretty big to give people decent speed by default. For a bigger model, like the l = 3 one you are using, though, you may need to set a lower batch size at the command line when you run nequip-evaluate.

[adam-norris]

Yes I meant the nequip-evaluate command. Thanks for the tip about setting the batch size in the command line. Previously, I was changing the batch size in the yaml file and that explains why I was still getting the "CUDA out of memory" error with a batch size of 1. Now, I have no issues

[Linux-cpp-lisp]

Great, glad this resolved your issue.