Supporting any Torch model (not just TorchScript)

[RaulPPelaez]

It would be great to be able to use a torch.compile'd model in OpenMM-Torch.

AFAIK there is no way to cross the Python-C++ barrier with a torch.compile'd model. Nothing like torch::jit::Module for TorchScript.

I can think of two solutions:

1. Let TorchForce accept any Python class as module (as long as it has the right Tensor input/outputs to forward).
   Doing this requires sending a generic Python class/function to C++ through SWIG, which I have not been successful in doing. It is really easy with pybind, but I cannot manage to mix pybind and swig for the life of me.
2. Take a TorchScript model like we do now, but internally call torch.compile.
   There is no way AFAIK to call torch.compile from C++, so we would have to invoke that via pybind at, perhaps, the TorchForce constructor. Then a py::object would be stored instead of a torch::jit::Module.

I think 2 is the simplest way with the current state of the codebase. Allowing something not easily serializable as the model (aka not TorchScript) would make serializing TorchForce an issue.

I would like to hear your thoughts on this!

[sef43]

here is a relevant PyTorch forum thread. At some point torch.export will exist as a option maybe
https://dev-discuss.pytorch.org/t/the-future-of-c-model-deployment/1282

Edit: maybe it is already possible, this looks promising: https://pytorch.org/docs/main/torch.compiler_aot_inductor.html

[peastman]

Any solution that requires a Python runtime to be present will be very limiting. Think of Folding@home, for example.

It sounds like this is all still in flux, and there are important things torch.compile can't do yet. Hopefully once everything settles down, there will be a clear migration path. I really hope they continue to support jit compilation, though. Having to rely on ahead-of-time compiled libraries would also be very limiting, and likely infeasible for important use cases.

[RaulPPelaez]

Steve's AOT thingy seems to be the only pytorch-endorsed way, but I have zero faith it is actually usable as of today -.-
I agree, Peter, things are super experimental still.
It is a same though to be able to, in the same Python script, run a torch.compiled model and run an OpenMM simulation but not mix the too. From an user perspective it is a bit frustrating.

[zhi-wang]

It would be great to be able to use a torch.compile'd model in OpenMM-Torch.

AFAIK there is no way to cross the Python-C++ barrier with a torch.compile'd model. Nothing like torch::jit::Module for TorchScript.

We recently encountered the same need to deploy torch compiled models in MD simulations and developed a solution that offloads some computations to Python. We wrote a short note focusing on this approach and included some very basic benchmarks.

We’re happy with the solution we arrived at, but in that note we did not address some of the less-than-ideal experiences we had with torch compiled models. Our main concerns (with the compiled models, not the deployment solution) are:

1. Some models are heavily decorated with torch.jit, making true out-of-the-box usage quite rare.
2. Models often need to be reimplemented with fixed-shape tensors, or torch’s dynamic-shape support requires further validation (which we didn’t have the bandwidth to investigate).
3. We discovered bugs (related to index_add) in the backend for torch 2.5 that caused numerical errors.
4. We are also aware of the limitations in AOT-based gradient calculations, so in our proof-of-concept we restricted ourselves to backpropagation for gradient computation.

All in all, we agree that the torch compile toolchain still needs further refinement. In parallel, existing MLFF models may also require non-trivial changes to accommodate torch compile. Nevertheless, I’m optimistic that future improvements will bring enough acceleration to justify updating models and perhaps more importantly to influence the new models under development.

The code and the note can be found at https://github.com/bytedance/OpenMM-Python-Force