Nested submodels slow down Turing ~7x (Mooncake AD)

Specific Domains
1

I just discovered that nesting several submodels in Turing.jl seems to considerably slow down sampling (I only tried Mooncake.jl so far).

The code below shows this with the eight schools dataset. I’ve deliberately made the sampling more expensive than it has to be to better show the speed discrepancy.

On my machine, the non-nested-submodel variant comes in at 4-5s, while the nested-submodel variant takes 30-35s. Not using any submodels is on par with the non-nested version.
I ran these on four threads, and took timings on the second (and later) runs.

Is this expected? I wanted to ask here before opening an issue against Turing.jl (maybe @penelopeysm knows?). I briefly searched Turing.jl and Mooncake.jl issues but found nothing related. Apologies if I missed something!

The example here is quite contrived of course, but in my actual modelling scenario the nested submodels would be a nice simplification.

# Example adapted from 
# https://gist.github.com/penelopeysm/5656697ea20c94d80a285f5f6a69b8ab

using Turing, Mooncake, LinearAlgebra

# Eight schools data
J = 8
y = [28, 8, -3, 7, -1, 1, 18, 12]
sigma = [15, 10, 16, 11, 9, 11, 10, 18]

# Submodels

@model function tau_prior()
    p ~ truncated(Cauchy(0, 5); lower = 0)
    return (; p)
end

@model function mu_prior()
    p ~ Normal(0, 5)
    return (; p)
end

@model function z_prior(J)
    p ~ MvNormal(zeros(J), I)
    return (; p)
end

@model function theta_prior(J)
    mu ~ to_submodel(mu_prior())
    tau ~ to_submodel(tau_prior())
    z ~ to_submodel(z_prior(J))
    p := z.p .* tau.p .+ mu.p
    return (; p)
end

# Nested submodel

@model function esc_subm_nested(J, y, sigma)
    theta ~ to_submodel(theta_prior(J))
    for i in 1:J
        y[i] ~ Normal(theta.p[i], sigma[i])
    end
end

# Same sampler for both models, deliberately high adapt_delta
sampler = NUTS(10000, 0.99; adtype=AutoMooncake())

model1 = esc_subm_nested(J, y, sigma)
chn1 = sample(model1, sampler, MCMCThreads(), 5000, 8)

# Single submodel

@model function esc_subm(J, y, sigma)
    mu ~ to_submodel(mu_prior())
    tau ~ to_submodel(tau_prior())
    z ~ to_submodel(z_prior(J))
    theta := z.p .* tau.p .+ mu.p
    for i in 1:J
        y[i] ~ Normal(theta[i], sigma[i])
    end
end

model2 = esc_subm(J, y, sigma)
chn2 = sample(model2, sampler, MCMCThreads(), 5000, 8)

# No submodel

@model function esc(J, y, sigma)
    mu ~ Normal(0, 5)
    tau ~ truncated(Cauchy(0, 5); lower = 0)
    z ~ MvNormal(zeros(J), I)
    theta := z .* tau .+ mu
    for i in 1:J
        y[i] ~ Normal(theta[i], sigma[i])
    end
end

model3 = esc(J, y, sigma)
chn3 = sample(model3, sampler, MCMCThreads(), 5000, 8)

These are the package versions used above:

(nested_submodels_mooncake) pkg> st
Status `~/Documents/Snippets/nested_submodels_mooncake/Project.toml`
  [da2b9cff] Mooncake v0.5.32
  [fce5fe82] Turing v0.45.0
  [37e2e46d] LinearAlgebra v1.12.0
2

In the nested model you probably have a type instability which makes mooncake slow.

Can you try API · DynamicPPL ?

Can you try ForwardDiff to see if it’s also slow?

3

Probably the first thing I’d investigate is benchmarking the other AD backends just to see if the slowdown is a Turing problem or an AD problem or both. It’s been a while since I did anything Turing but the easiest way is with API · DynamicPPL

using DynamicPPL, Distributions, ADTypes, (ALL_YOUR_AD_BACKENDS...)
using DynamicPPL.TestUtils.AD: run_ad

@model function eightsch(...) ... end
model = eightsch(...)

for adtypes in (AutoForwardDiff(), ...)
    run_ad(model, adtype; test=false, benchmark=true)
end

In general my expectation would be that there should be some performance losses from a submodel but 7x is a bit excessive, especially if your model is already nontrivial (so the added cost of using a submodel should become less significant compared to actually doing the work).

4

Thanks for pointing me towards the benchmarking code! I get the following results for grad_time (edit: I have added results for primal time now, too):

Model ADType Grad Slowdown: grad Primal Slowdown: primal
Flat ForwardDiff 837.82 ns ref 201.20 ns ref
Flat ReverseDiff 54.17 μs ref 212.99 ns ref
Flat Enzyme 635.64 ns ref 210.85 ns ref
Flat Mooncake 1.57 μs ref 203.02 ns ref
Submodel ForwardDiff 977.77 ns 1.2x 193.69 ns 1.0x
Submodel ReverseDiff 56.10 μs 1.0x 201.15 ns 0.9x
Submodel Enzyme 632.26 ns 1.0x 196.35 ns 0.9x
Submodel Mooncake 1.73 μs 1.1x 208.90 ns 1.0x
Nested submodel ForwardDiff 4.56 μs 5.4x 2.84 μs 14.1x
Nested submodel Enzyme 52.04 μs 81.9x 3.15 μs 14.9x
Nested submodel ReverseDiff 55.17 μs 1.0x 2.87 μs 13.5x
Nested submodel Mooncake 51.04 μs 32.5x 3.13 μs 15.4x

So it appears that all backends except for ReverseDiff are affected, but ForwardDiff the least. I assume that the slowdown being so much higher for Mooncake than what I see in my full MCMC runs is because of the other (more or less) nontrivial stuff that’s going on in the model, as Penny said.

The model that I initially discovered this issue with is much more complex – the runtime there went from < 10 min to slightly above 2 hours. So at least in this case the impact of the nested submodels didn’t shrink with growing model complexity.

5

That is kind of weird, I’m surprised that that’s the case. Some points:

  • Is the primal time (approximately) unchanged by introduction of submodels?
  • ReverseDiff is probably ‘unaffected’ because it’s already slow from the beginning. If you’re just looking for performance, it might be worth trying with compile=true (as long as there’s no control flow that depends on the value of the parameters).
  • As a band-aid you could try to stick @inline just inside every submodel you declare to see if that makes any difference:
    @model function tau_prior()
    @inline
    p ~ truncated(Cauchy(0, 5); lower = 0)
    return (; p)
end
    • It might well be that some definitions of DynamicPPL.evaluate!! or DynamicPPL._evaluate!! also could benefit from being inlined.
    • In general, my experience was that the Julia compiler was good enough at inferring when inlining was appropriate, but the inlining decisions are not made with AD in mind – only the primal, so maybe to get good AD performance it’s useful to hint at it a bit.
    • The other thing is that at the user level (i.e. without going deeply into DynamicPPL code) this is pretty much the only tool you have to try to affect things, so may as well try it…
  • I agree type instabilities could be an issue - the eight-schools example above doesn’t look like it should pose a problem but it would be worth checking for the real-world model.

It’s very hard to say whether this is a Turing or AD issue per se because it could be that Turing is generating code that’s very hard for AD to differentiate through, or it could be that Turing is generating perfectly fine code and AD should be better at differentiating through it. Figuring that out would require more in-depth investigation. But I’m not working on Turing any more, so I’ll leave it there…

6

Thank you for your input!

I have run the benchmarks again and added primal time to the table above. It’s unchanged for the basic submodel version, but roughly 15x slower for all AD backends in the nested submodel version.

I’ll have a look at the impact of adding @inline next! :slight_smile: I have now added @inline to each submodel as the first line after @model function .... If this is how you had intended it, it unfortunately didn’t solve the issue; the timings are approximately the same.

I’ve opened an issue Nested submodels slow down sampling · Issue #2844 · TuringLang/Turing.jl · GitHub.