Repeated optimization (differentiation?) of a System keeps incurring type inference / compilation time

[cstjean]

The foo() calls below optimize an MTK model's parameters via Optimization.jl, ForwardDiff.jl and remake. The (compilation?) time and allocation varies roughly on an odd/even pattern (I've plotted 100 such calls and it was clear). foo() itself builds a System in a fairly vanilla way --- I do not use eval or staged function or anything particularly exotic.

Why do I keep seeing compilation time, and such fluctuating allocation patterns?

julia> @time foo();
 10.598025 seconds (329.36 M allocations: 6.384 GiB, 16.52% gc time, 99.43% compilation time: <1% of which was recompilation)

julia> @time foo();
 96.695468 seconds (383.82 M allocations: 18.683 GiB, 6.77% gc time, 195.60% compilation time: <1% of which was recompilation)

julia> @time foo();
 10.885435 seconds (332.19 M allocations: 6.768 GiB, 17.51% gc time, 82.27% compilation time: <1% of which was recompilation)

julia> @time foo();
 11.131420 seconds (335.78 M allocations: 6.774 GiB, 16.02% gc time, 94.54% compilation time: <1% of which was recompilation)

julia> @time foo();
 13.010147 seconds (333.03 M allocations: 6.633 GiB, 14.73% gc time, 95.50% compilation time: <1% of which was recompilation)

julia> @time foo();
  0.724270 seconds (9.34 M allocations: 353.154 MiB, 108.19% compilation time: <1% of which was recompilation)

julia> @time foo();
  0.786725 seconds (9.03 M allocations: 337.088 MiB, 23.32% gc time, 98.37% compilation time: <1% of which was recompilation)

julia> @time foo();
 11.860633 seconds (332.62 M allocations: 6.766 GiB, 15.40% gc time, 84.53% compilation time: <1% of which was recompilation)

julia> @time foo();
  0.625859 seconds (9.15 M allocations: 347.657 MiB, 91.62% compilation time: <1% of which was recompilation)

julia> @time foo();
 10.476637 seconds (322.76 M allocations: 6.061 GiB, 11.95% gc time, 0.00% compilation time: 100% of which was recompilation)

I've tried using --trace-compile, and indeed I've seen new methods pop up at each foo() call, stuff like

precompile(Tuple{typeof(Base.ntupleany), Base.var"#fieldtypes##0#fieldtypes##1"{Type{OrdinaryDiffEqRosenbrock.GRK4T{CS, AD, F, P, FDT, ST, CJ} where CJ where ST where FDT where P where F where AD where CS}}, Int64})
precompile(Tuple{typeof(Base.get!), Revise.var"#fieldtypes_cached##0#fieldtypes_cached##1"{Type{OrdinaryDiffEqRosenbrock.ROK4a{CS, AD, F, P, FDT, ST, CJ} where CJ where ST where FDT where P where F where AD where CS}}, Base.IdDict{Type, Union{Nothing, Array{Any, 1}}}, Any})
precompile(Tuple{typeof(Base.ntupleany), Base.var"#fieldtypes##0#fieldtypes##1"{Type{OrdinaryDiffEqRosenbrock.ROK4a{CS, AD, F, P, FDT, ST, CJ} where CJ where ST where FDT where P where F where AD where CS}}, Int64})

(this is with Revise, but I'm not touching anything, and I've reproduced it without Revise).

ProfileView.jl on the fast runs looks like this:

Slow runs look like this, with the topmost red bars being solve_up calls, that trigger type inference and abstract interpretation.

On MTK11, Julia 1.12.2.

Is this a known issue? Has anyone seen something like this? I can provide an MWE if it helps, but this looks like several hours of work.

[ChrisRackauckas]

@AayushSabharwal do we still have tests on v11 that the generated code has a constant hash?

[AayushSabharwal]

I don't know if we ever had such tests 😅

@cstjean could you share the implementation of foo? This kind of stuff could vary with how the system is built, or how the optimization loop is implemented.

[ChrisRackauckas]

It was in Symbolics IIRC. The property is that you have to get exactly the same Julia expression for the same symbolic expression, so that it hashes the same. If it isn't constant then it generates a new function and thus a new compile. We at least had that awhile back, that's why we removed all gensym usage. But I see a gensym here https://github.com/JuliaSymbolics/Symbolics.jl/blob/589359bb08609065dc01c2e09390c14f90eadb43/src/build_function.jl#L213 which is not good...

[AayushSabharwal]

We never hit that code path in MTK

[AayushSabharwal]

Okay, yeah SCCNonlinearProblem uses gensym

[cstjean]

Thank you for the quick answer. Looks worthwhile to reduce it, so I'll do my best on Friday.

[cstjean]

I've measured GC usage, and while the official GC time is way smaller than the spikes, disabling the GC seems to make the big spikes go away. Naturally, each subsequent run is a bit slower than the last.

Manually calling GC.gc() and running eventually triggers

Internal error: during type inference of
var"#solve!#14"(Base.Pairs{Symbol, Float64, Nothing, NamedTuple{(:reltol,), Tuple{Float64}}}, typeof(CommonSolve.solve!), LinearSolveForwardDiffExt.DualLinearCache{ForwardDiff.Dual{ForwardDiff.Tag{OptimizationBase.var"#75#76"{SciMLBase.OptimizationFunction{true, ADTypes.AutoForwardDiff{nothing, Nothing}, typeof(SimDenkai.optim_loss), Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SimDenkai.cons_fn), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, OptimizationBase.var"#_cons_oop#72"{SciMLBase.OptimizationFunction{true, ADTypes.AutoForwardDiff{nothing, Nothing}, typeof(SimDenkai.optim_
...

on Julia 1.12.2. This is possibly an OOM error?

In any case, upgrading to Julia 1.12.5 fixed the issue. With GC on, the spikes are now way smaller than before, and are clearly due to GC.

[cstjean]

I'm not 100% sure that this is the above issue, but I pinpointed something very very similar down to System's tstops. Basically I have some plain solve(ODEProblem) script where repeated execution yields:

julia> @time include("scratch.jl");
  5.763103 seconds (41.24 M allocations: 1.277 GiB, 5.68% gc time, 87.64% compilation time: <1% of which was recompilation)

julia> @time include("scratch.jl");
  5.498104 seconds (40.51 M allocations: 1.241 GiB, 5.77% gc time, 87.32% compilation time: <1% of which was recompilation)

julia> @time include("scratch.jl");
  5.500900 seconds (40.51 M allocations: 1.241 GiB, 4.93% gc time, 88.48% compilation time: <1% of which was recompilation)

Why all the recompilation? It's always the same system, there shouldn't be any compilation. BUT, with this one-line change in some @components definition:

-    return System(eqs, t, [], ps; name, systems=I_outputs, metadata, tstops)
+    return System(eqs, t, [], ps; name, systems=I_outputs, metadata)

julia> @time include("scratch.jl");
  0.091518 seconds (1.57 M allocations: 37.367 MiB, 0.00% compilation time: 53% of which was recompilation)

julia> @time include("scratch.jl");
  0.091346 seconds (1.57 M allocations: 37.365 MiB, 0.00% compilation time: 47% of which was recompilation)

julia> @time include("scratch.jl");
  0.093742 seconds (1.57 M allocations: 37.367 MiB, 0.00% compilation time: 59% of which was recompilation)

No more compilation at all! (All runtimes above also specify tstops in the solve call)

Not sure what to attribute it to. length(sol.t) ends up being 10X shorter when I use solve(; tstops) than when it's inside System(; tstops).

See also #4330

Rerunning my test suite went from taking minutes to seconds 😓