pair_nequip compilation

[atulcthakur]

Hi,

I trained and froze a model for my system and I'm trying to run it by installing the pair_nequip enabled lammps. I cloned the lammps and pair_nequip distributions as mentioned in the git README and built them with externally compiled CUDA, cuDNN and libtorch. Compilation and linking all goes perfectly fine but when trying to run MD with lammps, I get the following error:

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 (18.284000 12.740000 11.778000)
1 by 1 by 1 MPI processor grid
reading atoms ...
256 atoms
read_data CPU = 0.002 seconds
NEQUIP is using device cpu
NequIP Coeff: type 1 is element H
NequIP Coeff: type 2 is element C
NequIP Coeff: type 3 is element N
Loading model from frozen_models/aceamo.pth
Freezing TorchScript model...
Neighbor list info ...
update every 1 steps, delay 0 steps, check yes
max neighbors/atom: 2000, page size: 100000
master list distance cutoff = 8
ghost atom cutoff = 8
binsize = 4, bins = 5 4 3
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 Verlet run ...
Unit style : real
Current step : 0
Time step : 0.5
terminate called after throwing an instance of 'std::runtime_error'
what(): The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
File "code/torch/nequip/nn/_grad_output.py", line 30, in forward
_5 = torch.append(wrt_tensors, data[k])
func0 = self.func
data0 = (func0).forward(data, )
~~~~~~~~~~~~~~ <--- HERE
of = self.of
_6 = [torch.sum(data0[of])]
File "code/torch/nequip/nn/_graph_mixin.py", line 36, in forward
input1 = (spharm_edges).forward(input0, )
input2 = (radial_basis).forward(input1, )
input3 = (chemical_embedding).forward(input2, )
~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
input4 = (layer0_convnet).forward(input3, )
input5 = (layer1_convnet).forward(input4, )
File "code/torch/nequip/nn/_atomwise.py", line 13, in forward
linear = self.linear
field = self.field
_0 = (linear).forward(data[field], None, None, )
~~~~~~~~~~~~~~~ <--- HERE
out_field = self.out_field
torch._set_item(data, out_field, _0)
File "code/torch/e3nn/o3/_linear.py", line 54, in forward
bias0 = unchecked_cast(Tensor, bias)
_compiled_main = self._compiled_main
_5 = (_compiled_main).forward(features, weight0, bias0, )
~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
return _5
File "code/torch/torch/fx/graph_module.py", line 25, in forward
reshape_3 = torch.reshape(reshape_1, [3, 32])
mul = torch.mul(reshape_2, 0.57735026918962573)
tensordot = _0(mul, reshape_3, ([1], [0]), None, )
~~ <--- HERE
permute = torch.permute(tensordot, [0, 2, 1])
reshape_4 = torch.reshape(permute, [getitem_1, 32])
File "code/torch/torch/functional/___torch_mangle_0.py", line 9, in tensordot
dims_a, dims_b, = dims1
if torch.is(out, None):
_1 = torch.tensordot(a, b, dims_a, dims_b)
~~~~~~~~~~~~~~~ <--- HERE
_0 = _1
else:

Traceback of TorchScript, original code (most recent call last):
File "/projects/f_rcr78_1/PACKAGES/NequIP/nequipgpu/lib/python3.9/site-packages/nequip/nn/_grad_output.py", line 82, in forward
wrt_tensors.append(data[k])
# run func
data = self.func(data)
~~~~~~~~~ <--- HERE
# Get grads
grads = torch.autograd.grad(
File "/projects/f_rcr78_1/PACKAGES/NequIP/nequipgpu/lib/python3.9/site-packages/nequip/nn/_graph_mixin.py", line 354, in forward
def forward(self, input: AtomicDataDict.Type) -> AtomicDataDict.Type:
for module in self:
input = module(input)
~~~~~~ <--- HERE
return input
File "/projects/f_rcr78_1/PACKAGES/NequIP/nequipgpu/lib/python3.9/site-packages/nequip/nn/_atomwise.py", line 56, in forward
def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
data[self.out_field] = self.linear(data[self.field])
~~~~~~~~~~~ <--- HERE
return data
File "/projects/f_rcr78_1/PACKAGES/NequIP/nequipgpu/lib/python3.9/site-packages/e3nn/o3/_linear.py", line 280, in forward
raise RuntimeError("Biases must be provided when internal_weights = False")
bias = self.bias
return self._compiled_main(features, weight, bias)
~~~~~~~~~~~~~~~~~~~ <--- HERE
File "<eval_with_key>.6", line 15, in forward
reshape_3 = reshape_1.reshape((3, 32)); reshape_1 = None
mul = reshape_2 * 0.5773502691896257; reshape_2 = None
tensordot = torch.functional.tensordot(mul, reshape_3, ([1], [0])); mul = reshape_3 = None
~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
permute = tensordot.permute([0, 2, 1]); tensordot = None
reshape_4 = permute.reshape(getitem_1, 32); permute = getitem_1 = None
File "/projects/f_rcr78_1/PACKAGES/NequIP/nequipgpu/lib/python3.9/site-packages/torch/functional.py", line 1159, in tensordot
if out is None:
return _VF.tensordot(a, b, dims_a, dims_b) # type: ignore[attr-defined]
~~~~~~~~~~~~~ <--- HERE
else:
return _VF.tensordot(a, b, dims_a, dims_b, out=out) # type: ignore[attr-defined]
RuntimeError: expected scalar type Float but found Double
Aborted (core dumped)

I'm not entirely sure what's going on. I have also tried using pip installed pytorch libraries but I still get the same error. There's also no version mismatch as I'm using CUDA 11.3 and libtorch 1.10.0 as recommended.

I found some discussion on StackOverflow about this error. Maybe this will be relevant: https://stackoverflow.com/questions/56741087/how-to-fix-runtimeerror-expected-object-of-scalar-type-float-but-got-scalar-typ
For your information, my training settings are same as in the example.yaml file so the default data type is float32 for my model.

Thanks in advance.

[Linux-cpp-lisp]

so the default data type is float32 for my model.

hm... are you sure? I don't really see how this could happen otherwise. The current implementation of pair_nequip always inputs floats, so if it found doubles they have to be coming from inside the network. Can you use Python to check, load the TorchScript of your deployed model from nequip-deploy, iterate over its buffers() and parameters(), and see if any are Double (torch.float64)?

[atulcthakur]

Hi,

Thanks for your response. You're right. I rechecked the training file and I did set the model data type to float64.
Is there any way to use a model trained with datatype float64 with pair_nequip ?

If not, could I somehow cast the float32 dtype on my trained model?

[Linux-cpp-lisp]

Hi @atulcthakur ,

I've been meaning to add an option to do that, but in the meantime something this simple will do:

model = torch.load("mydeployed-float64.pth")
model = model.to(torch.float32)
torch.save("mydeployed-float32.pth")

Out of curiosity, was there a specific need that made you choose to use float64?

[atulcthakur]

Hi @Linux-cpp-lisp,

Thanks. Looks easy enough.
And no. I think I just set it to have more accurate representation of the underlying data: pytorch/pytorch#18427
I guess the default float32 would have worked perfectly fine too.

Thanks again for the help.

[Linux-cpp-lisp]

I see, yeah... we have never yet seen a problem we have attributed conclusively to float32 numerics, though of course I'm sure its possible. We are pretty careful in the NequIP code do to the scaling and shifting of the targets in such a way to keep the vast majority of the computations in reasonable magnitudes for float32, which helps.

[atulcthakur]

Certainly. I think my training time might have been slightly higher because of the float64 but I haven't done any benchmarking to support this. I agree that with all the scaling and shifting it makes much more sense to keep the datatype to float32.

I think it'd be nice to have some sort of exception handling to enforce float32 dtype until doubles are implemented or simple note in docs would be helpful too.

[Linux-cpp-lisp]

We can add a note to the docs, yes.

Re training/inference times, I'd expect float64 to be significantly slower (worse than 2x).