PyTorch StableHLO Support

[csvance]

Continuing the discussion: https://discourse.julialang.org/t/node-training-performance-w-lux-vs-torchdiffeq/134860/5

Essentially it would be good if existing models, layers, etc in PyTorch could be converted to StableHLO and used within Reactant. This would reduce the barrier of entry to working with deep learning/neural networks in Julia for people who are heavily invested in the PyTorch ecosystem.

It looks like the easiest way to do this would be using TorchXLA and one of the two methods defined here: https://docs.pytorch.org/xla/master/features/stablehlo.html

Essentially we would just need the glue between ReactantPythonCallExt and one of these methods. I would partial to using the exported model approach since it likely is a smaller surface for things to go wrong and does not depend on on Jax.

@wsmoses am I understanding this correctly?

As an aside, it would be extremely useful if we could find a way to automatically bring over the weights from pretrained backbones or with whatever random initialization was used on the PyTorch side.

[csvance]

@wsmoses I managed to convert, compile, and execute a PyTorch model to run within Reactant. The interface may need to be different than the Jax version for a few reasons:

1. @compile complains about some of the array conversions in pytorch_to_stablehlo.
2. PyTorch differs from Jax since it associates modules with parameters. Users are going to want to get the parameters back from the module especially when they are pretrained / already have a good initialization from the PyTorch side. Handling this is what causes problems with @compile.

using PythonCall
using Reactant: Reactant
using OrderedCollections

const torchexportptr = Ref{Py}()
const torchaxexportptr = Ref{Py}()
torchexportptr[] = pyimport("torch.export")
torchaxexportptr[] = pyimport("torchax.export")

# When we convert to StableHLO we get Jax types back for the weights
jax_tensor_to_rarray(x) = Reactant.to_rarray(convert(Array, pyconvert(PyArray, x)))
pytorch_tensor_to_rarray(x) = Reactant.to_rarray(convert(Array, pyconvert(PyArray, x.detach().numpy())))

function ensure_py_import_path()
    path = [pyconvert(String, s) for s in sysptr[].path]
    if pwd() ∉ path
        sysptr[].path.append(pwd())
    end
end

function pytorch_to_stablehlo(pytorch_module::Py, sample_inputs)
    exported = torchexportptr[].export(pytorch_module, sample_inputs)
    weights, stablehlo = torchaxexportptr[].exported_program_to_stablehlo(exported)
    mlir = pyconvert(String, stablehlo.mlir_module())

    inps = [pytorch_tensor_to_rarray(inp) for inp in sample_inputs]

    params_k = [pyconvert(String, k) for (k, v) in exported.named_parameters()]
    buffers_k = [pyconvert(String, k) for (k, v) in exported.named_buffers()]
    # The convention used by the export is (parameters..., buffers..., inputs...)
    weights_k = vcat(params_k, buffers_k)
    weights_v = [jax_tensor_to_rarray(w) for w in weights]

    # Variables that are not used in the forward pass are discarded
    kept_var_idx = filter(
        x->x<=length(weights_k),
        [pyconvert(Int64, i)+1 for i in stablehlo.module_kept_var_idx]
    )
    weights_k = weights_k[kept_var_idx]
    weights_v = weights_v[kept_var_idx]

    state_dict = OrderedDict{String,Reactant.AbstractConcreteArray}()
    for (k, v) in zip(weights_k, weights_v)
        state_dict[k] = v
    end

    # Get all module functional arguments in the order expected by MLIR so we can call it with hlo_call
    linear_args = Vector{Reactant.AbstractConcreteArray}()

    # State (parameters + buffers)
    for v in values(state_dict)
        push!(linear_args, v)
    end

    # Forward Inputs
    for inp in inps
        push!(linear_args, inp)
    end

    return mlir, state_dict, linear_args
end

function compile_pytorch_module(pt_module::Py, sample_inputs, device)
    mlir, state_dict, linear_args = pytorch_to_stablehlo(pt_module, sample_inputs)

    # Send all the inputs to the correct device before tracing
    linear_args_device = [arg |> device for arg in linear_args]

    # Compile the module
    compile_fn(args...) = Reactant.Ops.@opcall hlo_call(mlir, args...)
    model_compiled = Reactant.@compile compile_fn(linear_args_device...)

    # Move the state dict to the device
    for (k, v) in state_dict 
        state_dict[k] = v |> device
    end

    # Return everything the user needs
    return model_compiled, state_dict
end

function main()
    timm = pyimport("timm")
    torch = pyimport("torch")

    Reactant.set_default_backend("cpu")

    # Load the model and create example inputs
    model = timm.create_model(
        "resnetv2_50x1_bit.goog_in21k"; pretrained=true, num_classes=0, in_chans=1
    )

    inp_pytorch = torch.randn(1, 1, 224, 224)
    sample_inputs = (inp_pytorch,)
    oup_pytorch = pytorch_tensor_to_rarray(model.forward(inp_pytorch))

    xdevice = reactant_device()

    # nn.Module -> ExportedProgram -> StableHLO -> @compile
    model_compiled, state_dict = compile_pytorch_module(model, sample_inputs, xdevice)

    # Try calling it
    inp = pytorch_tensor_to_rarray(inp_pytorch) |> xdevice

    args = [_ for _ in values(state_dict)]
    # The last args are the non weight inputs to nn.Module.forward(...)
    push!(args, inp)

    oups = model_compiled(args...)
    oup = first(oups)

    @info "difference" oup_pytorch .- oup
end

main()

┌ Info: difference
│   oup_pytorch .- oup =
│    1×2048 Reactant.ConcretePJRTArray{Float32,2}:
└     0.0  0.0  1.62125f-5  3.57628f-6  7.60891f-7  0.0  -3.40492f-6  -2.24113f-5  0.0  -4.29153f-6  …  -5.43594f-5  0.0  5.55068f-7  -1.22469f-7  0.0  -3.50177f-7  -4.75906f-7  4.05312f-6  0.0

[csvance]

So the above doesn't work when there are buffers or unused items in the state dictionary in the forward pass. I'm working on an update now to properly handle that.