[BUG] out of memory by se_e3

[QuantumMisaka]

Bug summary

Out of memory error when running deepmd by using se_e3 descriptors instead of se_e2_r

DeePMD-kit Version

2.2.2

TensorFlow Version

2.9.0

How did you download the software?

conda

Input Files, Running Commands, Error Log, etc.

Input file

{
 "model": {
  "descriptor": {
   "type": "se_e3",
   "sel": [80,120,40,40],
   "rcut_smth": 0.5,
   "rcut": 6.0,
   "neuron": [
    25,
    50,
    100
   ],
   "resnet_dt": false,
   "seed": 19090,
   "_activation_function": "tanh"
  },
  "fitting_net": {
   "neuron": [
    240,
    240,
    240
   ],
   "resnet_dt": true,
   "_coord_norm": true,
   "_type_fitting_net": false,
   "seed": 2375417769,
   "_activation_function": "tanh"
  },
  "type_map": [
   "Fe",
   "C",
   "H",
   "O"
  ]
 },
 "learning_rate": {
  "type": "exp",
  "start_lr": 0.001,
  "decay_steps": 80000,
  "stop_lr": 5e-08,
  "_decay_rate": 0.95
 },
 "loss": {
  "start_pref_e": 0.02,
  "limit_pref_e": 2,
  "start_pref_f": 1000,
  "limit_pref_f": 1,
  "start_pref_v": 0.01,
  "limit_pref_v": 1
 },
 "training": {
  "training_data": {
   "systems": "../dpdata_23265/",
   "batch_size": "auto"
  },
  "numb_steps": 16000000,
  "_comment": "that's all",
  "disp_file": "lcurve.out",
  "disp_freq": 2000,
  "numb_test": 1,
  "save_freq": 20000,
  "save_ckpt": "model.ckpt",
  "disp_training": true,
  "time_training": true,
  "profiling": false,
  "profiling_file": "timeline.json"
 }
}

running by

dp train input.json

wrong message:

2023-08-05 21:18:18.237808: W tensorflow/core/common_runtime/bfc_allocator.cc:479] Allocator (GPU_0_bfc) ran out of memory trying to allocate 1.27GiB (rounded to 1359360000)requested by op gradients_1/gradients/filter_type_all/concat_9_grad/Slice_grad/Pad
If the cause is memory fragmentation maybe the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' will improve the situation. 
Current allocation summary follows.
Current allocation summary follows.
2023-08-05 21:18:18.240210: W tensorflow/core/common_runtime/bfc_allocator.cc:491] *************_***************************************_**********************************************
2023-08-05 21:18:18.240278: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at pad_op.cc:130 : RESOURCE_EXHAUSTED: OOM when allocating tensor with shape[1699200,100] and type double on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
2023-08-05 21:18:28.240814: W tensorflow/core/common_runtime/bfc_allocator.cc:479] Allocator (GPU_0_bfc) ran out of memory trying to allocate 864.26MiB (rounded to 906240000)requested by op gradients_1/gradients/filter_type_all/concat_3_grad/Slice_grad/Pad
If the cause is memory fragmentation maybe the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' will improve the situation. 
Current allocation summary follows.
Current allocation summary follows.
2023-08-05 21:18:28.243087: W tensorflow/core/common_runtime/bfc_allocator.cc:491] *************_*****_********************************************************************************
2023-08-05 21:18:28.243140: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at pad_op.cc:130 : RESOURCE_EXHAUSTED: OOM when allocating tensor with shape[1132800,100] and type double on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
2023-08-05 21:18:38.244500: W tensorflow/core/common_runtime/bfc_allocator.cc:479] Allocator (GPU_0_bfc) ran out of memory trying to allocate 1.27GiB (rounded to 1359360000)requested by op gradients_1/gradients/filter_type_all/Tanh_14_grad/TanhGrad_grad/mul
If the cause is memory fragmentation maybe the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' will improve the situation. 
Current allocation summary follows.
Current allocation summary follows.
2023-08-05 21:18:38.246796: W tensorflow/core/common_runtime/bfc_allocator.cc:491] ********************************_********************_**********************************************
2023-08-05 21:18:38.246873: W tensorflow/core/framework/op_kernel.cc:1733] RESOURCE_EXHAUSTED: failed to allocate memory
Traceback (most recent call last):
  File "/home/liuzq/apps/anaconda3/envs/deepmd-2.2.2/lib/python3.10/site-packages/tensorflow/python/client/session.py", line 1377, in _do_call
    return fn(*args)
  File "/home/liuzq/apps/anaconda3/envs/deepmd-2.2.2/lib/python3.10/site-packages/tensorflow/python/client/session.py", line 1360, in _run_fn
    return self._call_tf_sessionrun(options, feed_dict, fetch_list,
  File "/home/liuzq/apps/anaconda3/envs/deepmd-2.2.2/lib/python3.10/site-packages/tensorflow/python/client/session.py", line 1453, in _call_tf_sessionrun
    return tf_session.TF_SessionRun_wrapper(self._session, options, feed_dict,
tensorflow.python.framework.errors_impl.ResourceExhaustedError: OOM when allocating tensor with shape[1699200,100] and type double on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
         [[{{node gradients_1/gradients/filter_type_all/concat_9_grad/Slice_grad/Pad}}]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info. This isn't available when running in Eager mode.

---- omitted -----

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/liuzq/apps/anaconda3/envs/deepmd-2.2.2/bin/dp", line 10, in <module>
    sys.exit(main())
  File "/home/liuzq/apps/anaconda3/envs/deepmd-2.2.2/lib/python3.10/site-packages/deepmd/entrypoints/main.py", line 623, in main
    train_dp(**dict_args)
  File "/home/liuzq/apps/anaconda3/envs/deepmd-2.2.2/lib/python3.10/site-packages/deepmd/entrypoints/train.py", line 164, in train
    _do_work(jdata, run_opt, is_compress)
  File "/home/liuzq/apps/anaconda3/envs/deepmd-2.2.2/lib/python3.10/site-packages/deepmd/entrypoints/train.py", line 278, in _do_work
    model.train(train_data, valid_data)
  File "/home/liuzq/apps/anaconda3/envs/deepmd-2.2.2/lib/python3.10/site-packages/deepmd/train/trainer.py", line 998, in train
    _, next_train_batch_list = run_sess(
  File "/home/liuzq/apps/anaconda3/envs/deepmd-2.2.2/lib/python3.10/site-packages/deepmd/utils/sess.py", line 49, in run_sess
    raise OutOfMemoryError(MESSAGE) from e
deepmd.utils.errors.OutOfMemoryError: Your memory may be not enough, thus an error has been raised above. You need to take the following actions:
1. Check if the network size of the model is too large.
2. Check if the batch size of training or testing is too large. You can set the training batch size to `auto`.
3. Check if the number of atoms is too large.
4. Check if another program is using the same GPU by execuating `nvidia-smi`. The usage of GPUs is controlled by `CUDA_VISIBLE_DEVICES` environment variable (current value: None).

Steps to Reproduce

dp train in A100-40G machine

Further Information, Files, and Links

same setting can be used in se_e2_r descriptors with 5000MiB memory

[wanghan-iapcm]

se_3 usually has a higher memory consumption than se_2_r. You may reduce the cutoff radius to reduce the memory requirement.