Determining chunksize to use in ForwardDiff.jl for a Turing.jl model

benchmarking-turing.jl

julia> using Turing, TuringBenchmarking, ADTypes

julia> @model function f_model(n)
           μ ~ truncated(Normal(), lower=0)
           σ² ~ InverseGamma(2, 3)
           x ~ filldist(truncated(Normal(μ, sqrt(σ²)); lower=0), n)
       end

f_model (generic function with 2 methods)

julia> max_chunksizes = [
           0, # corresponds to the default heuristic
           1,
           4,
           8,
           16,
           32
       ];

julia> # Turing.jl now uses ADTypes.jl for the AD backend configuration.
       adbackends=map(c -> ADTypes.AutoForwardDiff(chunksize=c), chunksizes)
6-element Vector{AutoForwardDiff{_A, Nothing} where _A}:
 AutoForwardDiff{0, Nothing}(nothing)
 AutoForwardDiff{1, Nothing}(nothing)
 AutoForwardDiff{4, Nothing}(nothing)
 AutoForwardDiff{8, Nothing}(nothing)
 AutoForwardDiff{16, Nothing}(nothing)
 AutoForwardDiff{32, Nothing}(nothing)

julia> # Benchmark the model with different backends.
       results = TuringBenchmarking.benchmark_model(f_model(100); adbackends)
2-element BenchmarkTools.BenchmarkGroup:
  tags: []
  "evaluation" => 2-element BenchmarkTools.BenchmarkGroup:
          tags: []
          "linked" => Trial(4.727 μs)
          "standard" => Trial(1.906 μs)
  "gradient" => 6-element BenchmarkTools.BenchmarkGroup:
          tags: []
          "AutoForwardDiff{4, Nothing}(nothing)" => 2-element BenchmarkTools.BenchmarkGroup:
                  tags: ["ForwardDiff"]
                  "linked" => Trial(251.729 μs)
                  "standard" => Trial(123.806 μs)
          "AutoForwardDiff{32, Nothing}(nothing)" => 2-element BenchmarkTools.BenchmarkGroup:
                  tags: ["ForwardDiff"]
                  "linked" => Trial(89.686 μs)
                  "standard" => Trial(46.526 μs)
          "AutoForwardDiff{8, Nothing}(nothing)" => 2-element BenchmarkTools.BenchmarkGroup:
                  tags: ["ForwardDiff"]
                  "linked" => Trial(141.020 μs)
                  "standard" => Trial(74.930 μs)
          "AutoForwardDiff{16, Nothing}(nothing)" => 2-element BenchmarkTools.BenchmarkGroup:
                  tags: ["ForwardDiff"]
                  "linked" => Trial(93.164 μs)
                  "standard" => Trial(46.667 μs)
          "AutoForwardDiff{1, Nothing}(nothing)" => 2-element BenchmarkTools.BenchmarkGroup:
                  tags: ["ForwardDiff"]
                  "linked" => Trial(635.984 μs)
                  "standard" => Trial(241.602 μs)
          "AutoForwardDiff{0, Nothing}(nothing)" => 2-element BenchmarkTools.BenchmarkGroup:
                  tags: ["ForwardDiff"]
                  "linked" => Trial(110.319 μs)
                  "standard" => Trial(56.452 μs)

julia> # Extract the results for the gradient computation with linking.
       results_gradient = results[@tagged ("gradient" && "linked")]
1-element BenchmarkTools.BenchmarkGroup:
  tags: []
  "gradient" => 6-element BenchmarkTools.BenchmarkGroup:
          tags: []
          "AutoForwardDiff{4, Nothing}(nothing)" => 1-element BenchmarkTools.BenchmarkGroup:
                  tags: ["ForwardDiff"]
                  "linked" => Trial(251.729 μs)
          "AutoForwardDiff{32, Nothing}(nothing)" => 1-element BenchmarkTools.BenchmarkGroup:
                  tags: ["ForwardDiff"]
                  "linked" => Trial(89.686 μs)
          "AutoForwardDiff{8, Nothing}(nothing)" => 1-element BenchmarkTools.BenchmarkGroup:
                  tags: ["ForwardDiff"]
                  "linked" => Trial(141.020 μs)
          "AutoForwardDiff{16, Nothing}(nothing)" => 1-element BenchmarkTools.BenchmarkGroup:
                  tags: ["ForwardDiff"]
                  "linked" => Trial(93.164 μs)
          "AutoForwardDiff{1, Nothing}(nothing)" => 1-element BenchmarkTools.BenchmarkGroup:
                  tags: ["ForwardDiff"]
                  "linked" => Trial(635.984 μs)
          "AutoForwardDiff{0, Nothing}(nothing)" => 1-element BenchmarkTools.BenchmarkGroup:
                  tags: ["ForwardDiff"]
                  "linked" => Trial(110.319 μs)

julia> # Sort according to the minimum time for each backend.
       sort(
           BenchmarkTools.leaves(results_gradient),
           by=minimum ∘ last
       )
6-element Vector{Any}:
 (Any["gradient", "AutoForwardDiff{32, Nothing}(nothing)", "linked"], Trial(89.686 μs))
 (Any["gradient", "AutoForwardDiff{16, Nothing}(nothing)", "linked"], Trial(93.164 μs))
 (Any["gradient", "AutoForwardDiff{0, Nothing}(nothing)", "linked"], Trial(110.319 μs))
 (Any["gradient", "AutoForwardDiff{8, Nothing}(nothing)", "linked"], Trial(141.020 μs))
 (Any["gradient", "AutoForwardDiff{4, Nothing}(nothing)", "linked"], Trial(251.729 μs))
 (Any["gradient", "AutoForwardDiff{1, Nothing}(nothing)", "linked"], Trial(635.984 μs))