CUDA out of memory

[ken1944]

I have a model trained on 4000 frames on a 900 atoms surface-water system. After I have the model, I want to finetune it on a smaller dataset with 400 frames. While doing the training crashes with the OOM error:

Torch device: cuda^M
Loaded data: Batch(batch=[727200], edge_cell_shift=[878670204, 3], edge_index=[2, 878670204], forces=[727200, 3], pos=[727200, 3], ptr=[601], total_energy=[600, 1])^M
Cached processed data to disk^M
Successfully loaded the data set of type NpzDataset(600)...^M
Replace string dataset_forces_rms to 20.033674240112305^M
Replace string dataset_per_atom_total_energy_mean to -4.708922863006592^M
Atomic outputs are scaled by: [H, C, N, O, I, Pb: 20.033674], shifted by [H, C, N, O, I, Pb: -4.708923].^M
Replace string dataset_forces_rms to 20.033674240112305^M
Initially outputs are globally scaled by: 20.033674240112305, total_energy are globally shifted by None.^M
Successfully built the network...^M
Number of weights: 384728^M
! Starting training ...^M
Processing...^M
Done!^M
Traceback (most recent call last):^M
Traceback (most recent call last):^M
  File "/codes/miniconda/envs/nequip/bin/nequip-train", line 8, in <module>^M
    sys.exit(main())^M
  File "/codes/miniconda/envs/nequip/lib/python3.9/site-packages/nequip/scripts/train.py", line 71, in main^M
    trainer.train()^M
  File "j/codes/miniconda/envs/nequip/lib/python3.9/site-packages/nequip/train/trainer.py", line 764, in train^M
    self.epoch_step()^M
  File "/codes/miniconda/envs/nequip/lib/python3.9/site-packages/nequip/train/trainer.py", line 899, in epoch_step^M
    self.batch_step(^M
  File "/codes/miniconda/envs/nequip/lib/python3.9/site-packages/nequip/train/trainer.py", line 803, in batch_step^M
    out = self.model(input_data)^M
  File "/codes/miniconda/envs/nequip/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl^M
    return forward_call(*input, **kwargs)^M
  File "/codes/miniconda/envs/nequip/lib/python3.9/site-packages/nequip/nn/_rescale.py", line 140, in forward^M
    data = self.model(data)^M
  File "/codes/miniconda/envs/nequip/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl^M
    return Variable._execution_engine.run_backward(^M
RuntimeError: CUDA out of memory. Tried to allocate 1.75 GiB (GPU 0; 44.35 GiB total capacity; 41.57 GiB already allocated; 811.38 MiB free; 41.60 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^M

I am using l_max: 1 and r_max: 5.0 and num_layers: 6 with a batch size of 1. I'm not sure if I should change any other parameters to solve this error

[Linux-cpp-lisp]

Are the frames in the smaller dataset also 900 atoms? Are they significantly more dense? Are you running on the same GPU?

[ken1944]

Yes the number of atoms in the smaller dataset same. I just grabbed those frames from enhanced sampling of the same simulation. Yes, I am running it on the same GPU

[ken1944]

@Linux-cpp-lisp any suggestions to fix this?