[ANN] ReactantNitro.jl: Reactant-first training framework

Package Announcements
, ,
1

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 massive 2-10x 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 (Host by 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 GraphConst values are the only thing beyond code changes that need a recompile, which allows the framework to know when we should invalidate a cached program. Host values are guaranteed to be invisible to the compiled program, so they cannot change it. Changes to Device values do not need a recompile unless the shape or type of it changes.
  • The framework knows exactly when Device values 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.

function ReactantNitro.forward(e::MnistMLP, model, ps, st; img)
    logits, st_new = Lux.apply(model, img, ps, st)
    return logits ./ e.smoothing, st_new
end

function ReactantNitro.loss(e::MnistMLP, logits; label)
    logp = logits .- log.(sum(exp, logits; dims = 1))
    return -sum(label .* logp) / 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.
2

Very cool!