Performance Tips

This section briefly summarises a few common techniques to ensure good performance when using Turing. We refer to the Julia documentation for general techniques to ensure good performance of Julia programs.

Benchmark and profile your model

Before changing your model for performance reasons, it is worth measuring where the time actually goes. When sampling with a gradient-based method such as NUTS, almost all of the time is spent evaluating the log density of the model and its gradient. You can time both with DynamicPPL.TestUtils.AD.run_ad:

using Turing
using ADTypes: AutoReverseDiff
using DynamicPPL.TestUtils.AD: run_ad
using ReverseDiff

@model function bmodel(x)
    m ~ Normal()
    s ~ truncated(Normal(); lower=0)
    return x ~ MvNormal(fill(m, length(x)), s^2 * I)
end

result = run_ad(bmodel(randn(100)), AutoReverseDiff(); benchmark=true)
result.primal_time, result.grad_time
[ Info: Running AD on bmodel with ADTypes.AutoReverseDiff()
       params : [-2.177290315677856, -0.8106415275776712]
       actual : (-1367.2824392853547, [1072.4396815742316, 2604.858851605788])
     expected : (-1367.2824392853547, [1072.439681574232, 2604.858851605788])
   evaluation : 214.964 ns (4 allocs: 1.812 KiB)
     gradient : 91.986 μs (1681 allocs: 64.359 KiB)
  grad / eval : 427.9

(2.1496428571428573e-7, 9.198600000000001e-5)

Both times are median seconds per evaluation. If the gradient takes many times longer than the log density itself, the AD backend is the thing to change: see Choose your AD backend below, and the Automatic Differentiation page for how to compare several backends on your model. If the log density itself is slow, the model is the thing to change, and the remaining sections on this page apply.

To find out where the time goes within a single evaluation, you can use Julia’s built-in profiler on the model’s log density:

using DynamicPPL, LogDensityProblems, Profile

model = bmodel(randn(100))
ldf = LogDensityFunction(model, getlogjoint_internal, LinkAll(); adtype=AutoReverseDiff())
x = rand(ldf)

Profile.@profile for _ in 1:10_000
    LogDensityProblems.logdensity_and_gradient(ldf, x)
end
Profile.print(; mincount=10)

This profiles the model evaluation and the gradient computation together, which is what a sampler actually runs. We do not run the profiler on this page because its output is long and machine-specific: Profile.print shows a tree of function calls, each annotated with the number of samples in which it appeared, and the hot spots are the lines with the largest counts (mincount=10 hides rarely-seen calls). Graphical viewers such as @profview in VS Code, ProfileView.jl, or PProf.jl make the output much easier to read than Profile.print.

Use multivariate distributions

It is generally preferable to use multivariate distributions if possible.

The following example:

using Turing
@model function gmodel(x)
    m ~ Normal()
    for i in eachindex(x)
        x[i] ~ Normal(m, 0.2)
    end
end
gmodel (generic function with 2 methods)

can be directly expressed more efficiently with a simple transformation:

using FillArrays

@model function gmodel(x)
    m ~ Normal()
    return x ~ MvNormal(Fill(m, length(x)), 0.04 * I)
end
gmodel (generic function with 2 methods)

Choose your AD backend

Automatic differentiation (AD) makes it possible to use modern, efficient gradient-based samplers like NUTS and HMC. This, however, also means that using a performant AD system is incredibly important. Turing currently supports several AD backends, including ForwardDiff (the default), Mooncake, and ReverseDiff.

For many common types of models, the default ForwardDiff backend performs well, and there is no need to worry about changing it. However, if you need more speed, you can try different backends via the standard ADTypes interface by passing an AbstractADType to the sampler with the optional adtype argument, e.g. NUTS(; adtype = AutoMooncake()).

Generally, adtype = AutoForwardDiff() is likely to be the fastest and most reliable for models with few parameters (say, less than 20 or so), while reverse-mode backends such as AutoMooncake() or AutoReverseDiff() will perform better for models with many parameters or linear algebra operations. If in doubt, you can benchmark your model with different backends to see which one performs best. See the Automatic Differentiation page for details.

Special care for ReverseDiff with a compiled tape

For large models, the fastest option is often ReverseDiff with a compiled tape, specified as adtype=AutoReverseDiff(; compile=true). However, it is important to note that if your model contains any branching code, such as if-else statements, the gradients from a compiled tape may be inaccurate, leading to erroneous results. If you use this option for the (considerable) speedup it can provide, make sure to check your code for branching and ensure that it does not affect the gradients. It is also a good idea to verify your gradients with another backend.

Note that compile=true is not supported for variational inference: passing AutoReverseDiff(; compile=true) to vi raises an error, because a compiled tape can silently give incorrect gradients when reused across optimisation steps. For VI, use AutoReverseDiff(; compile=false) or a different reverse-mode backend such as AutoMooncake().

Ensure that types in your model can be inferred

For efficient gradient-based inference, e.g. using HMC, NUTS or ADVI, it is important to ensure the types in your model can be inferred.

The following example with abstract types

@model function tmodel(x, y)
    p, n = size(x)
    params = Vector{Real}(undef, n)
    for i in 1:n
        params[i] ~ truncated(Normal(); lower=0)
    end

    a = x * params
    return y ~ MvNormal(a, I)
end
tmodel (generic function with 2 methods)

can be transformed into the following representation with concrete types:

@model function tmodel(x, y, ::Type{T}=Float64) where {T}
    p, n = size(x)
    params = Vector{T}(undef, n)
    for i in 1:n
        params[i] ~ truncated(Normal(); lower=0)
    end

    a = x * params
    return y ~ MvNormal(a, I)
end
tmodel (generic function with 4 methods)

Alternatively, you could use filldist in this example:

@model function tmodel(x, y)
    params ~ filldist(truncated(Normal(); lower=0), size(x, 2))
    a = x * params
    return y ~ MvNormal(a, I)
end
tmodel (generic function with 4 methods)

You can use DynamicPPL’s debugging utilities to find types in your model definition that the compiler cannot infer. These will be marked in red in the Julia REPL (much like when using the @code_warntype macro).

For example, consider the following model:

@model function tmodel(x)
    p = Vector{Real}(undef, 1)
    p[1] ~ Normal()
    p = p .+ 1
    return x ~ Normal(p[1])
end
tmodel (generic function with 6 methods)

Because the element type of p is an abstract type (Real), the compiler cannot infer a concrete type for p[1]. To detect this, we can use

model = tmodel(1.0)

using DynamicPPL
DynamicPPL.DebugUtils.model_warntype(model)

In this particular model, the following call to getindex should be highlighted in red (the exact numbers may vary):

[...]
│    %120 = p::AbstractVector
│    %121 = Base.getindex(%120, 1)::Any
[...]
Back to top