# Optimizing Lux.jl and Reactant.jl performance

**URL:** <https://discourse.julialang.org/t/optimizing-lux-jl-and-reactant-jl-performance/135942>\
**Category:** Machine Learning\
**Tags:** lux, reactant\
**Created:** [March 1, 2026, 4:42pm UTC](https://discourse.julialang.org/t/optimizing-lux-jl-and-reactant-jl-performance/135942 "2026-03-01T16:42:49Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![KSepetanc](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ksepetanc/32/10841_2.png) [@KSepetanc](https://discourse.julialang.org/u/KSepetanc)\
**Post date:** [March 1, 2026, 4:42pm UTC](https://discourse.julialang.org/t/optimizing-lux-jl-and-reactant-jl-performance/135942/1 "2026-03-01T16:42:49Z")

</div>

I am facing two separate issues. One is related to setting Reactant initialization options to improve performance and the other is about avoiding unnecessary latencies in neural network training loop via compilation.

Note that my code works normally without additional options that I am trying here to tweak to improve performance. Used cpu is 9900x (12 core/24 threads) and gpu is 5070 Ti (Linux, Julia 1.11.9).

**Reactant initialization options**

1. When run with gpu backend, XLA allocates 75% of device vram regardless of how much it has. I have tried several things to disable it, but non worked:

```julia-auto
using Preferences, UUIDs
Preferences.set_preferences!(
    UUID("3c362404-f566-11ee-1572-e11a4b42c853"),
    "xla_gpu_fraction" => 0.4
)
using Lux, Random, Optimisers, Reactant, Enzyme, MLUtils, ProgressMeter, XLSX

```

another try:

```julia-auto
using Preferences, UUIDs
ENV["XLA_FLAGS"] = "--xla_gpu_fraction=0.4"
using Lux, Random, Optimisers, Reactant, Enzyme, MLUtils, ProgressMeter, XLSX

```

Preferences approach has no effect and ENV approach freezes Julia. Both approaches when run with `"XLA_PYTHON_CLIENT_PREALLOCATE"` set to `false` did not have effect either.

1. With cpu backend, the training uses only 6 threads resulting in 25% cpu usage as observed in system monitor. I do not know how change that. Setting Julia threads has no impact on this. Both `Reactant.addressable_devices()` and `Reactant.devices()` output:

```julia-auto
1-element Vector{Reactant.XLA.PJRT.Device}:
 Reactant.XLA.PJRT.Device(Ptr{Nothing} @0x000000002ac8ee50, "CPU:0 cpu")

```

**Avoid train latencies by compiling**  
Following several tutorials from Lux examples on homepage, initially I got to following form:

Training part:

```julia-auto
total_loss = ConcretePJRTNumber{Float32, 1}(0.0f0)

for (x, y) in train_loader
    (_, loss, _, train_state) = Training.single_train_step!(
    AutoReactant(), MSELoss(agg=sum), (x, y), train_state
    )
    total_loss += loss
end

```

Validation part:

```julia-auto
total_val_loss = 0.0f0

st_ = Lux.testmode(train_state.states)
for (x, y) in val_loader
    ŷ, st_ = model_compiled(x, train_state.parameters, st_)
    ŷ, y = cdev(ŷ), cdev(y)
    total_val_loss += sum(abs2, ŷ .- y)
end

```

Both accumulation of `total_loss` in train part and `cdev()` and `total_val_loss` in validation part slow down computation for about 20% each for my test case. I managed to resolve validation segment (20% loss is basically fully annulled).

Compile:

```julia-auto
function val_step(x, y, p, s, tl)
    ŷ, s = model(x, p, s)
    loss = sum(abs2, ŷ .- y)
    return tl + loss, s
end

val_step_compiled = @compile val_step(first(val_loader)[1], first(val_loader)[2], ps, Lux.testmode(st), ConcretePJRTNumber{Float32, 1}(0.0f0))

```

Loop:

```julia-auto
total_val_loss = ConcretePJRTNumber{Float32, 1}(0.0f0)

st_ = Lux.testmode(train_state.states)
for (x, y) in val_loader
    total_val_loss, st_ = val_step_compiled(x, y, train_state.parameters, st_, total_val_loss)
end

```

While this works for validation, I do not know how to do it for training part as `single_train_step!` can not be put in `@compile` as it has some device wrapping and compilation already. Also `compute_gradients` can not be used in this way.

---

<div class="post-metadata">

**Author:** ![wsmoses](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/wsmoses/32/26497_2.png) [@wsmoses](https://discourse.julialang.org/u/wsmoses)\
**Post date:** [March 1, 2026, 8:06pm UTC](https://discourse.julialang.org/t/optimizing-lux-jl-and-reactant-jl-performance/135942/2 "2026-03-01T20:06:48Z")

</div>

As for the first issue, see [Reactant.jl/docs/src/api/config.md at ca406f850272f17cb4659d995f8703aa5cacae5c · EnzymeAD/Reactant.jl · GitHub](https://github.com/EnzymeAD/Reactant.jl/blob/ca406f850272f17cb4659d995f8703aa5cacae5c/docs/src/api/config.md?plain=1#L47)

For your second part you don’t need total\_loss to be a concreternumber, you can just use a regular float32.

@avikpal may have thoughts on the lux api usage

---

<div class="post-metadata">

**Author:** ![KSepetanc](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ksepetanc/32/10841_2.png) [@KSepetanc](https://discourse.julialang.org/u/KSepetanc)\
**Post date:** [March 1, 2026, 10:16pm UTC](https://discourse.julialang.org/t/optimizing-lux-jl-and-reactant-jl-performance/135942/3 "2026-03-01T22:16:42Z")

</div>

Thank you for reply. It does fix gpu memory allocation:  
`ENV["XLA_REACTANT_GPU_PREALLOCATE"] = "false"`

No option is listed to control cpu threads use when run with cpu backend (explained above in issue 1, point 2). Do you think you can make such option if there is no other way?

For `total_loss` yes, it is not needed. I experimented with a lot so it got left in the code that way and the issue persist regardless. However, for `total_val_loss` when used with `val_step_compiled` it is needed (notice it is used as function parameter) as otherwise Reactant complains that either function is not compiled with the same type (if it compiled with concreternumber) or that it can not do +(Float32, concreternumber) if is it compiled with Float32. So concreternumber might be needed in train part too depending how we solve it.

---

<div class="post-metadata">

**Author:** ![KSepetanc](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ksepetanc/32/10841_2.png) [@KSepetanc](https://discourse.julialang.org/u/KSepetanc)\
**Post date:** [March 6, 2026, 8:58pm UTC](https://discourse.julialang.org/t/optimizing-lux-jl-and-reactant-jl-performance/135942/4 "2026-03-06T20:58:07Z")

</div>

@wsmoses sorry for pinging you. I did not find solution for setting CPU threads. If there is no way to set it using any flags or options, should I open an issue?

Edit: It is already open issue [How to control the number of threads? · Issue #63 · EnzymeAD/Reactant.jl · GitHub](https://github.com/EnzymeAD/Reactant.jl/issues/63)