Within-chain parallelization with Turing.jl

Hi everyone,

I’m a new Turing user and I read that there is a possible way to parallelize the MCMC of a single chain on multiple threads but I’m not able to do that. Do you have any references that I can use to have a better understanding of that?

Thank you,

Could you upload your attempt so that we could take a look into what’s not working?

You should be able to utilize whatever @threads and @spawn and other parallel features of Julia from within a model. There is no way I know of to automatically parallelize a chain in some way. It is possible to run several chains in parallel easily though.

I was trying to:
sample(model, NUTS(), MCMCDistributed(), 50000, 1)
using the @everywhere and I also tried:
sample(model, NUTS(), MCMCThreads(), 50000, 1)

None of the two approaches seemed working for the single chain while they are working well running multiple chains in parallel.

Thank you,

The follwing is probably more like what you’re referring to:

@model function demo_threading(x)
    s ~ InverseGamma(2, 3)
    m ~ Normal(0, sqrt(s))

    # This will then be parallelized over the available threads.
    Threads.@threads for i in eachindex(x)
        x[i] ~ Normal(m, sqrt(s))

Note that this only works for observe-statements, i.e. when the LHS of ~ is “fixed” / not random (in the above example x is considered an observation because it’s part of the model arguments).

1 Like

Yes! Thank you so much!

1 Like

Note that this is also mainly just compatible with ForwardDiff.jl in case you’re using samplers which require differentiation :confused: Unfortunately reverse-mode AD backends aren’t threadsafe (when used in this “naive” way).

Also, you can use API · DynamicPPL! to add log-prob computations performed by hand.