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!
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.
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
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.
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?
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 out behind a reverse proxy. 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.
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.
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.
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!
