[Blog] Replacing PyTorch with Julia and Reactant: One Year In

Specific Domains Machine Learning
, ,
1

Over the past year, Medical Metrics Inc. has been migrating our internal model training/serving stack to Julia and Reactant. This resulted in several open source projects building on top of Reactant, and I thought it would be a good idea to write a deep dive style blog post on how these came to be and how they fit together, including an end-to-end example starting from a dataset and finishing with calling a served model from Julia over gRPC:

The tutorial part of the post covers the following among other things:

  • Training CNN+NeuralODE classifier w/ ReactantNitro.jl
  • Exporting the model and creating a bundle for serving with ReactantServerExport.jl
  • Loading the model with ReactantServer.jl for serving
  • Using ReactantServerClient.jl to evaluate the served model
  • Benchmarking training speed against PyTorch / torch.compile

I don’t claim to be an authority on Julia/Reactant by any means, but hopefully this provides a glimpse into the ecosystem and how the different parts of it fit together from a practical/applied standpoint. Enjoy!

2

Very interesting! I’m pleasantly surprised that Julia turned out to be the more practical language for ML model development compared to PyTorch. Usually, what I hear is that Julia’s AD/ML ecosystem is immature and buggy, and slower than PyTorch because all work is done in pre-compiled kernels anyway and Julia suffers from huge TTFX for ML stacks.

Also an excellent read with beautiful graphics!

3

TLDR: it was difficult getting started, but I don’t really spend time fighting with the ecosystem in my day-to-day work; all of that time is spent being productive working on the actual problems I’m trying to solve.

Fresh from PyTorch I hit a large number of snags and a couple bugs trying to get my first model training with Reactant. While I found success and a workflow I was satisfied with in the end, reaching that point did require significant time investment learning the ins and outs of working with Reactant and Enzyme. All of the lessons learned went into ReactantNitro.jl, which is explicitly designed around avoiding these sorts of issues. My hope is that this makes using Reactant far more accessible for ML work.

Reactant development can also move rather fast, which means sometimes things get broken with new versions. I don’t aggressively update to the latest version, and when I do upgrade I usually run a reference training pipeline that is known to be working and make sure everything still looks ok. Personally I haven’t seen any of my own workloads break this way, but practicing some degree of caution never hurts when you are dealing with bleeding edge compiler tech :sweat_smile:

PyTorch’s kernel advantage isn’t that large, and when you are working with scientific machine learning applications, better kernel fusion often beats out faster individual kernels. Enzyme producing a backwards pass over the already optimized forwards pass, optimizing that, and then handing it over to XLA for whole program optimization/fusion is essentially the best possible case for effective kernel fusion. IMO, that’s why you see Reactant pull so far ahead in the pure NeuralODE example.

4

Great blog post. If I ever have to train a NN again I’d love to test out reactantserver.jl.
I’m curious though: What motivated you to go this way instead of using onnx? Is it because of no support for reactant models?

5

Three main reasons we decided to build our own:

  1. We have many infrequently used models that don’t all fit in a single pod’s GPU memory. Hot swapping models in milliseconds is the difference between a single-pod deployment and splitting inference serving out into a cluster; with a smaller unit of deployment, scaling horizontally can be done with more precision and redundancy becomes less expensive. Keeping everything in one pod also lets us pass tensors through shared memory, since our production service and the inference server share an IPC namespace on the same physical machine.
  2. I wanted the StableHLO a model compiles to be exactly what gets served, with no conversion to another format and no separate lowering to validate. This means that if the model compiles during training, we are good to go for serving.
  3. We need bit-identical inference run to run for production. ReactantServer’s regulated profile provides this as long as the image, GPU class, driver and numerics mode are held constant and the autotuning cache is persistent: Reproducibility & Regulated Use · ReactantServer.jl. Our CI/CD pipelines assert this at the bit level.
6

What did you use to make those visuals?!?!

7

The animated/interactive parts were Javascript driving a <canvas>. Those were all produced by Claude. Every static figure in the post was produced with CairoMakie.jl. One thing I would like to experiment with in the future is doing some of the animations in WGLMakie.jl and exporting them as standalone HTML/JS via Bonito.jl. The more things I can do in Julia the better!

8

That sounds exciting, keep us posted on your experiences!

9

My superficial understanding is that the situation started changing when Reactant.jl became mature. Rectant is basically a GPU compiler that compiles your kernels, and array-oriented code is compiled in a way that largely avoids GC. The efficient machine code from the GPU compilation gives you a speedup compared with PyTorch, just as Google’s JAX in Python does.

I only briefly dabbled in deep learning, and the OP may have more insights about this.

10

I need to finish some Reactant integrations. I’ve gotten pretty far but it’s not done. I hope by December we are fully compliant. Billy has given me the tools I need.

11

Out of curiosity what tools are those?

12

mostly a giant pile of bugfixes

13

Look forward to seeing this in a tutorial or book form on how to use Julia for the traditional AI workflow usually covered by the Python tools. Especially for those of us who haven’t used the Python based tools and don’t plan to start now that we have Julia. Well written. Thanks for the work to move things forward!

14

@csvance Very good to see Julia shine in deep learning tasks, excellent work!

Since in the blog post you focus on CNNs+Neural ODEs, I was wondering what is your assessment of Lux + Reactant’s performance on more standard architectures.

Lately, I have been landing a lot of optimizations in Flux.jl and its ecosystem (Zygote.jl, NNlib.jl, MLUtils.jl, HuggingFaceDatasets.jl, HuggingFaceTokenizers.jl) and filed tons of issues to CUDA.jl, Enzyme.jl, and Reactant.jl (whose maintainers I thank a lot, they have been very helpful and responsive). So far, I am very satisfied with the improvements on non-reactant backends, slightly less so with Reactant.

For instance, in Refresh ResNet and GPT-2 examples and benchmarks - Pull Request #2794 - FluxML/Flux.jl - GitHub I benchmarked two standard training examples, a vision one (ResNet18 trained on Tiny ImageNet) and a language one (GPT-2 trained on FineWeb, similar to Karpathy’s nanogpt).

Training measurements on an RTX 5090 (higher is better):

Example Flux + CUDA Flux + Reactant PyTorch PyTorch compiled
ResNet, images/s 2,830 3,270 3,450 3,950
GPT-2, tokens/s 167k 131k 158k 163k

So for GPT-2, Flux + CUDA (+ Zygote) is on par with PyTorch, even slightly better,
while Reactant lags behind, most likely due to a missing fused/flash attention StableHLO operator, which forces reactant to materialize attention weights Fused (flash) attention for NNlib.scaled_dot_product_attention under tracing · Issue #3510 · EnzymeAD/Reactant.jl · GitHub

On the vision task instead, Reactant performs slightly better than Flux + CUDA, but still lags behind eage and compiled pytorch. It is likely that the gap between Flux + CUDA and eager pytorch will mostly fixed once cuDNN: add measured convolution plan selection for fixed-shape workloads · Issue #3506 · JuliaGPU/CUDA.jl · GitHub is resolved. I wouldn’t know how to improve Reactant’s performance though. Maybe Reactant: pipeline batch prefetch, transfers, and loss readback · Issue #2801 · FluxML/Flux.jl · GitHub (which I filed after reading your blog post) could help, but I don’t think it will be enough to match pytorch compiled, let’s see.

I would be curious to see the Lux (+ ReactantNitro.jl?) port of those examples, especially if done by you, since I could be missing something.

15

@CarloLucibello wow, those numbers are much closer than I expected to see for Zygote/CUDA! It’s been around half a year since I last tried Lux/Zygote/CUDA, but I would imagine after all of these improvements it could be a viable option for me for certain tasks that require the sorts of control flow / dynamic shapes which are not possible in Reactant. Appreciate all the hard work that went into improving the ecosystem for everyone!

Most of my workloads are on smaller backbones where kernel launch overhead tends to dominate, especially when doing CNN+NeuralODE training. Usually I won’t use a standard ResNet, instead opting for something like ResNetV2 50 BiT which uses weight standardization for every conv and group norm instead of batch norm. So you have a ton of tiny/smaller kernel launches in eager mode which tends to slow things down significantly. torch.compile is actually a little bit faster than Reactant for that specific backbone on an A6000, but Reactant is roughly twice as fast PyTorch as eager mode.

For GPT-2, kernel selection is probably the most important aspect since a single kernel can do sufficient work to where launch overhead is small in comparison. I’ve actually noticed for some much larger backbones, I don’t get as good performance with Reactant as I do in PyTorch even in eager mode. That seems to line up exactly with your results.

Both of those tasks would make a good example in the ReactantNitro.jl repo. I don’t have a 5090, but you could run it once I put it together and update your comparison here.

EDIT: maybe some of what you are observing is due to TF32 performance on your 5090. On non workstation/datacenter cards, TF32 performance is essentially the same as FP32; nVidia seems to have intentionally hobbled it. When you get the huge compute speedup from TF32 (~2x in my experience), kernel launch overhead becomes that much more important to minimize. All of the numbers I’m reporting are using TF32 on a card where it significantly speeds things up.

16

I couldn’t wait and did the port :slight_smile: Add Lux + ReactantNitro versions of the GPT-2 and ResNet examples - Pull Request #2803 - FluxML/Flux.jl - GitHub

ResNet-18, images/s GPT-2 124M, tokens/s
Flux + CUDA 2830 167k
Flux + Reactant 3270 131k
Lux + ReactantNitro 3510 138k
PyTorch 3450 158k
PyTorch compiled 3950 163k

Something can be still gained in Flux’s Reactant to match ReactantNitro performance. I wonder if we can catch up with torch compiled though on these (not tiny) networks, and if plain Flux + CUDA can also be improved on resnets. I’ll redo the benchmark on an H100 when I get a chance.

17

Thanks for this, wait for flash attention arriving in Enzyme cpp. This is what’s missing.
For images would need to play with the reactant compile options I guess.

18

I ran the ResNet18 example on an A6000. Your 5090 is ~2.1x as fast, so the kernel launch overhead actually becomes more relevant not less!

There is likely room to upstream performance improvements around XLA’s CUDA backend; it’s well known that it lags behind a bit compared to TPU. There may also be potential for Enzyme-JAX rewrite rules targeting these sorts of workloads in places as well as additional Reactant specializations.