Memory management

[kkly1995]

Hi,

I would like to try MACE on my own bulk hydrogen dataset and ran out of memory on my first attempt. I am well aware that my setup is somewhat big (samples with up to N = 250, local environments with up to ~150 neighbors if the cutoff is 3.0) but it is not clear to me how much of this issue is due to my setup vs the implementation. Below I've pasted the input, log, and error messages. Any advice would be appreciated!

input:

python /home/ly1/mace/scripts/run_train.py \
	--name="MACE_model" \
	--train_file="training.xyz" \
	--valid_fraction=0.05 \
	--test_file="testing.xyz" \
	--config_type_weights='{"Default":1.0}' \
	--E0s="average" \
	--model="MACE" \
	--hidden_irreps='128x0e + 128x1o' \
	--r_max=3.0 \
	--batch_size=5 \
	--max_num_epochs=1000 \
	--swa \
	--start_swa=900 \
	--ema \
	--ema_decay=0.99 \
	--amsgrad \
	--restart_latest \
	--device=cuda

log:

2022-09-22 14:30:19.082 INFO: MACE version: 0.1.0
2022-09-22 14:30:19.082 INFO: Configuration: Namespace(name='MACE_model', seed=123, log_dir='logs', checkpoints_dir='checkpoints', results_dir='results', downloads_dir='downloads', device='cuda', default_dtype='float64', log_level='INFO', error_table='PerAtomRMSE', model='MACE', r_max=3.0, num_radial_basis=8, num_cutoff_basis=5, interaction='RealAgnosticResidualInteractionBlock', interaction_first='RealAgnosticResidualInteractionBlock', max_ell=3, correlation=3, num_interactions=2, MLP_irreps='16x0e', hidden_irreps='128x0e + 128x1o', gate='silu', scaling='rms_forces_scaling', avg_num_neighbors=1, compute_avg_num_neighbors=True, train_file='training.xyz', valid_file=None, valid_fraction=0.05, test_file='testing.xyz', E0s='average', energy_key='energy', forces_key='forces', loss='weighted', forces_weight=10.0, swa_forces_weight=1.0, energy_weight=1.0, swa_energy_weight=1000.0, config_type_weights='{"Default":1.0}', optimizer='adam', batch_size=5, valid_batch_size=10, lr=0.01, swa_lr=0.001, weight_decay=5e-07, amsgrad=True, scheduler='ReduceLROnPlateau', lr_factor=0.8, scheduler_patience=50, lr_scheduler_gamma=0.9993, swa=True, start_swa=900, ema=True, ema_decay=0.99, max_num_epochs=1000, patience=2048, eval_interval=2, keep_checkpoints=False, restart_latest=True, save_cpu=False, clip_grad=10.0)
2022-09-22 14:30:19.172 INFO: CUDA version: 11.2, CUDA device: 0
2022-09-22 14:30:29.908 INFO: Loaded 7905 training configurations from 'training.xyz'
2022-09-22 14:30:29.909 INFO: Using random 5.0% of training set for validation
2022-09-22 14:30:31.109 INFO: Loaded 877 test configurations from 'testing.xyz'
2022-09-22 14:30:31.109 INFO: Total number of configurations: train=7510, valid=395, tests=[Default: 877]
2022-09-22 14:30:31.405 INFO: AtomicNumberTable: (1,)
2022-09-22 14:30:31.406 INFO: Atomic Energies not in training file, using command line argument E0s
2022-09-22 14:30:31.406 INFO: Computing average Atomic Energies using least squares regression
2022-09-22 14:30:31.474 INFO: Atomic energies: [-12.852237399374355]
2022-09-22 15:33:01.340 INFO: WeightedEnergyForcesLoss(energy_weight=1.000, forces_weight=10.000)
2022-09-22 15:33:13.986 INFO: Average number of neighbors: 108.026
2022-09-22 15:33:13.987 INFO: Building model
2022-09-22 15:33:22.461 WARNING: Cannot find checkpoint with tag 'MACE_model_run-123' in 'checkpoints'
2022-09-22 15:33:22.561 INFO: Using stochastic weight averaging (after 900 epochs) with energy weight : 1000.0, forces weight : 1.0 and learning rate : 0.001
2022-09-22 15:33:22.563 INFO: ScaleShiftMACE(
  (node_embedding): LinearNodeEmbeddingBlock(
    (linear): Linear(1x0e -> 128x0e | 128 weights)
  )
  (radial_embedding): RadialEmbeddingBlock(
    (bessel_fn): BesselBasis(r_max=3.0, num_basis=8, trainable=False)
    (cutoff_fn): PolynomialCutoff(p=5.0, r_max=3.0)
  )
  (spherical_harmonics): SphericalHarmonics()
  (atomic_energies_fn): AtomicEnergiesBlock(energies=[-12.8522])
  (interactions): ModuleList(
    (0): RealAgnosticInteractionBlock(
      (linear_up): Linear(128x0e -> 128x0e | 16384 weights)
      (conv_tp): TensorProduct(128x0e x 1x0e+1x1o+1x2e+1x3o -> 128x0e+128x1o+128x2e+128x3o | 512 paths | 512 weights)
      (conv_tp_weights): FullyConnectedNet[8, 64, 64, 64, 512]
      (linear): Linear(128x0e+128x1o+128x2e+128x3o -> 128x0e+128x1o+128x2e+128x3o | 65536 weights)
      (skip_tp): FullyConnectedTensorProduct(128x0e+128x1o+128x2e+128x3o x 1x0e -> 128x0e+128x1o+128x2e+128x3o | 65536 paths | 65536 weights)
      (reshape): reshape_irreps()
    )
    (1): RealAgnosticResidualInteractionBlock(
      (linear_up): Linear(128x0e+128x1o -> 128x0e+128x1o | 32768 weights)
      (conv_tp): TensorProduct(128x0e+128x1o x 1x0e+1x1o+1x2e+1x3o -> 256x0e+384x1o+384x2e+256x3o | 1280 paths | 1280 weights)
      (conv_tp_weights): FullyConnectedNet[8, 64, 64, 64, 1280]
      (linear): Linear(256x0e+384x1o+384x2e+256x3o -> 128x0e+128x1o+128x2e+128x3o | 163840 weights)
      (skip_tp): FullyConnectedTensorProduct(128x0e+128x1o x 1x0e -> 128x0e | 16384 paths | 16384 weights)
      (reshape): reshape_irreps()
    )
  )
  (products): ModuleList(
    (0): EquivariantProductBasisBlock(
      (symmetric_contractions): SymmetricContraction(
        (contractions): ModuleDict(
          (128x0e): Contraction(
            (weights): ParameterDict(
                (1): Parameter containing: [torch.cuda.DoubleTensor of size 1x1x128 (GPU 0)]
                (2): Parameter containing: [torch.cuda.DoubleTensor of size 1x4x128 (GPU 0)]
                (3): Parameter containing: [torch.cuda.DoubleTensor of size 1x23x128 (GPU 0)]
            )
          )
          (128x1o): Contraction(
            (weights): ParameterDict(
                (1): Parameter containing: [torch.cuda.DoubleTensor of size 1x1x128 (GPU 0)]
                (2): Parameter containing: [torch.cuda.DoubleTensor of size 1x6x128 (GPU 0)]
                (3): Parameter containing: [torch.cuda.DoubleTensor of size 1x51x128 (GPU 0)]
            )
          )
        )
      )
      (linear): Linear(128x0e+128x1o -> 128x0e+128x1o | 32768 weights)
    )
    (1): EquivariantProductBasisBlock(
      (symmetric_contractions): SymmetricContraction(
        (contractions): ModuleDict(
          (128x0e): Contraction(
            (weights): ParameterDict(
                (1): Parameter containing: [torch.cuda.DoubleTensor of size 1x1x128 (GPU 0)]
                (2): Parameter containing: [torch.cuda.DoubleTensor of size 1x4x128 (GPU 0)]
                (3): Parameter containing: [torch.cuda.DoubleTensor of size 1x23x128 (GPU 0)]
            )
          )
        )
      )
      (linear): Linear(128x0e -> 128x0e | 16384 weights)
    )
  )
  (readouts): ModuleList(
    (0): LinearReadoutBlock(
      (linear): Linear(128x0e+128x1o -> 1x0e | 128 weights)
    )
    (1): NonLinearReadoutBlock(
      (linear_1): Linear(128x0e -> 16x0e | 2048 weights)
      (non_linearity): Activation [x] (16x0e -> 16x0e)
      (linear_2): Linear(16x0e -> 1x0e | 16 weights)
    )
  )
  (scale_shift): ScaleShiftBlock(scale=1.044676, shift=0.000000)
)
2022-09-22 15:33:22.576 INFO: Number of parameters: 558608
2022-09-22 15:33:22.576 INFO: Optimizer: Adam (
Parameter Group 0
    amsgrad: True
    betas: (0.9, 0.999)
    eps: 1e-08
    initial_lr: 0.01
    lr: 0.01
    name: embedding
    swa_lr: 0.001
    weight_decay: 0.0

Parameter Group 1
    amsgrad: True
    betas: (0.9, 0.999)
    eps: 1e-08
    initial_lr: 0.01
    lr: 0.01
    name: interactions_decay
    swa_lr: 0.001
    weight_decay: 5e-07

Parameter Group 2
    amsgrad: True
    betas: (0.9, 0.999)
    eps: 1e-08
    initial_lr: 0.01
    lr: 0.01
    name: interactions_no_decay
    swa_lr: 0.001
    weight_decay: 0.0

Parameter Group 3
    amsgrad: True
    betas: (0.9, 0.999)
    eps: 1e-08
    initial_lr: 0.01
    lr: 0.01
    name: products
    swa_lr: 0.001
    weight_decay: 5e-07

Parameter Group 4
    amsgrad: True
    betas: (0.9, 0.999)
    eps: 1e-08
    initial_lr: 0.01
    lr: 0.01
    name: readouts
    swa_lr: 0.001
    weight_decay: 0.0
)
2022-09-22 15:33:22.576 INFO: Using gradient clipping with tolerance=10.000
2022-09-22 15:33:22.576 INFO: Started training

error:

Traceback (most recent call last):
  File "/home/ly1/mace/scripts/run_train.py", line 375, in <module>
    main()
  File "/home/ly1/mace/scripts/run_train.py", line 318, in main
    tools.train(
  File "/home/ly1/.conda/envs/mace/lib/python3.9/site-packages/mace/tools/train.py", line 68, in train
    _, opt_metrics = take_step(
  File "/home/ly1/.conda/envs/mace/lib/python3.9/site-packages/mace/tools/train.py", line 174, in take_step
    output = model(batch, training=True)
  File "/home/ly1/.conda/envs/mace/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ly1/.conda/envs/mace/lib/python3.9/site-packages/mace/modules/models.py", line 221, in forward
    node_feats, sc = interaction(
  File "/home/ly1/.conda/envs/mace/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ly1/.conda/envs/mace/lib/python3.9/site-packages/mace/modules/blocks.py", line 516, in forward
    mji = self.conv_tp(
  File "/home/ly1/.conda/envs/mace/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ly1/.conda/envs/mace/lib/python3.9/site-packages/e3nn/o3/_tensor_product/_tensor_product.py", line 529, in forward
    return self._compiled_main_left_right(x, y, real_weight)
  File "/home/ly1/.conda/envs/mace/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "<eval_with_key>.67", line 115, in forward
    einsum_23 = torch.functional.einsum('edab,ecad->ecb', tensordot_3, einsum_22);  tensordot_3 = einsum_22 = None
    reshape_34 = einsum_23.reshape(getitem_4, 640);  einsum_23 = getitem_4 = None
    cat = torch.cat([reshape_12, reshape_26, reshape_15, reshape_24, reshape_30, reshape_18, reshape_28, reshape_34, reshape_21, reshape_32], dim = 1);  reshape_12 = reshape_26 = reshape_15 = reshape_24 = reshape_30 = reshape_18 = reshape_28 = reshape_34 = reshape_21 = reshape_32 = None
          ~~~~~~~~~ <--- HERE
    reshape_35 = cat.reshape(add_3);  cat = add_3 = None
    return reshape_35
RuntimeError: CUDA out of memory. Tried to allocate 3.87 GiB (GPU 0; 15.00 GiB total capacity; 10.04 GiB already allocated; 581.38 MiB free; 10.87 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

[gabor1]

I am wondering if it is ever the right regime for any system to have 150 neighbors in one layer. Maybe decrease the cutoff?

[davkovacs]

I think your GPU seems to have quite small memory. Few things to try are float32 and smaller batch size, maybe go to even 2 or 3. And of course you can also decrease the size of the model to 80x0e+80x1o.
I don’t think the number of neighbours should be the source of a memory overflow problem.

[kkly1995]

Apparently the V100 comes in two variants, 16GB and 32GB, and this computer happens to have the former. I have tried reducing r_max, yielding INFO: Average number of neighbors: 79.814, changing to float32, and decreasing batch_size to 3. This time the training actually proceeds, but then it runs out of memory at a different point:

Traceback (most recent call last):
  File "/home/ly1/mace/scripts/run_train.py", line 375, in <module>
    main()
  File "/home/ly1/mace/scripts/run_train.py", line 318, in main
    tools.train(
  File "/home/ly1/.conda/envs/mace/lib/python3.9/site-packages/mace/tools/train.py", line 85, in train
    valid_loss, eval_metrics = evaluate(
  File "/home/ly1/.conda/envs/mace/lib/python3.9/site-packages/mace/tools/train.py", line 207, in evaluate
    output = model(batch, training=False)
  File "/home/ly1/.conda/envs/mace/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ly1/.conda/envs/mace/lib/python3.9/site-packages/mace/modules/models.py", line 249, in forward
    "forces": compute_forces(
  File "/home/ly1/.conda/envs/mace/lib/python3.9/site-packages/mace/modules/utils.py", line 24, in compute_forces
    gradient = torch.autograd.grad(
  File "/home/ly1/.conda/envs/mace/lib/python3.9/site-packages/torch/autograd/__init__.py", line 234, in grad
    return Variable._execution_engine.run_backward(
RuntimeError: CUDA out of memory. Tried to allocate 1.20 GiB (GPU 0; 15.00 GiB total capacity; 9.45 GiB already allocated; 1.12 GiB free; 12.85 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

Here is a copy of the input from stdout:

2022-09-24 22:13:25.651 INFO: Configuration: Namespace(name='MACE_model', seed=123, log_dir='logs', checkpoints_dir='checkpoints', results_dir='results', downloads_dir='downloads', device='cuda', default_dtype='float32', log_level='INFO', error_table='PerAtomRMSE', model='MACE', r_max=2.7, num_radial_basis=8, num_cutoff_basis=5, interaction='RealAgnosticResidualInteractionBlock', interaction_first='RealAgnosticResidualInteractionBlock', max_ell=3, correlation=3, num_interactions=2, MLP_irreps='16x0e', hidden_irreps='128x0e + 128x1o', gate='silu', scaling='rms_forces_scaling', avg_num_neighbors=1, compute_avg_num_neighbors=True, train_file='training.xyz', valid_file=None, valid_fraction=0.05, test_file='testing.xyz', E0s='average', energy_key='energy', forces_key='forces', loss='weighted', forces_weight=10.0, swa_forces_weight=1.0, energy_weight=1.0, swa_energy_weight=1000.0, config_type_weights='{"Default":1.0}', optimizer='adam', batch_size=3, valid_batch_size=10, lr=0.01, swa_lr=0.001, weight_decay=5e-07, amsgrad=True, scheduler='ReduceLROnPlateau', lr_factor=0.8, scheduler_patience=50, lr_scheduler_gamma=0.9993, swa=True, start_swa=900, ema=True, ema_decay=0.99, max_num_epochs=1000, patience=2048, eval_interval=2, keep_checkpoints=False, restart_latest=True, save_cpu=False, clip_grad=10.0)

[davkovacs]

Try changing hidden_irreps='128x0e + 128x1o' to hidden_irreps='64x0e + 64x1o'

[kkly1995]

That works, thank you! I also tried keeping hidden_irreps as is, but reducing valid_batch_size to 3, just like batch_size, and that also worked.

By the way, is there an easy way to convert a checkpoint into a useable model?

[ilyes319]

To have a final model, you need to wait until the end of training after the test evaluation. To get a model from a checkpoint, you can use the hack by restarting a training for just one epochs and then it will create the model for you.