I’m excited to finally release ReactantNitro.jl; a Reactant-first training framework inspired by PyTorch Lightning, but specifically designed around making working with Reactant.jl easy while still allowing for a large amount of flexibility without sacrificing performance. While Reactant and Enzyme have some truly incredible potential to speed up machine learning workflows, there are a number of stumbling blocks researchers and ML practitioners often hit coming from other ML ecosystems.
We developed ReactantNitro.jl at Medical Metrics Inc. to explicitly avoid these sorts of problems while keeping the signifigant speedup we were seeing for our models compared to PyTorch eager mode and torch.compile. It’s been in the works for around half a year now and has been used internally to successfully train many different models.
Experiment DSL designed specifically for Reactant
One of the core features of the framework is an Experiment DSL which lets you categorize your models parameters in a way that the framework can ensure correct and optimal behavior:
| Marker | Reaches traced code as | In the compile key? | Changing the value |
|---|---|---|---|
GraphConst{T} |
a baked literal | yes | recompiles, correctly: a different value is a different program |
Device{T} |
a device-resident traced input | no, by construction | never recompiles: sweep it, schedule it, rewrite it live |
unmarked, i.e. Host{T} |
not at all | no | never recompiles: driver-only, invisible to the tracer |
The frameworks callback hooks are keyed on your @experiment struct:
@experiment struct MyExp
"Structural: changes the emitted graph, so it bakes and is part of the compile key."
width::GraphConst{Int} = 128
"A traced input: sweep it or schedule it without recompiling."
smoothing::Device{Float32} = 0.05f0
"Unmarked, therefore Host: driver-only and invisible to the tracer."
max_epochs::Int = 20
end
Structuring things this way eliminates entire categories of pitfalls:
- The tracer will never trace over your dataset (
Hostby default); for large datasets, this can massively increase Reactant’s compile time - One of the first principles of the framework is reducing wasted compile time. Changes to
GraphConstvalues are the only thing beyond code changes that need a recompile, which allows the framework to know when we should invalidate a cached program.Hostvalues are guaranteed to be invisible to the compiled program, so they cannot change it. Changes toDevicevalues do not need a recompile unless the shape or type of it changes. - The framework knows exactly when
Devicevalues should be transferred to the device. For scheduled values it happens every step, for constant buffers it’s done a single time at the start of training. More on flexible scheduling later.
Eliminate Boilerplate
We take inspiration from PyTorch Lightning when it comes to eliminating boilerplate by providing hooks that are dispatched on your experiment type. You handle your forwards pass, loss function, and metrics; the framework handles the boilerplate glue that connects it all while aiming to still provide the user with sufficient flexibility. A manual optimization mode is also provided which can be used to train models like GANs.
ReactantNitro.forward(::MnistMLP, model, ps, st; img) = Lux.apply(model, img, ps, st)
function ReactantNitro.loss(e::MnistMLP, logits; label)
smoothed = (1f0 - e.smoothing) .* label .+ e.smoothing / 10f0
return -sum(smoothed .* logsoftmax(logits; dims = 1)) / size(label, 2)
end
These hooks are used to compile three different programs: training forwards/backwards, training optimizer step, and validation/test/infer. The reason we don’t roll the optimizer step into the forwards/backwards is to provide flexibility to do things like gradient accumulation and allow you to switch optimisers without recompiling the entire program.
REPL Driven Workflow + Kaimon MCP Extension
Every part of the interface was designed to be used in a REPL making use of Revise.jl to allow you to quickly iterate and test new ideas without unnecessary recompile. The invalidation system is Revise aware and tracks changes that would invalidate compiled programs, so compile only happens when its actually needed.
ReactantNitro.jl also has a Kaimon.jl extension which lets your agents drive training jobs entirely through MCP and a full set of skills for your agents to use. Thanks to @kahliburke for all the amazing work he has done on making this possible in the Julia ecosystem! We have experimented some with integrating the framework into KaimonSlate.jl as well, some follow up work is likely needed on my end to make it entirely seamless.
Flexible Metrics
Metrics can either be handled on the device side for maximum performance or host side for increased flexibility and avoiding re-compile after changing the metrics. This can be defined separately between for train and validation/test. Reductions are structured in a normalization aware way for validation/test, you return a tuple each step with (scalar, n) and the n is used in ReactantNitro.finalize_metrics. You can also return (arr, nothing) at each step and handle the reduction yourself for metrics like AUROC or R^2 that need to be done over all of the individual samples.
Parameter Schedules
You can schedule pretty much any kind of Device (including optimizer parameters like learning rate, weight decay, etc)
@experiment struct Seq2SeqExp
teacher_forcing::Device{Float32} = 1f0
max_epochs::Int = 20
end
ReactantNitro.schedules(::Seq2SeqExp) = (;
device = (; teacher_forcing = total -> t -> max(0f0, 1f0 - 2f0 * t / total)),
opt = (; eta = total -> OneCycle(total, 1f-3)), # ParameterSchedulers.jl, no dependency
)
This even works with more complicated configurations such as multiple optimisers keyed by the individual optimizer name.
Export Models to ReactantServer.jl
A ReactantServerExport.jl extension is provided which makes it straightforward to export your nitro models to ReactantServer.jl for serving in production. In the future we are also looking at different export paths such as tensorflow/jax which would allow for deployment in existing serving stacks.
Status
I plan on registering the package shortly. The API should remain stable as I’m fairly confident it can handle most common use cases, but eventually there will be a 1.0.0 release where we can make improvements to the interface based on feedback.
Work is ongoing to provide end-to-end examples, but the framework is simple enough that getting something off of the ground shouldn’t be terribly difficult.
AI Disclaimer
- This post was written entirely by me with no AI assistance.
- Coding agents were used heavily during the development of ReactantNitro.jl, but the design is entirely mine down to the last detail. The framework has been used to produce many models that had their predictions independently verified during internal testing. Of course as with all software there may be bugs; if you encounter any issues we would love an issue report or a PR addressing them.
- Documentation is largely AI generated but reviewed / refined by me; the process of improving it is ongoing.
No need to think, when you skip the compile it’s because it wasn’t needed and when you don’t its because something was invalidated.