Mooncake performance

[SamuelBrand1]

I saw there was a general thread for Zygote performance, so it seemed logical to have one for mooncake.

Running locally (Mac with M1 cores, julia 1.11.1) the linear regression example https://turinglang.org/docs/tutorials/05-linear-regression/ takes about x15 the time compared to either AutoForwardDiff or AutoReverseDiff(compile=true).

I'm happy to do a MWE on this but its simply the tutorial code with extra ReverseDiff and Mooncake deps and a different adtype kwarg passed to NUTS constructor.

Any immediate ideas about where the performance hit it happening here?

[torfjelde]

@willtebbutt !

I also encountered this, and I believe we concluded it was likely due to a lack of custom rules for Mooncake.jl (these are currently somewhat annoying and life, for the most part, in DistributionsAD.jl). So technically, nothing should be stopping us from achieving the same perf here. Would be useful to confirm this though by implementing the corresponding AD rules.

[SamuelBrand1]

Btw, just to clarify @SamuelBrand1 , when you say "15 slower", are you talking about benchmarking the model with the AD backends or are you talking about the sample call itself?

The whole sample call (on second run to avoid compilation time).

[SamuelBrand1]

Ah, right, so I think this is actually TuringLang/DynamicPPL.jl#727 issue -- in short Turing.jl isn't currently caching rules properly, which will cause really quite poor performance in this case. So, sadly, I think I accidentally told an untruth in the chat the other day @SamuelBrand1 -- Mooncake.jl is almost ready for use with Turing.jl. models via AutoMooncake.

I thought this was happening but didn't have time to burrow in.

[SamuelBrand1]

Note that what Tor says about things from DistributionsAD.jl. is correct (Mooncake.jl doesn't use it at all), however, it is also true that it doesn' have to. In short, I have no intention at all of incorporating rules from DistributionsAD.jl.

Yes, this was the impression I was getting and super happy its going this way!

[seabbs]

Any interest in us throwing a few simple models together to use in benchmarking or some co-creation?

[willtebbutt]

Very interested. Did you have anything in particular in mind?