RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

[isabel-s-rosa]

Hey all!

I'm trying to call forward in C++ on a nequip model (stored in a .pth file, acquired by running nequip-train). The call to forward does not terminate, however, and prints this error:

libc++abi: terminating with uncaught exception of type std::runtime_error: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/nequip/nn/_rescale/___torch_mangle_173.py", line 2859, in forward
    _610 = torch.scatter_add_(out6, 0, index3, _603)
    torch._set_item(data, "total_energy", _610)
    grads = torch.autograd.grad([torch.sum(data["total_energy"])], wrt_tensors)
            ~~~~~~~~~~~~~~~~~~~ <--- HERE
    for _611 in range(ops.prim.min([1, torch.len(grads)])):
      out7 = (["forces"])[_611]

Traceback of TorchScript, original code (most recent call last):
  File "/home/gridsan/isrosa/anaconda3/envs/nequip/lib/python3.7/site-packages/nequip/nn/_grad_output.py", line 77, in forward
        data = self.func(data)
        # Get grads
        grads = torch.autograd.grad(
                ~~~~~~~~~~~~~~~~~~~ <--- HERE
            # TODO:
            # This makes sense for scalar batch-level or batch-wise outputs, specifically because d(sum(batches))/d wrt = sum(d batch / d wrt) = d my_batch / d wrt
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

Has anyone seen this error before? Any clue as to what this means or how to solve it would be greatly appreciated:) thanks!

[simonbatzner]

Hi @isabel-s-rosa, would you mind sharing in what context you are doing this? Is this a "regular" NequIP model (energies, forces, ...) or is this something custom? Could you share a bit more of how you are loading + calling the model?

The error appears when you call the gradient in PyTorch w.r.t. a tensor that does not have requires_grad=True. This would not happen in our default positions/energy/force setup if you use it for inference, molecular dynamics, etc. hence my asking.

[isabel-s-rosa]

Thanks for the quick response! It should be a regular model. I've attached the yaml file so you can take a look to confirm. Here's the code for how I load the model:

    torch::jit::script::Module module;
    try {
      // Deserialize the ScriptModule from a file using torch::jit::load().
      module = torch::jit::load(argv[1]);
    }
    catch (const c10::Error& e) {
      std::cerr << "error loading the model\n";
      return -1;
    }

Then, I create a dictionary (called inputs) mapping strings to tensors and call the model as follows: at::Tensor output = module.forward(inputs).toTensor();.

Let me know if you have any ideas! Thanks again.

[Linux-cpp-lisp]

Hi @isabel-s-rosa ,

A few things in no particular order:

1. full.yaml is meant as a reference, not a starting point. We recommend starting with minimal.yaml and adding options in as you need them. This makes your setup clearer.
2. at::Tensor output = module.forward(inputs).toTensor(); will not work, since the output of the model is Dict[str, torch.Tensor]. An example of how to process that can be seen in our official pair_nequip LAMMPS plugin here: https://github.com/mir-group/pair_nequip/blob/main/pair_nequip.cpp#L424
3. Are you running the model under no_grad or inference_mode? While these are recommended for inference for normal PyTorch models, we require the autograd engine to compute forces as the gradient of energy.
4. Does this model only fail in C++, but not in training or Python inference?
5. What is your usecase? While writing up the C++ inference yourself is certainly possible, we strongly recommend the use of our offical LAMMPS plugin pair_nequip for C++ inference/simulations. If your only goal is to evaluate the model on some fixed dataset, we strongly recommend our nequip-evaluate command. Both of these options take care of some of the subtleties and pain points of using the model correctly in a transparent way--- in the case where you really do need to write your own C++ inference code, I would recommend looking through the model loading and setup in pair_nequip to make sure that your code is setting global JIT and TensorFloat32 settings correctly.

[Linux-cpp-lisp]

This also may be relevant:

stored in a .pth file, acquired by running nequip-train

Do you run nequip-deploy first? This is the only way we support going from a training session to a deployed .pth TorchScript model. While there is (or at least was) a compile_model option that saves one during training, but that option has been recently removed.

[isabel-s-rosa]

After deploying, I see this error:

libc++abi: terminating with uncaught exception of type std::runtime_error: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/nequip/nn/_rescale.py", line 19, in forward
    data: Dict[str, Tensor]) -> Dict[str, Tensor]:
    model = self.model
    data0 = (model).forward(data, )
             ~~~~~~~~~~~~~~ <--- HERE
    training = self.training
    if training:
  File "code/__torch__/nequip/nn/_grad_output.py", line 28, in forward
    _5 = [torch.sum(data0[of])]
    training = self.training
    grads = torch.autograd.grad(_5, wrt_tensors, None, None, training)
            ~~~~~~~~~~~~~~~~~~~ <--- HERE
    out_field = self.out_field
    _6 = [torch.len(out_field), torch.len(grads)]

Traceback of TorchScript, original code (most recent call last):
  File "/Users/isabelrosa/Library/Python/3.8/lib/python/site-packages/nequip/nn/_rescale.py", line 124, in forward
    def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
        data = self.model(data)
               ~~~~~~~~~~ <--- HERE
        if self.training:
            return data
  File "/Users/isabelrosa/Library/Python/3.8/lib/python/site-packages/nequip/nn/_grad_output.py", line 77, in forward
        data = self.func(data)
        # Get grads
        grads = torch.autograd.grad(
                ~~~~~~~~~~~~~~~~~~~ <--- HERE
            # TODO:
            # This makes sense for scalar batch-level or batch-wise outputs, specifically because d(sum(batches))/d wrt = sum(d batch / d wrt) = d my_batch / d wrt
RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.

Currently looking through your suggestions above for more ideas! In the meantime, though, do you know what's going on with this output?

[Linux-cpp-lisp]

OK I'm a little confused... this is slightly different than your error before. How did you compile your model to TorchScript before?

The error is basically the same in terms of content though...

Can you try to load your deployed model in Python and run a frame of data through it there to see if you have the same issue? That would help us understand if this is something about your model or the C++ code around it.

(FYI I will be somewhat unreachable for the next few days, so if we can't get this resolved now just know that I will be back and am not ignoring the thread)

[Linux-cpp-lisp]

Another guess: when you build your input in C++, are you explicitly providing edge_vectors? That would break the autograd graph and give exactly this error. The network should only be provided with pos, edge_index, an edge_cell_shift (in the periodic case) so that it can build edge_vectors out of pos so that taking grad w.r.t. pos gives the right answer.

Even if this isn't your bug, it is probably something I should try to make an error rather than a silent failure... if that's possible.

[isabel-s-rosa]

This was it!! Thank you so much:) This discussion was super helpful!

[simonbatzner]

@isabel-s-rosa, glad Alby was able to resolve the issue! On a side note, I noticed a few settings in your input file that I think are not optimal or outdated, I'll list them below and I hope they will give you better training. As Alby pointed out, we also recommend to start with example.yaml, not full.yaml. Either way, here are some suggestions:

• allow_tf32: true --> I would set this to false for now, this is an experimental feature and only makes sense on an A100
• append: false --> If you want to start/restart runs arbitrarily, set this to true, and it will simply restart the run + append to your log file
• resnet: true --> we moved the default for this to false, since it usually doesn't help. The important thing is to turn on use_sc: true, which acts similar to a resnet
• num_basis: 20 --> this might work well, but is something that I would definitely be careful with. We found our default of 8 to work quite well across a larger number of materials. If 20 works better for you then of course use it, but just know that this is an important parameter and might cause unstable results (the Bessel basis has a 1/r term in, so with higher Bessel roots you will have larger values as you go to 0, which can be a bad thing)
• max_epochs: 20: I typically just set this to a very very large value. We only save the best model on the validation set, so training longer doesn't hurt
• Your loss function, you use forces: 100 together with - PerAtomMSELoss --> this will put a lot of weight on the forces. If you use PerAtomMSELoss then that will normalize your total potential energy to a per-atom quantity. That means that now both the forces and the atoms are local quantities, and so after this normalization it's a good idea to start with giving them equal weighting (the 100:1 makes more sense when the forces are local but the energy is a global quantity). If you use the PerAtomMSELoss (which is a good idea), then a setting of 1:1 in the energy and force terms should be a good starting point, i.e.:

loss_coeffs:                                                                       
  forces: 1                                                                       
  total_energy:                                                               
    - 1
    - PerAtomMSELoss

Let me know if you have questions + good luck!