Using GPU-accelerated Ensembles with Automatic Differentiation

EnsembleGPUArray comes with derivative overloads for reverse mode automatic differentiation, and thus can be thrown into deep learning training loops. The following is an example of this use:

using OrdinaryDiffEq, SciMLSensitivity, Lux, Optimisers, Zygote, DiffEqGPU, CUDA, Random

CUDA.allowscalar(false)

# A tiny Lux model whose parameters are what we train. It maps a constant
# input to the two ODE parameters.
const dense = Dense(1 => 2)
const x = Float32[1.0]
rng = Random.default_rng()
Random.seed!(rng, 0)
ps, st = Lux.setup(rng, dense)

u0 = Float32[3.0]

function modelf(du, u, p, t)
    du[1] = 1.01f0 * u[1] * p[1] * p[2]
    return
end

function model(p)
    prob = ODEProblem(modelf, u0, (0.0f0, 1.0f0), p)

    function prob_func(prob, ctx)
        return remake(prob, u0 = 0.5f0 .+ Float32(ctx.sim_id) / 100 .* prob.u0)
    end

    ensemble_prob = EnsembleProblem(prob; prob_func)
    return solve(
        ensemble_prob, Tsit5(), EnsembleGPUArray(CUDA.CUDABackend()),
        saveat = 0.1f0, trajectories = 10
    )
end

# loss function: run the Lux model to produce ODE parameters, then score the ensemble
function loss(ps)
    p_vec, _ = dense(x, ps, st)
    return sum(abs2, 1.0f0 .- Array(model(p_vec)))
end

println("Starting to train")

l1 = loss(ps)
@show l1

# Optimisers.jl handles parameter updates; Zygote.jl handles gradients
opt_state = Optimisers.setup(Optimisers.Adam(0.1f0), ps)
for epoch in 1:10
    grads = Zygote.gradient(loss, ps)
    Optimisers.update!(opt_state, ps, grads[1])
    @show loss(ps)
end

Forward-mode automatic differentiation works as well, as demonstrated by its capability to recompile for Dual number arithmetic:

using OrdinaryDiffEq, DiffEqGPU, ForwardDiff, Test, CUDA

function lorenz(du, u, p, t)
    du[1] = p[1] * (u[2] - u[1])
    du[2] = u[1] * (p[2] - u[3]) - u[2]
    du[3] = u[1] * u[2] - p[3] * u[3]
    return
end

u0 = [
    ForwardDiff.Dual(1.0f0, (1.0, 0.0, 0.0)),
    ForwardDiff.Dual(0.0f0, (0.0, 1.0, 0.0)),
    ForwardDiff.Dual(0.0f0, (0.0, 0.0, 1.0)),
]
tspan = (0.0f0, 100.0f0)
p = (10.0f0, 28.0f0, 8 / 3.0f0)
prob = ODEProblem{true, SciMLBase.FullSpecialize}(lorenz, u0, tspan, p)
prob_func = (prob, ctx) -> remake(prob, p = rand(Float32, 3) .* p)
monteprob = EnsembleProblem(prob; prob_func)
@time sol = solve(
    monteprob, Tsit5(), EnsembleGPUArray(CUDA.CUDABackend()),
    trajectories = 10_000,
    saveat = 1.0f0
)

EnsembleSolution Solution of length 10000 with uType:
SciMLBase.ODESolution{ForwardDiff.Dual{Nothing, Float64, 3}, 2, uType, Nothing, Nothing, Vector{Float32}, rateType, Nothing, SciMLBase.ODEProblem{Vector{ForwardDiff.Dual{Nothing, Float64, 3}}, Tuple{Float32, Float32}, true, Vector{Float32}, SciMLBase.ODEFunction{true, SciMLBase.FullSpecialize, typeof(Main.lorenz), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Nothing, @NamedTuple{}}, SciMLBase.StandardODEProblem}, OrdinaryDiffEqTsit5.Tsit5{typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!), FastBroadcast.Serial}, IType, SciMLBase.DEStats, Nothing, Nothing, Nothing, Nothing, Nothing} where {uType, rateType, IType}

Enzyme reverse mode with EnsembleGPUKernel

For GPUTsit5, use Enzyme to differentiate a scalar loss of the ensemble's final states. Set cpu_offload to zero and save_everystep = false. Keep the parameters on the host when constructing the problems; the ensemble solver transfers the compatible problems to the selected backend.

using DiffEqGPU, Enzyme, KernelAbstractions, SciMLBase, StaticArrays

function ensemble_loss(p, backend)
    rhs(u, p, t) = p[1] * u
    prob = ODEProblem{false}(rhs, SVector(1.0), (0.0, 1.0), SVector(p[1]))
    prob_func = (prob, ctx) -> remake(prob; p = SVector(p[ctx.sim_id]))
    ensemble = EnsembleProblem(prob; prob_func, safetycopy = false)
    sol = solve(
        ensemble, GPUTsit5(), EnsembleGPUKernel(backend, 0.0);
        trajectories = length(p), adaptive = false, dt = 0.05,
        save_everystep = false
    )
    return sum(s -> sum(s.u[end]), sol.u)
end

p = [0.2, -0.3]
backend = CPU()
dp = zero(p)
Enzyme.autodiff(Reverse, ensemble_loss, Active, Duplicated(p, dp), Const(backend))
@assert dp ≈ exp.(p)
dp

2-element Vector{Float64}:
 1.2214027581602014
 0.740818220681565

For NVIDIA GPUs, load CUDA and pass CUDA.CUDABackend() instead of CPU(). CUDA gradients with released dependencies remain unverified. The upstream GPUCompiler compatibility work and compatible CUDA allocation rules are required for that configuration. Loading Enzyme activates DiffEqGPU's transfer rules automatically. Reset the shadow buffer dp to zero before each independent reverse-mode call: Enzyme accumulates into it. Fixed-step gradients differentiate the numerical steps; adaptive gradients differentiate the executed solver path and are not a record-and-replay adjoint with a frozen mesh.

Callbacks and DAE initialization are solver features; their absence from this example does not imply that they are unsupported. Representative Float64 CPU checks also validate Enzyme gradients through a fixed-time discrete callback, nonlinear initialization, the final integration time, and requested saveat times. These checks do not establish differentiation support for every callback or a complete singular-mass DAE solve. In particular, a continuous callback with a parameter-dependent event time can produce incorrect gradients; event-time sensitivities need separate validation.