Within-chain parallelization with Turing.jl

Specific Domains Probabilistic programming
1

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,
Nicolo’

2

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

3

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.

4

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,
Nicolo’

5

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))
    end
end

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
6

Yes! Thank you so much!

1 Like
7

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.