[WIP] PyTorch Import

[csvance]

Note: this isn't quite ready for full review yet.

Implements #2065

It includes monkey patching that is potentially brittle due to an inputs ordering difference in torch.export and torchax. It might be possible to avoid the monkey patch, I'm looking into alternative solutions. For torchscript there likely is not an alternative for the workarounds I implemented here. At the end of the day we are dealing with a format that has been deprecated for over 5 years now.

TODO:

• Investigate potential alternatives to monkey patch for input/param ordering
• Document matmul precision behavior when torchax sees a GPU vs when it doesn't - potential footgun
• FP64 support w/ tests
• Scalar type promotion matches PyTorch w/ Jax x64 mode
• Find a way to make torch install without requiring CPU version workaround

[csvance]

Whenever Jax sees a GPU on startup, it appears to lower matmul precision to something less than FP32. For the test with the linear layer, relerror is around 1e-7 or so, when Jax sees a GPU on startup its around 1e-4, which was causing one of the tests to fail. I'm still getting familiar with how all of the tests are wired up in Reactant and what to expect under different testing conditions.

Also the tests don't seem run with a GPU version of torch installed. I'm trying to figure out a way to automate install of the CPU version via CondaPkg. Installing the same version manually via pip in the environment works.

EDIT: I added a workaround to install the CPU version of torch. Here is the error for the GPU version:

ImportError: .../torch/lib/libtorch_cuda.so: undefined symbol: ncclCommResume

[codecov]

Codecov Report

❌ Patch coverage is 13.33333% with 52 lines in your changes missing coverage. Please review.
✅ Project coverage is 35.87%. Comparing base (b39a1fc) to head (df6383d).
⚠️ Report is 1150 commits behind head on main.

Files with missing lines	Patch %	Lines
ext/ReactantPythonCallExt/pytorch.jl	2.08%	47 Missing ⚠️
ext/ReactantPythonCallExt/ReactantPythonCallExt.jl	55.55%	4 Missing ⚠️
ext/ReactantPythonCallExt/overlays.jl	66.66%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #2928       +/-   ##
===========================================
- Coverage   68.16%   35.87%   -32.30%     
===========================================
  Files         109      220      +111     
  Lines       11779    31391    +19612     
===========================================
+ Hits         8029    11260     +3231     
- Misses       3750    20131    +16381     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[mofeing]

Whenever Jax sees a GPU on startup, it appears to lower matmul precision to something less than FP32.

right, you need to do the following:

import jax
jax.config.update("jax_enable_x64", True)

[mofeing]

for the symbol clashes, it's a problem that reappears periodically. I have an idea on how to solve it but @gbaraldi and I are gonna try it first on MLIR.jl

[mofeing]

minor thing, but you can skip the Python interpretation and directly call the C API for isinstance if you use pyisinstance directly

Suggested change

	is_torch_module(f::Py) =
	TORCH_EXPORT_SUPPORTED[] && pyconvert(Bool, pybuiltins.isinstance(f, torchptr[].nn.Module))
	is_torchscript_module(f::Py) =
	TORCH_EXPORT_SUPPORTED[] &&
	pyconvert(Bool, pybuiltins.isinstance(f, torchptr[].jit.ScriptModule))
	is_torch_module(f::Py) =
	TORCH_EXPORT_SUPPORTED[] && pyisinstance(f, torchptr[].nn.Module)
	is_torchscript_module(f::Py) =
	TORCH_EXPORT_SUPPORTED[] && pyisinstance(f, torchptr[].jit.ScriptModule)

[csvance]

I'm going to rework the warning about matmul precision. It shouldn't be written from the assumption that the user is using FP32. We should also clearly communicate Jax's behavior and how what devices it detects impacts the precision used for lowering. If the user wants TF32, they need to have an Ampere or later GPU visible on the machine used for the import.

[csvance]

Think I found a way to make the input ordering work without monkey patching. Testing on ~100 different models before pushing changes here.

[csvance]

I'm reading that the x64 toggle per model might be an anti pattern / lead to unexpected behavior. I was running into some issues with scalars causing type promotions they shouldn't have been based on the original model semantics. Working on fixing this up now. The goal is the type promotion behavior with the x64 flag enabled should be the same as the original model for the same inputs.

EDIT 1: we might need to demote floating point scalars on the PyTorch side to single precision to preserve weak type semantics. If we don't do that Jax will treat weakly typed float scalars as double precision when using the x64 mode, which then causes type promotion to double precision for non scalars.

EDIT 2: Just removed the need for monkey patching + fixed scalar type promotion causing issues with x64 mode.

[wsmoses]

alternatively to jax, could we use torchtpu to import the stablehlo directly? https://developers.googleblog.com/torchtpu-running-pytorch-natively-on-tpus-at-google-scale/

cc @GleasonK

[csvance]

I think I’ve addressed most of the open issues at this point, aside from the known CUDA symbol problem. The torchtpu path sounds promising.

Keep in mind with the following I'm not a compiler guy. So sorry if I say something truly asinine 😅

While testing, I ran all of my models through the current PR’s import system and noticed a large jump in compile time compared to how I handle this in my own setup. I suspect it comes from inlining weights as constants on bigger models, particularly once you get past the 1e8 parameter range. I believe the current Jax important code does this as well, which is why I implemented it the same way for the PyTorch path.

Is inlining weights as constants the best default for general model import, rather than passing them in as arguments? For large models it seems to make compile time explode, and if a user wants to compile the same function across multiple input sizes, baking the weights into the MLIR isn’t ideal unless they’re small enough to be negligible.

For context, I’m building an inference server, and my approach has been to share a single .safetensors as additional model arguments alongside one .mlir per batch size / input shape. In my case it also makes it easy to load weights to GPU on demand.

[GleasonK]

could we use torcht

Once its publicly available, yes you should be able to. Code looks well modularized enough where you just have some layer that exports to mlir and go from there, should be easy enough to swap in

[wsmoses]

lgtm, but let @avik-pal and/or @mofeing give it a once over before merging

[mofeing]

... aside from the known CUDA symbol problem.

Do you mean the ncclCommReduce symbol clash?

I suspect it comes from inlining weights as constants on bigger models, particularly once you get past the 1e8 parameter range.

Is inlining weights as constants the best default for general model import, rather than passing them in as arguments? For large models it seems to make compile time explode,

I cannot say something for sure without a profile, but I would bet that at that size, the constant propagation / folding mechanism is to blame because it is super simple and running in 1-thread, probably no vectorisation, ...

If it's an option, I would pass the model as arguments. If that's not an option, I would try add a "no-inline" func.func that contains and returns the constant.

[csvance]

If it's an option, I would pass the model as arguments. If that's not an option, I would try add a "no-inline" func.func that contains and returns the constant.

I think the Jax implementation might be signifigantly less aggressive about inlining than I originally thought. If that's the case I should just be able to make the PyTorch implementation match it, although the problem is its not totally intuitive from a user perspective as to what the order of the inputs should be. Perhaps we could pass the user back a helper function that would take the state dictionary and map it to the correct order. Regardless, it looks like a slightly different interface will be required between the two because of the differences in their semantics (PyTorch uses encapsulated parameters/buffers, Jax uses pure functions). It's obvious what the input ordering should be for a converted Jax function, not so much with PyTorch.