From 9acd0eb4376c90c53a43e216829173f98099fe42 Mon Sep 17 00:00:00 2001 From: Markus Hauru Date: Thu, 6 Mar 2025 16:35:06 +0000 Subject: [PATCH 1/4] Make benchmarks not depend on TuringBenchmarking.jl --- benchmarks/Project.toml | 5 +- benchmarks/benchmarks.jl | 18 ++----- benchmarks/src/DynamicPPLBenchmarks.jl | 67 +++++++++++++++++++++----- benchmarks/src/Models.jl | 2 +- 4 files changed, 64 insertions(+), 28 deletions(-) diff --git a/benchmarks/Project.toml b/benchmarks/Project.toml index fab8aee23..0eafc1c12 100644 --- a/benchmarks/Project.toml +++ b/benchmarks/Project.toml @@ -3,9 +3,12 @@ uuid = "d94a1522-c11e-44a7-981a-42bf5dc1a001" version = "0.1.0" [deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d" -TuringBenchmarking = "0db1332d-5c25-4deb-809f-459bc696f94f" +ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" diff --git a/benchmarks/benchmarks.jl b/benchmarks/benchmarks.jl index 359026db6..52c84ff77 100644 --- a/benchmarks/benchmarks.jl +++ b/benchmarks/benchmarks.jl @@ -62,22 +62,12 @@ end results_table = Tuple{String,String,String,Bool,Float64,Float64}[] for (model_name, model, varinfo_choice, adbackend, islinked) in chosen_combinations - suite = make_suite(model, varinfo_choice, adbackend) + suite = make_suite(model, varinfo_choice, adbackend, islinked) results = run(suite) - result_key = islinked ? "linked" : "standard" - - eval_time = median(results["evaluation"][result_key]).time + eval_time = median(results["evaluation"]).time relative_eval_time = eval_time / reference_time - - grad_group = results["gradient"] - if isempty(grad_group) - relative_ad_eval_time = NaN - else - grad_backend_key = first(keys(grad_group)) - ad_eval_time = median(grad_group[grad_backend_key][result_key]).time - relative_ad_eval_time = ad_eval_time / eval_time - end - + ad_eval_time = median(results["gradient"]).time + relative_ad_eval_time = ad_eval_time / eval_time push!( results_table, ( diff --git a/benchmarks/src/DynamicPPLBenchmarks.jl b/benchmarks/src/DynamicPPLBenchmarks.jl index f98aede2b..774891b07 100644 --- a/benchmarks/src/DynamicPPLBenchmarks.jl +++ b/benchmarks/src/DynamicPPLBenchmarks.jl @@ -1,16 +1,39 @@ module DynamicPPLBenchmarks using DynamicPPL: VarInfo, SimpleVarInfo, VarName -using BenchmarkTools: BenchmarkGroup -using TuringBenchmarking: make_turing_suite +using BenchmarkTools: BenchmarkGroup, @benchmarkable +using DynamicPPL: DynamicPPL +using ADTypes: ADTypes +using LogDensityProblems: LogDensityProblems + +# Load some the default backends to trigger conditional loading. +using ForwardDiff: ForwardDiff +using ReverseDiff: ReverseDiff include("./Models.jl") using .Models: Models export Models, make_suite +# Utility functions for representing AD backends using symbols. +# Copied from TuringBenchmarking.jl. +const SYMBOL_TO_BACKEND = Dict( + :forwarddiff => ADTypes.AutoForwardDiff(; chunksize=0), + :reversediff => ADTypes.AutoReverseDiff(; compile=false), + :reversediff_compiled => ADTypes.AutoReverseDiff(; compile=true), + :mooncake => ADTypes.AutoMooncake(; config=nothing), +) + +to_backend(x) = error("Unknown backend: $x") +to_backend(x::ADTypes.AbstractADType) = x +function to_backend(x::Union{AbstractString,Symbol}) + k = Symbol(lowercase(string(x))) + haskey(SYMBOL_TO_BACKEND, k) || error("Unknown backend: $x") + return SYMBOL_TO_BACKEND[k] +end + """ - make_suite(model, varinfo_choice::Symbol, adbackend::Symbol) + make_suite(model, varinfo_choice::Symbol, adbackend::Symbol, islinked::Bool) Create a benchmark suite for `model` using the selected varinfo type and AD backend. Available varinfo choices: @@ -20,8 +43,10 @@ Available varinfo choices: • `:simple_dict` → builds a `SimpleVarInfo{Float64}` from a Dict (pre-populated with the model’s outputs) The AD backend should be specified as a Symbol (e.g. `:forwarddiff`, `:reversediff`, `:zygote`). + +`islinked` determines whether to link the VarInfo for evaluation. """ -function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol) +function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol, islinked::Bool) suite = BenchmarkGroup() vi = if varinfo_choice == :untyped @@ -40,14 +65,32 @@ function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol) error("Unknown varinfo choice: $varinfo_choice") end - # Add the AD benchmarking suite. - suite = make_turing_suite( - model; - adbackends=[adbackend], - varinfo=vi, - check_grads=true, - error_on_failed_backend=true, - ) + adbackend = to_backend(adbackend) + context = DynamicPPL.DefaultContext() + + if islinked + vi = DynamicPPL.link(vi, model) + end + + # We construct `LogDensityFunction` using different values + # than the ones we're going to use for the test. Some of the AD backends + # compile the tape upon `LogDensityFunction` construction, and we want to + # evaluate using inputs different from those that the tape was compiled for. + f = DynamicPPL.LogDensityFunction(model, vi, context; adtype=adbackend) + + # The parameters at which we evaluate f. + θ = if islinked + randn(length(vi[:])) + else + rand(Vector, model) + end + + # Run once to trigger compilation. + LogDensityProblems.logdensity_and_gradient(f, θ) + suite["gradient"] = @benchmarkable $(LogDensityProblems.logdensity_and_gradient)($f, $θ) + + # Also benchmark just standard model evaluation because why not. + suite["evaluation"] = @benchmarkable $(DynamicPPL.evaluate!!)($model, $vi, $context) return suite end diff --git a/benchmarks/src/Models.jl b/benchmarks/src/Models.jl index 25d90f5e1..589af3039 100644 --- a/benchmarks/src/Models.jl +++ b/benchmarks/src/Models.jl @@ -17,7 +17,7 @@ using Distributions: logpdf, product_distribution, truncated -using DynamicPPL: @model, to_submodel +using DynamicPPL: DynamicPPL, @model, to_submodel using LinearAlgebra: cholesky export simple_assume_observe_non_model, From cfd781c741bdd082c4f69e7d28bc9ebea118a43d Mon Sep 17 00:00:00 2001 From: Markus Hauru Date: Thu, 6 Mar 2025 16:40:24 +0000 Subject: [PATCH 2/4] Make benchmarks.jl dev the local DDPPL version --- benchmarks/benchmarks.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/benchmarks/benchmarks.jl b/benchmarks/benchmarks.jl index 52c84ff77..d13255695 100644 --- a/benchmarks/benchmarks.jl +++ b/benchmarks/benchmarks.jl @@ -1,3 +1,7 @@ +using Pkg +# To ensure we benchmark the local version of DynamicPPL, dev the folder above. +Pkg.develop(; path=joinpath(@__DIR__, "..")) + using DynamicPPLBenchmarks: Models, make_suite using BenchmarkTools: @benchmark, median, run using PrettyTables: PrettyTables, ft_printf From 5f5a3d2402db564ca6d720fa0158a18e58ab1d14 Mon Sep 17 00:00:00 2001 From: Markus Hauru Date: Thu, 6 Mar 2025 16:45:10 +0000 Subject: [PATCH 3/4] Add benchmarks compat bounds --- benchmarks/Project.toml | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/benchmarks/Project.toml b/benchmarks/Project.toml index 0eafc1c12..14b8a39b2 100644 --- a/benchmarks/Project.toml +++ b/benchmarks/Project.toml @@ -12,3 +12,13 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" + +[compat] +ADTypes = "1.14.0" +BenchmarkTools = "1.6.0" +Distributions = "0.25.117" +# DynamicPPL deliberately has no compat bound, to allow using the local version. +ForwardDiff = "0.10.38" +LogDensityProblems = "2.1.2" +PrettyTables = "2.4.0" +ReverseDiff = "1.15.3" From 8d6e7d4809433225aefa5ca950620f97278f7c5a Mon Sep 17 00:00:00 2001 From: Markus Hauru Date: Thu, 6 Mar 2025 17:14:12 +0000 Subject: [PATCH 4/4] Use ForwardDiff with dynamic benchmark model --- benchmarks/benchmarks.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/benchmarks/benchmarks.jl b/benchmarks/benchmarks.jl index d13255695..265ba6f39 100644 --- a/benchmarks/benchmarks.jl +++ b/benchmarks/benchmarks.jl @@ -51,7 +51,9 @@ chosen_combinations = [ ("Multivariate 1k", multivariate1k, :typed, :reversediff, true), ("Loop univariate 10k", loop_univariate10k, :typed, :reversediff, true), ("Multivariate 10k", multivariate10k, :typed, :reversediff, true), - ("Dynamic", Models.dynamic(), :typed, :reversediff, true), + # TODO(mhauru) Would like to use :reversediff here, but see + # https://github.com/TuringLang/DynamicPPL.jl/issues/835 + ("Dynamic", Models.dynamic(), :typed, :forwarddiff, true), ("Submodel", Models.parent(randn()), :typed, :reversediff, true), ("LDA", lda_instance, :typed, :reversediff, true), ]