Replacing PyTorch with Julia and Reactant: One Year In

Nearly a year ago, I decided to give the Julia programming language a shot; not even for machine learning work initially, but as a way to speed up a production service which was CPU bound. Now I know many people are thinking “Julia has correctness issues” and “dynamic typing in mah prod”. In my experience the first is overstated; the second is a real trade-off, and one I’ve chosen to make, because the sort of things I value when deploying services and models for scientific computing are different than the sorts of things many other engineers do.

These were the parts of the language that sold me:

This post is about my journey, mostly focusing on the ML aspect, but also delving into my journey through the “production” side of things.

We greatly accelerated our main backend service by rewriting all of the performance sensitive code in Julia. After tasting the initial success I thought to myself maybe we can do something similar with our ML stack? While PyTorch is an absolutely incredible piece of software, there were a couple of pain points that made me feel like my use case wasn’t something that was going to get much support going forward:

Together, these made me want a deployment format I could count on for years and a compiler that sped up the models I was actually training. We had already adopted Julia for our backend and were building on it extensively, so the natural next step was to see whether it could handle our ML work as well.

While I initially found a robust array of packages for scientific computing, I could immediately tell fewer people were doing “PyTorch” style ML for lack of a better term. In medical imaging it’s still quite common to use CNNs as backbones are not the ceiling on most problems, believe it or not. So I surveyed the ecosystem and here is what I found:

The core proposition of Zygote seems compelling: automatic differentiation at the language level, through almost any Julia code you write. But for PyTorch-style training workloads, its speed and device memory usage weren’t competitive with PyTorch eager mode, so I decided to give Reactant a shot; there was a lot of talk around Reactant.jl and Enzyme being the next step forward for ML in Julia.

Reactant compiles Julia code into an MLIR program, Enzyme differentiates that program at the compiler level (Moses and Churavy 2020), and XLA turns the result into fused CPU/GPU/TPU kernels. StableHLO is a dialect of MLIR used by XLA, and it has a number of benefits such as promising five years of backwards compatibility. If we can compile our models to StableHLO, that means that we should have a relatively stable platform for model training and serving going forward, so I took the dive into Reactant.

On the Julia Dispatch podcast, William Moses, who leads Enzyme and Reactant, and Avik Pal, the creator of Lux, tell how Reactant came to be in their own words.

With Reactant, I don’t write backward passes for custom layers. Reactant traces Julia code into a small set of StableHLO operations, Enzyme’s optimization passes simplify that program, and Enzyme differentiates the optimized program, so the gradient gets the same simplifications as the forward pass. In PyTorch, a layer built from loops and indexing is often slow under autograd, and making it fast can mean writing a custom autograd.Function with its own backward. Now when I write a new layer, I only need to consider the forward pass, and Enzyme takes care of the backward pass without giving up performance.

To be fair, none of this is exclusive to Julia. JAX users can run a function through Enzyme-JAX’s optimization pipeline with the enzyme-ad package’s enzyme_jax_ir decorator, and the same package lets JAX differentiate external C++ code. Enzyme itself works on LLVM IR, so it can differentiate code from languages that compile to LLVM: C and C++ through Clang, and Rust and Fortran through their own bindings. What Reactant adds is the frontend: it traces ordinary Julia code, so an existing Julia package can become input to these compilers by being made traceable rather than rewritten against a framework’s array API.

The next two weeks of my life were a living hell. I had no idea what I was doing with Reactant, constantly hitting tracing errors, going OoM, dealing with hour long compile times. But I persevered and kept learning from my mistakes. I had a particularly difficult problem to iterate on that I just couldn’t get PyTorch to speed up: a large backbone combined with a NeuralODE, using torchdiffeq. The model took nearly a week to converge across 4x NVIDIA A6000 (Ampere). I kept going, convinced that the optimizations provided by XLA would prove decisive for the NeuralODE part.

Then the eureka moment came: effectively a 10x speedup over what I was seeing with PyTorch eager. All that time I had spent learning the basics of Reactant had paid off, and now I was in a position to quickly iterate on the model and solved the problem later that week. For the sake of transparency, it’s worth noting that I have since revisited a similar workload with torch.compile, and got much more comparable numbers with Reactant still holding the lead, but not an “order of magnitude”. Regardless, I was sold: it was time to build a full training and serving stack on top of Reactant.

ReactantNitro.jl is an ML framework with a training loop designed around making Reactant easier to work with and avoiding unnecessary compiling. It has a REPL + Revise.jl first workflow: you can hot reload your data processing code, sweep hyperparameters, etc. without restarting your REPL or re-compiling.

Why another framework, you say? There are two main reasons. The first is that a compiler like Reactant warrants its own framework, one that handles the pitfalls I ran into before I created ReactantNitro: datasets ending up in the trace and slowing compiles, programs recompiling when nothing they depend on has changed, stale programs after a code edit, device transfers at unexpected times, and device memory building up across batches.

The second is a training workflow built around Revise. In October 2025 I wrote on Hacker News that waiting 15+ seconds to test small changes to my PyTorch training code was rather annoying, and that “sometimes I wish we could have a training workflow similar to how Revise works.” That is how I train now. I keep one Julia session open, edit the model, loss or data pipeline in my editor, and Revise patches the change into the running session. The next run picks it up without restarting Julia or reloading packages, and Nitro recompiles only the programs the edit reaches. The @experiment declaration makes this possible by stating which settings shape the compiled program and which can change freely. Each stage of the workflow, from setting up an experiment through training, validation, evaluation, prediction and export, is a function call at the REPL, and a Ctrl-C stops a run cleanly and hands the session back. It turns out that an interface designed for a REPL is only a few steps away from one an agent can use over MCP, so Nitro’s MCP tools, covered in Working with agents, call the same entry points.

Nitro also builds in training techniques that are easy to get subtly wrong when you implement them yourself: an exponential moving average of the weights, checkpoint averaging, and sharpness-aware minimization (SAM and ASAM) (Foret et al. 2021; Kwon et al. 2021). SAM is a good example. It takes a second gradient at perturbed weights, and that has to be combined correctly with gradient accumulation, the inner optimizer and multiple devices. In Nitro the perturbation is a first-class part of the gradient step, so SAM comes at no implementation cost: turning it on is one line, with no manual optimization loop to write and none of those details to keep in your head. The same goes for an agent: trying SAM is a one-line change it can make and evaluate, instead of an optimization loop it has to write and get right, so its effort goes into the parts of the problem that matter.

ReactantServer grew out of one question: can we serve more models than their weights can fit on the GPU by hot swapping them in real time? That solved a real problem for us. We have a lot of models that are called infrequently, and together they were close to no longer fitting on our pod. Beyond that, it’s a general serving platform for StableHLO.

ReactantServer.jl uses Reactant’s PJRT interface to serve StableHLO produced by Reactant.@compile using the KServe V2 protocol. It allows you to express arbitrary pre and post processing in pure Julia without having to set up a Python environment or compile any stubs, and supports hot swapping more models than you can fit in memory at PCIe speed (typically between 10-100 ms for most CNN vision models).

Neither of these got built overnight. Three smaller projects came first, and ReactantNitro and ReactantServer both depend on them.

Luximm.jl does for Lux.jl what timm does for PyTorch. It ports timm’s image backbones to Lux and loads their pretrained weights directly from Hugging Face Hub. It covers ResNet, SE-ResNet, BiT ResNetV2, ConvNeXt, ConvNeXt V2, VGG, ViT and CoAtNet, and parity tests check each port against the PyTorch original. The ConvNeXt V2 backbone in the walkthrough below comes from Luximm.

gRPCClient.jl and gRPCServer.jl give Julia a native gRPC stack, which ReactantServer uses to speak KServe V2. I rewrote gRPCClient.jl on top of libcurl’s HTTP/2 support. Sébastien Celles and I wrote gRPCServer.jl.

What would a good blog post be without actually showing how you can train and serve your own models end-to-end with the Reactant stack? Yes, it’s a real thing, you can do it today if you want. All of the code is in the companion repo, so you can follow along or run it yourself. Since we are dealing with Julia, there must be a reason to add a NeuralODE somewhere. Honestly, I don’t need one: I still find these models fascinating, especially the functions and representations they learn. Let’s start with how the full stack fits together.

Train: one Julia process

Serve: a container

@experiment + hooks your code: one struct, a handful of hooks

Lux.jl explicit-parameter model Luximm pretrained backbone

ReactantNitro.jl compiles the gradient, optimizer, and eval programs data pipeline · checkpoints · export

Reactant.jl traces Julia code to MLIR (StableHLO)

Enzyme reverse-mode AD on the traced program

XLA, through PJRT optimizes and compiles for CPU, GPU, or TPU: any XLA backend C++ libraries; no Python anywhere

model bundle written by ReactantServerExport StableHLO, weights, IO spec

load

ReactantServer.jl KServe V2 over gRPC (:8001) compiled ahead of time, CUDA

gRPC

ReactantServerClient.jl or any Triton or KServe gRPC client

The problem: PathMNIST

We are pointing the stack at PathMNIST: 224x224 tissue images in 9 classes. It’s distantly related to the problems I work on, none of which have had an “ImageNet” moment with a high quality public dataset, and it’s the closest I could find to a medical imaging dataset that isn’t a toy like MNIST but is still small enough for a demo. If you know of a better one, I’d like to hear about it.

juliaup add 1.13.1 git clone https://github.com/csvance/PathMNISTReactant.jl cd PathMNISTReactant.jl julia +1.13.1 --project=.

using Pkg Pkg.instantiate()

ReactantNitro.jl: Declaring the @experiment

ReactantNitro uses a small experiment macro to ensure things are wired up correctly and Reactant’s tracer doesn’t trace over things it shouldn’t, which can be a source of slow compiles. We declare an @experiment for PathMNIST. All of ReactantNitro’s hooks dispatch on this type.

Each field’s type says how it reaches the compiled program. A GraphConst field is baked into the program as a constant, so changing it recompiles; it is for values that change the graph, such as the number of classes or a layer width. A Device field