torch.autograd.grad

[kudoryutaro]

Can I use autograd.grad with torch-tensorrt?
We pass the coordinates of the atoms to the model, calculate the potential energy, and then differentiate the potential energy by pos to calculate the force. in model.forward(), we use torch.autograd.grad, but Is it possible to transform this model with torch-tensorRT?

Pseudo Code

class Model():
def forward(pos: torch.Tensor):
potential_energy = torch.sum(pos) # potential_energy is scalar
force = torch.autograd.grad(potential_energy, pos)
return potential_energy, force

[narendasan]

You could potentially use autograd as a pytorch op in conjunction with TensorRT ops via torch-tensorrt. But gradients will not be propagated through the TensorRT components. So like in this use case, I'm not sure it would work since you need the information from sum. @frank-wei do you have any ideas about how we could handle cases like this?

[lmuellender]

I'm interested in a similar use case. Could you elaborate on what you mean by "use autograd as a pytorch op in conjunction with TensorRT ops via torch-tensorrt"? Would we have to register autogrid as a custom op in tensorrt?

The separate torch.sum() is not the critical issue, as it could be built into the model. What we need would be the output of the model, as well as the gradients w.r.t. the model inputs. But, if gradients aren't backpropagated through the model, autogrid won't be much use in any case, no?

[narendasan]

Theres a few approaches we have thought about but have not spent anytime implementing. You can keep the components that need gradients in PyTorch and anything else in TensorRT, this works for cases where you are doing fine-tuning or something where most weights are fixed. For layouts where you are just fine-tuning network heads, this could improve training time. The other workflow we have thought about is supporting the backwards pass of AOTAutograd, which creates an additional graph for the backwards pass out of core aten ops which TensorRT could run but this would take much more work inside torch-tensorrt, could be space for the community to contribute if there is interest. cc: @cehongwang

[lmuellender]

Thanks for the update @narendasan! I figured as much. I'd be very interested in any progress regarding this, but I probably don't have the time (or know-how) to contribute myself here.