Conjugate gradient minimization fails with NequIP model despite low forces/energy error - possible compile error

[gavinwinter]

I followed the procedure outlined here for training nequip. The only difference was that I used a dataset for a periodic crystal, and unfortunately conjugate gradient minimization as run through LAMMPS with pair_nequip fails with the deployed model (which had a test error of 0.5kcal/mol/Angs for forces). The following is the output from the job. Do you think that these errors are related to a problem with how I compiled pair_nequip (it seemed successful, but it looks like default_program is giving compile errors)?

[d-9-16-2:73721] 1 more process has sent help message help-mpi-btl-openib.txt / no device params found
[d-9-16-2:73721] Set MCA parameter "orte_base_help_aggregate" to 0 to see all help / error messages
terminate called after throwing an instance of 'std::runtime_error'
  what():  The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
RuntimeError: default_program(18): error: expected a ")"

default_program(20): error: expected a ";"

default_program(22): error: expression must have class type

default_program(23): error: expression must have class type

default_program(24): error: expression must have class type

default_program(25): error: expression must have class type

default_program(27): error: expression must have class type

7 errors detected in the compilation of "default_program".

nvrtc compilation failed: 

#define NAN __int_as_float(0x7fffffff)
#define POS_INFINITY __int_as_float(0x7f800000)
#define NEG_INFINITY __int_as_float(0xff800000)


template<typename T>
__device__ T maximum(T a, T b) {
  return isnan(a) ? a : (a > b ? a : b);
}

template<typename T>
__device__ T minimum(T a, T b) {
  return isnan(a) ? a : (a < b ? a : b);
}

extern "C" __global__
void fused_mul_div_sin_div_mul_mul(float* t_, float* t__, float* aten_mul, float* aten_mul_1, float* aten_sin, float* aten_div, float* aten_mul_2, float* const_self.model.func.radial_basis.basis.bessel_weights) {
{
  float const_self.model.func.radial_basis.basis.bessel_weights_1 = __ldg(const_self.model.func.radial_basis.basis.bessel_weights + (512 * blockIdx.x + threadIdx.x) % 8);
  float t___1 = __ldg(t__ + (512 * blockIdx.x + threadIdx.x) / 8);
  aten_mul_2[512 * blockIdx.x + threadIdx.x] = const_self.model.func.radial_basis.basis.bessel_weights_1 * t___1;
  aten_div[512 * blockIdx.x + threadIdx.x] = (float)((double)(const_self.model.func.radial_basis.basis.bessel_weights_1 * t___1) / 4.0);
  aten_sin[512 * blockIdx.x + threadIdx.x] = sinf((float)((double)(const_self.model.func.radial_basis.basis.bessel_weights_1 * t___1) / 4.0));
  aten_mul_1[512 * blockIdx.x + threadIdx.x] = (float)((double)((sinf((float)((double)(const_self.model.func.radial_basis.basis.bessel_weights_1 * t___1) / 4.0))) / t___1) * 0.5);
  float v = __ldg(t_ + (512 * blockIdx.x + threadIdx.x) / 8);
  aten_mul[512 * blockIdx.x + threadIdx.x] = (float)((double)((sinf((float)((double)(const_self.model.func.radial_basis.basis.bessel_weights_1 * t___1) / 4.0))) / t___1) * 0.5) * v;
}
}


Aborted
/home/gridsan/gawinter/projects/labutil/src/plugins/lammps.py:96: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray.
  outrows = numpy.transpose(numpy.array(output))

The following is the lammps output, which is pretty much the same as in the toluene example here (except for the fact that my output does not have anything related to ghost atoms whereas the example output does).

--------------------------------------------------------------------------
WARNING: No preset parameters were found for the device that Open MPI
detected:

  Local host:            d-9-16-2
  Device name:           mlx5_1
  Device vendor ID:      0x02c9
  Device vendor part ID: 4117

Default device parameters will be used, which may result in lower
performance.  You can edit any of the files specified by the
btl_openib_device_param_files MCA parameter to set values for your
device.

NOTE: You can turn off this warning by setting the MCA parameter
      btl_openib_warn_no_device_params_found to 0.
--------------------------------------------------------------------------
--------------------------------------------------------------------------
No OpenFabrics connection schemes reported that they were able to be
used on a specific port.  As such, the openib BTL (OpenFabrics
support) will be disabled for this port.

  Local host:           d-9-16-2
  Local device:         mlx5_0
  Local port:           1
  CPCs attempted:       udcm
--------------------------------------------------------------------------
LAMMPS (29 Sep 2021 - Update 2)
OMP_NUM_THREADS environment is not set. Defaulting to 1 thread. (src/comm.cpp:98)
  using 1 OpenMP thread(s) per MPI task
Reading data file ...
  orthogonal box = (0.0000000 0.0000000 0.0000000) to (35.150584 35.150584 50.630184)
  1 by 1 by 1 MPI processor grid
  reading atoms ...
  3200 atoms
  read_data CPU = 0.007 seconds
NEQUIP is using device cuda
NequIP Coeff: type 1 is element Li
NequIP Coeff: type 2 is element Ge
NequIP Coeff: type 3 is element P
NequIP Coeff: type 4 is element S
Loading model from /home/gridsan/gawinter/projects/other_nff/results/lgps/example-lgps/lgps-deployed.pth
Freezing TorchScript model...
WARNING: Using 'neigh_modify every 1 delay 0 check yes' setting during minimization (src/min.cpp:188)
Neighbor list info ...
  update every 1 steps, delay 0 steps, check yes
  max neighbors/atom: 2000, page size: 100000
  master list distance cutoff = 5
  ghost atom cutoff = 5
  binsize = 2.5, bins = 15 15 21
  1 neighbor lists, perpetual/occasional/extra = 1 0 0
  (1) pair nequip, perpetual
      attributes: full, newton off
      pair build: full/bin/atomonly
      stencil: full/bin/3d
      bin: standard
Setting up cg style minimization ...
  Unit style    : real
  Current step  : 0
Per MPI rank memory allocation (min/avg/max) = 4.340 | 4.340 | 4.340 Mbytes
Step Temp E_pair E_mol TotEng Press 
       0            0    27554.828            0    27554.828            0 
       1            0    22022.125            0    22022.125            0 
[d-9-16-2:73719] *** Process received signal ***
[d-9-16-2:73719] Signal: Aborted (6)
[d-9-16-2:73719] Signal code:  (-6)
[d-9-16-2:73719] [ 0] /lib/x86_64-linux-gnu/libpthread.so.0(+0x12980)[0x7f968d053980]
[d-9-16-2:73719] [ 1] /lib/x86_64-linux-gnu/libc.so.6(gsignal+0xc7)[0x7f963c5e8fb7]
[d-9-16-2:73719] [ 2] /lib/x86_64-linux-gnu/libc.so.6(abort+0x141)[0x7f963c5ea921]
[d-9-16-2:73719] [ 3] /usr/lib/x86_64-linux-gnu/libstdc++.so.6(+0x8c957)[0x7f963cfdd957]
[d-9-16-2:73719] [ 4] /usr/lib/x86_64-linux-gnu/libstdc++.so.6(+0x92ae6)[0x7f963cfe3ae6]
[d-9-16-2:73719] [ 5] /usr/lib/x86_64-linux-gnu/libstdc++.so.6(+0x92b21)[0x7f963cfe3b21]
[d-9-16-2:73719] [ 6] /usr/lib/x86_64-linux-gnu/libstdc++.so.6(+0x92d54)[0x7f963cfe3d54]
[d-9-16-2:73719] [ 7] /home/gridsan/gawinter/.conda/envs/nff/lib/python3.7/site-packages/torch/lib/libtorch_cpu.so(+0x3558b2c)[0x7f9687097b2c]
[d-9-16-2:73719] [ 8] /home/gridsan/gawinter/.conda/envs/nff/lib/python3.7/site-packages/torch/lib/libtorch_cpu.so(+0x3562a26)[0x7f96870a1a26]
[d-9-16-2:73719] [ 9] /home/gridsan/gawinter/.conda/envs/nff/lib/python3.7/site-packages/torch/lib/libtorch_cpu.so(_ZN5torch3jit16InterpreterState3runERSt6vectorIN3c106IValueESaIS4_EE+0x1c)[0x7f96870926cc]
[d-9-16-2:73719] [10] /home/gridsan/gawinter/.conda/envs/nff/lib/python3.7/site-packages/torch/lib/libtorch_cpu.so(+0x354bc7b)[0x7f968708ac7b]
[d-9-16-2:73719] [11] /home/gridsan/gawinter/.conda/envs/nff/lib/python3.7/site-packages/torch/lib/libtorch_cpu.so(_ZN5torch3jit13GraphFunctionclESt6vectorIN3c106IValueESaIS4_EERKSt13unordered_mapISsS4_St4hashISsESt8equal_toISsESaISt4pairIKSsS4_EEE+0x39)[0x7f9686dfaf89]
[d-9-16-2:73719] [12] /home/gridsan/gawinter/.conda/envs/nff/lib/python3.7/site-packages/torch/lib/libtorch_cpu.so(_ZN5torch3jit6MethodclESt6vectorIN3c106IValueESaIS4_EERKSt13unordered_mapISsS4_St4hashISsESt8equal_toISsESaISt4pairIKSsS4_EEE+0x138)[0x7f9686e085a8]
[d-9-16-2:73719] [13] /home/gridsan/gawinter/projects/other_nff/lammps/build/lmp(+0x39b7ff)[0x561fd5f357ff]
[d-9-16-2:73719] [14] /home/gridsan/gawinter/projects/other_nff/lammps/build/lmp(+0xc7360)[0x561fd5c61360]
[d-9-16-2:73719] [15] /home/gridsan/gawinter/projects/other_nff/lammps/build/lmp(+0x36026d)[0x561fd5efa26d]
[d-9-16-2:73719] [16] /home/gridsan/gawinter/projects/other_nff/lammps/build/lmp(+0x35713b)[0x561fd5ef113b]
[d-9-16-2:73719] [17] /home/gridsan/gawinter/projects/other_nff/lammps/build/lmp(+0xc52be)[0x561fd5c5f2be]
[d-9-16-2:73719] [18] /home/gridsan/gawinter/projects/other_nff/lammps/build/lmp(+0xc946d)[0x561fd5c6346d]
[d-9-16-2:73719] [19] /home/gridsan/gawinter/projects/other_nff/lammps/build/lmp(+0x86ece)[0x561fd5c20ece]
[d-9-16-2:73719] [20] /home/gridsan/gawinter/projects/other_nff/lammps/build/lmp(+0x87605)[0x561fd5c21605]
[d-9-16-2:73719] [21] /home/gridsan/gawinter/projects/other_nff/lammps/build/lmp(+0x5f9aa)[0x561fd5bf99aa]
[d-9-16-2:73719] [22] /lib/x86_64-linux-gnu/libc.so.6(__libc_start_main+0xe7)[0x7f963c5cbbf7]
[d-9-16-2:73719] [23] /home/gridsan/gawinter/projects/other_nff/lammps/build/lmp(+0x5ffea)[0x561fd5bf9fea]
[d-9-16-2:73719] *** End of error message ***

I would assume the dataset is not the problem since I trained PaiNN (should be similar to NequIP) using this dataset and was able to successfully run MD with PaiNN. My .yaml for training nequip looks like this:

# !! PLEASE NOTE: `minimal.yaml` is meant as a _minimal_ example of a tiny, fast
#                 training that can be used to verify your nequip install,
#                 the syntax of your configuration edits, etc.
#                 These are NOT recommended hyperparameters for real applications!
#                 Please see `example.yaml` for a reasonable starting point.

# general
root: results/lgps
run_name: example-lgps
seed: 123
dataset_seed: 456

# network
r_max: 4.0                                                                        # cutoff radius in length units, here Angstrom, this is an important hyperparamter to scan
num_layers: 4                                                                     # number of interaction blocks, we find 3-5 to work best
l_max: 1                                                                          # the maximum irrep order (rotation order) for the network's features, l=1 is a good default, l=2 is more accurate but slower
parity: true                                                                      # whether to include features with odd mirror parity; often turning parity off gives equally good results but faster networks, so do consider this
num_features: 32                                                                  # the multiplicity of the features, 32 is a good default for accurate network, if you want to be more accurate, go larger, if you want to be faster, go lower
nonlinearity_type: gate                                                           # may be 'gate' or 'norm', 'gate' is recommended

# scalar nonlinearities to use — available options are silu, ssp (shifted softplus), tanh, and abs.
# Different nonlinearities are specified for e (even) and o (odd) parity;
# note that only tanh and abs are correct for o (odd parity)
# silu typically works best for even 
nonlinearity_scalars:
  e: silu
  o: tanh

nonlinearity_gates:
  e: silu
  o: tanh

# radial network basis
num_basis: 8                                                                      # number of basis functions used in the radial basis, 8 usually works best
BesselBasis_trainable: true                                                       # set true to train the bessel weights
PolynomialCutoff_p: 6                                                             # p-exponent used in polynomial cutoff function, smaller p corresponds to stronger decay with distance

# radial network
invariant_layers: 2                                                               # number of radial layers, usually 1-3 works best, smaller is faster
invariant_neurons: 64                                                             # number of hidden neurons in radial function, smaller is faster
avg_num_neighbors: auto                                                           # number of neighbors to divide by, null => no normalization, auto computes it based on dataset 
use_sc: true                                                                      # use self-connection or not, usually gives big improvement

# data set
# the keys used need to be stated at least once in key_mapping, npz_fixed_field_keys or npz_keys
# key_mapping is used to map the key in the npz file to the NequIP default values (see data/_key.py)
# all arrays are expected to have the shape of (nframe, natom, ?) except the fixed fields
# note that if your data set uses pbc, you need to also pass an array that maps to the nequip "pbc" key
dataset: npz                                                                       # type of data set, can be npz or ase
dataset_file_name: ./datasets/lgps_dft.npz                                        # path to data set file
key_mapping:
  pbc: pbc
  cell: cell                                                                       # cell
  z: atomic_numbers                                                                # atomic species, integers
  E: total_energy                                                                  # total potential eneriges to train to
  F: forces                                                                        # atomic forces to train to
  R: pos                                                                           # raw atomic positions
npz_fixed_field_keys:                                                              # fields that are repeated across different examples
  - atomic_numbers

chemical_symbols:
  - Li
  - Ge
  - P
  - S

# logging
wandb: false                                                                        # we recommend using wandb for logging
wandb_project: lgps                                                            # project name used in wandb


verbose: info                                                                      # the same as python logging, e.g. warning, info, debug, error; case insensitive
log_batch_freq: 10                                                                 # batch frequency, how often to print training errors withinin the same epoch
log_epoch_freq: 1                                                                  # epoch frequency, how often to print 
save_checkpoint_freq: -1                                                           # frequency to save the intermediate checkpoint. no saving of intermediate checkpoints when the value is not positive.
save_ema_checkpoint_freq: -1                                                       # frequency to save the intermediate ema checkpoint. no saving of intermediate checkpoints when the value is not positive.

# training
n_train: 100                                                                       # number of training data
n_val: 50                                                                          # number of validation data
learning_rate: 0.005                                                               # learning rate, we found values between 0.01 and 0.005 to work best - this is often one of the most important hyperparameters to tune
batch_size: 5                                                                      # batch size, we found it important to keep this small for most applications including forces (1-5); for energy-only training, higher batch sizes work better
max_epochs: 100                                                                    # stop training after _ number of epochs, we set a small number here to have an example that finished within a few minutes, but in practice we recommend using a very large number, as e.g. 1million and then to just use early stopping and not train the full number of epochs
train_val_split: random                                                            # can be random or sequential. if sequential, first n_train elements are training, next n_val are val, else random, usually random is the right choice
shuffle: true                                                                      # if true, the data loader will shuffle the data, usually a good idea
metrics_key: validation_loss                                                       # metrics used for scheduling and saving best model. Options: `set`_`quantity`, set can be either "train" or "validation, "quantity" can be loss or anything that appears in the validation batch step header, such as f_mae, f_rmse, e_mae, e_rmse
use_ema: true                                                                      # if true, use exponential moving average on weights for val/test, usually helps a lot with training, in particular for energy errors
ema_decay: 0.99                                                                    # ema weight, typically set to 0.99 or 0.999
ema_use_num_updates: true                                                          # whether to use number of updates when computing averages
report_init_validation: true                                                       # if True, report the validation error for just initialized model

# early stopping based on metrics values.
early_stopping_patiences:                                                          # stop early if a metric value stopped decreasing for n epochs
  validation_loss: 50

early_stopping_lower_bounds:                                                       # stop early if a metric value is lower than the bound
  LR: 1.0e-5

# loss function
loss_coeffs:                                                                        
  forces: 1                                                                        # if using PerAtomMSELoss, a default weight of 1:1 on each should work well
  total_energy:                                                                    
    - 1
    - PerAtomMSELoss

# output metrics
metrics_components:
  - - forces                               # key 
    - mae                                  # "rmse" or "mae"
  - - forces
    - rmse
  - - forces
    - mae
    - PerSpecies: True                     # if true, per species contribution is counted separately
      report_per_component: False          # if true, statistics on each component (i.e. fx, fy, fz) will be counted separately
  - - forces                                
    - rmse                                  
    - PerSpecies: True                     
      report_per_component: False    
  - - total_energy
    - mae    
  - - total_energy
    - mae
    - PerAtom: True                        # if true, energy is normalized by the number of atoms

# optimizer, may be any optimizer defined in torch.optim
# the name `optimizer_name`is case sensitive
optimizer_name: Adam                                                               # default optimizer is Adam 
optimizer_amsgrad: false

# lr scheduler, currently only supports the two options listed in full.yaml, i.e. on-pleteau and cosine annealing with warm restarts, if you need more please file an issue
# here: on-plateau, reduce lr by factory of lr_scheduler_factor if metrics_key hasn't improved for lr_scheduler_patience epoch
lr_scheduler_name: ReduceLROnPlateau
lr_scheduler_patience: 100
lr_scheduler_factor: 0.5

# we provide a series of options to shift and scale the data
# these are for advanced use and usually the defaults work very well
# the default is to scale the atomic energy and forces by scaling them by the force standard deviation and to shift the energy by the mean atomic energy
# in certain cases, it can be useful to have a trainable shift/scale and to also have species-dependent shifts/scales for each atom

# whether the shifts and scales are trainable. Defaults to False. Optional
per_species_rescale_shifts_trainable: false
per_species_rescale_scales_trainable: false

# initial atomic energy shift for each species. default to the mean of per atom energy. Optional
# the value can be a constant float value, an array for each species, or a string that defines a statistics over the training dataset
per_species_rescale_shifts: dataset_per_atom_total_energy_mean

# initial atomic energy scale for each species. Optional.
# the value can be a constant float value, an array for each species, or a string
per_species_rescale_scales: dataset_forces_rms

# if explicit numbers are given for the shifts/scales, this parameter must specify whether the given numbers are unitless shifts/scales or are in the units of the dataset. If ``True``, any global rescalings will correctly be applied to the per-species values.
# per_species_rescale_arguments_in_dataset_units: True

[simonbatzner]

Hi @gavinwinter can you share the LAMMPS input file you ran + Pytorch version + pair_nequip version? Thanks. It's odd that it fails on the 3rd step.

[simonbatzner]

Confirming that you're saying you saw the same error with pytorch 1.11 in a clean build? We have so far not observed this and it seems unlikely given that this was a bug in 1.9 that was fixed in 1.10.

[gavinwinter]

It was a separate dependency issue with PyTorch 1.11 that caused an issue different from the bug in 1.9. In one of these discussions or issues forums, your colleague Albert had mentioned an issue with 1.11, if I remember correctly. I tried to find that issue and link it but couldn’t.

[gavinwinter]

PyTorch 1.10 is all compatible and works

[Linux-cpp-lisp]

as far as I know PyTorch 1.11 (current stable) works on the Python side but doesn't work in LAMMPS yet in the stable releases

[Linux-cpp-lisp]

unless I already merged those fixes 👀 I will check