Over the past year, PyTorch team has done a lot of work to improve the user experience of torch.compile and torch.export. As PyTorch Compiler Deployment team, we’ve seen a lot of interesting use cases around “minimal” deployment where people try to compile PyTorch models down to self-contained executables without dependencies because:
- While running torch.compile is useful, but the dependencies involved are a non-starter for a large variety of use cases.
- A lot of people want to just download a binary and run it (e.g. project like Llamfile is a cool project to demo LLMs with very little installation steps)
Although we have a lot of tools for enabling native deployment of PyTorch models (export/aoti), almost nobody is actually using them for building/shipping binary in practice.
Therefore, I want to understand what the obstacles were there, and so I tried packing Whisper into a single binary, and this is what I found.
Summary
- Although we have a lot of tools to compete with llamafile, the libtorch dependency and export’s top-down UX make it unappealing to use torch.export to turn an existing PyTorch model to a distributable executable.
- To address these issues we propose to push on the following projects:
- libtorch-free aotinductor.
- An eager-style torch.export API.
- TorchNative: a set of tech demos to showcase using torch.export+aoti on realworld PyTorch models.
Experiment on OpenAI Whisper
To concretely understand what could be helpful for export+aoti (and follow the spirit of PyTorch team’s User Empathy Day™), I decided to have a hack-a-week to see what I can come up with when using export+aoti on some random open source model. I went the extra mile for a whole week’s time because I want to truly understand what it would take to fully get rid of the Python environment to run whisper, and a single user empathy day is definitely not enough to take me from the Python wonderland to the no-Python wonderland 
It turns out that, it’s totally feasible to produce a native binary which takes an audio file and produces the transcription using export and aoti, i.e. a fully reproduced e2e whisper run with only torch compiled kernels and torch ops without Python! I put the demo and code here: https://github.com/zhxchen17/torchnative/tree/main/whisper_aoti (note: I still need to locally patch PyTorch to make this work. Link to the patch: gist:48bcf4bf5af77302a1047f6ba67e8568 · GitHub)
Here is a architecture overview of the produced binary:
Observation 1
Problem: Libtorch dependency makes us non-competitive for single-binary setting, because people are trying to get other users to download this dependency, while libtorch is an extra couple gigabytes (very bad).
Solution: Libtorch-free aotinductor.
Remember the architecture diagram we just saw, everything looks perfect, right? Well I intentionally leave 1 big part out of the diagram: libtorch! A more complete architecture looks like the following:
In our example, the binary size is on the order of magnitude of hundreds of megabytes (and can be optimized further), but libtorch itself takes a few gigabytes, which makes the distribution of the whole binary much harder.
As a side note, today the generated code from AOTI also cannot be easily statically linked because different model binaries tend to have symbol conflict in a single binary setting. Currently this is being worked around by dlopen-ing each shared object separately so their namespaces are isolated from each other.
Another issue is that with libtorch as dependency, we need to carefully maintain the ABI compatibility between the produced binary and the libtorch we’re shipping with. With the design of libtorch.so ↔ model.so shim layer, this problem is managed but the ABI issue will be eliminated all at once if we don’t dynamically link with libtorch.so in the first place, resulting in a much simpler deployment pipeline.
Overall, I strongly agree with Bin Bao’s vision that we should work on libtorch-free codegenerate in AOTI.
Observation 2
Problem: Export’s UX is too “top-down”, meaning either it works on the first try, or it’s hard to make it work. However, for “real” models like Whisper, it’s much easier to do it from a “bottom-up” way.
Solution: Eager style export API.
Today when a user wants to AOTI compile a model, the process is repetitive and tedious: First we need to acquire sample inputs from somewhere, which results in hacky code to “breakpoint” at model’s entry point and observe the inputs there. Secondly you need to pass ExportedProgram and sample inputs to AOTI’s API and package it manually.
For Whisper, I actually needed to export the model into 5 different smaller parts, so people can imagine how much boilerplate code I ended up with. Therefore, based on the sticky cache idea floating around PyTorch team, I want to explore the idea of a decorator style like the following:
In prototype phase, API design subject to change.
@_compile_cache("model", out="/tmp/zhxchen17/", input_matchers=({"x": {0: Dim("a")}, y: None}))
def model_code(x, y):
return x + y
model_code(x, y)
torch.export team has started the discussion of this API and we believe it will make users’ experience much better when exporting models with smaller granularity.
Observation 3
Problem: Almost nobody is using export/aotinductor in the public community and it’s unclear what is the standard workflow to deploy exported models in distributable binaries.
Solution: TorchNative: A set of tech demos to showcase how to use export/aotinductor and compile PyTorch models into single binaries.
It becomes more and more clear, after the Whisper experiment, that there’s an interesting problem space we can explore to further achieve export+aoti’s potential: reducing the dependency to the absolutely bare minimum and making it much easier to adopt. To give some example what seems to be useful to be built so far:
- no-libtorch deployment as the first class feature.
- Easier export API adoption by decorator style API.
- Better support for profiling.
- Weight sharing.
- Better CPU performance.
Feeling there should be a better umbrella name for these related projects, I coined a concept name for what we can achieve here: TorchNative, meaning we are using torch native tools to produce natively runnable binaries on any platform.
Under the umbrella name of TorchNative, our goal is to produce a bunch of highly informative and simple tech demos on how to apply PyTorch Compiler technologies to a variety of open source PyTorch models. While TorchNative’s key deliverables are a set of tech demos, our hope is that we can uncover many interesting issues from using export+aoti workflow like a real end user, and generate a bunch of new ideas like eager-style export API and no-libtorch deployment. For example, we are going to upstream a bunch of local patches we’ve made to PyTorch main branch to make sure everything we put in the demos/tutorials are working and ready to go for end users.
Takeaways
By working more on the end users’ side, we are able to make reproducible and holistic decisions on what to build in PyTorch Compiler to benefit end users, and we should continue working on enabling more models with export+aoti.
This post aims to share some of my thoughts on export+aoti since September and hopefully this brings some awareness to PyTorch developers/users on what we want to build next.