computation time of `returned` scales with number of model parameters

[tiemvanderdeure]

MWE:

using Turing
@model function mymodel(n)
    x ~ filldist(Normal(), n)
    return x[1]
end

for n in [10, 100, 1000, 10_000]
    mod = mymodel(n)
    chn = sample(mod, Prior(), 10)
    @time returned(mod, chn)
end

gives

  0.000808 seconds (2.95 k allocations: 129.688 KiB)
  0.003559 seconds (14.14 k allocations: 1.431 MiB)
  0.114613 seconds (146.33 k allocations: 81.974 MiB)
  3.866223 seconds (1.66 M allocations: 7.510 GiB, 28.82% gc time)

For models with many parameters returned is basically unusable because of this.

Profview shows that all the computation time is happening in call to DynamicPPL.setval_and_resample! inside the loop that has a TODO

DynamicPPL.jl/ext/DynamicPPLMCMCChainsExt.jl

Lines 251 to 254 in ea6b6de

	# TODO: Use `fix` once we've addressed https://github.com/TuringLang/DynamicPPL.jl/issues/702.
	# Update the varinfo with the current sample and make variables not present in `chain`
	# to be sampled.
	DynamicPPL.setval_and_resample!(varinfo, chain, sample_idx, chain_idx)

[penelopeysm]

So setval_and_resample! is on its way out: #984

using DynamicPPL, MCMCChains, Distributions, DistributionsAD, Random, OrderedCollections

# Note we are not importing Turing because there's no version of Turing that works
# with the DynamicPPL PR -- so we need this function to make the chain -- it does
# exactly the same thing as `sample(model, Prior(), n)`
# See DynamicPPL's `test/test_util.jl` for more info
function make_chain_from_prior(rng::Random.AbstractRNG, model::Model, n_iters::Int)
    varinfos = [VarInfo(rng, model) for _ in 1:n_iters]
    varnames = OrderedSet{VarName}()
    dicts = map(varinfos) do t
        vals = DynamicPPL.values_as(t, OrderedDict)
        iters = map(DynamicPPL.varname_and_value_leaves, keys(vals), values(vals))
        tuples = mapreduce(collect, vcat, iters)
        push!(varnames, map(first, tuples)...)
        OrderedDict(tuples)
    end
    varnames = collect(varnames)
    vals = [get(dict, vn, missing) for dict in dicts, vn in varnames]
    vn_to_sym_dict = Dict(zip(varnames, map(Symbol, varnames)))
    return Chains(vals, varnames; info=(; varname_to_symbol=vn_to_sym_dict))
end

@model function mymodel(n)
    x ~ filldist(Normal(), n)
    return x[1]
end

for n in [10, 100, 1000, 10_000]
    mod = mymodel(n)
    chn = make_chain_from_prior(Random.default_rng(), mod, 10)
    @time returned(mod, chn)
end

Good news: if you check out the code in #984 and run the above, this looks somewhat better for N = 10, 100, 1000

 0.000156 seconds (1.04 k allocations: 49.047 KiB)
  0.000441 seconds (5.37 k allocations: 305.734 KiB)
  0.006227 seconds (57.09 k allocations: 3.244 MiB)

Bad news: for N = 10000

ERROR: StackOverflowError:
Stacktrace:
  [1] getsym
    @ ~/.julia/packages/AbstractPPL/jn19m/src/varname.jl:136 [inlined]
  [2] ==
    @ ~/.julia/packages/AbstractPPL/jn19m/src/varname.jl:191 [inlined]
  [3] isequal
    @ ./operators.jl:134 [inlined]
  [4] ht_keyindex2(h::OrderedDict{VarName, Nothing}, key::VarName{:x, Accessors.IndexLens{Tuple{Int64}}})
    @ OrderedCollections ~/.julia/packages/OrderedCollections/Xihhq/src/ordered_dict.jl:283
  [5] setindex!
    @ ~/.julia/packages/OrderedCollections/Xihhq/src/ordered_dict.jl:334 [inlined]
  [6] push!
    @ ~/.julia/packages/OrderedCollections/Xihhq/src/ordered_set.jl:40 [inlined]
  [7] push!(A::OrderedSet{VarName}, a::VarName{:x, Accessors.IndexLens{…}}, b::VarName{:x, Accessors.IndexLens{…}})
    @ Base ./abstractarray.jl:3536
  [8] push!(::OrderedSet{…}, ::VarName{…}, ::VarName{…}, ::VarName{…}, ::VarName{…}, ::Vararg{…}) (repeats 4523 times)
    @ Base ./abstractarray.jl:3537
  [9] (::var"#8#11")(t::VarInfo{@NamedTuple{x::DynamicPPL.Metadata{…}}, DynamicPPL.AccumulatorTuple{3, @NamedTuple{…}}})
    @ Main ./REPL[9]:14
 [10] iterate
    @ ./generator.jl:48 [inlined]
 [11] _collect(c::Vector{VarInfo{…}}, itr::Base.Generator{Vector{…}, var"#8#11"}, ::Base.EltypeUnknown, isz::Base.HasShape{1})
    @ Base ./array.jl:811
 [12] collect_similar(cont::Vector{VarInfo{…}}, itr::Base.Generator{Vector{…}, var"#8#11"})
    @ Base ./array.jl:720
 [13] map(f::Function, A::Vector{VarInfo{@NamedTuple{…}, DynamicPPL.AccumulatorTuple{…}}})
    @ Base ./abstractarray.jl:3371
 [14] make_chain_from_prior(rng::TaskLocalRNG, model::Model{…}, n_iters::Int64)
    @ Main ./REPL[9]:10
 [15] top-level scope
    @ ./REPL[12]:3
Some type information was truncated. Use `show(err)` to see complete types.

which should probably be fixed somewhere.

[penelopeysm]

More generally though, I'm curious what kind of behaviour do you expect?

I don't think it's unsurprising that the time taken to calculate returned scales somehow with the number of parameters. After all, returned has to evaluate the model to get the return value and if the model is larger it's going to take more time. Even if we stored x as a single vector (which unfortunately we don't........ this is a fundamental flaw with MCMCChains and is not fixable in the short term 😭), the time it takes to calculate logpdf of filldist(Normal(), n) with respect to x is going to scale with the length of x.

Of course in this example you're only returning x[1] so you could argue that there's no need to evaluate the other N - 1 elements of x, but in general DynamicPPL has no way of knowing that, it doesn't analyse the return value to see which values it needs.

So I think at the very least returned has to take O(n) time.

That's not to say that the current behaviour is perfect, because

# run this on DynamicPPL main
# other imports and `make_chain_from_prior` are in the previous comment

@model function mymodel(n)
    x ~ filldist(Normal(), n)
    return x[1]
end

using Chairmarks, Statistics, Plots
ns = map(x -> round(Int, x), logrange(1, 1000, length=30))
ts = Float64[]
for n in ns
    mod = mymodel(n)
    chn = make_chain_from_prior(Random.default_rng(), mod, 10)
    be = @be returned(mod, chn)
    t = median(be).time
    push!(ts, t)
end

scatter(log.(ns), log.(ts), xlabel="log(n)", ylabel="log(time)")

The slope towards the right is ~ 2. On #984 the same thing happens but just with a much smaller constant factor.

[tiemvanderdeure]

I haven't thought very deeply about this. I just had to abort a julia session because returned on my 50 thousand parameter model was taking forever. Clearly that isn't great - so I thought I'd make an issue to track this.

I'm curious what kind of behaviour do you expect?

Good question. I guess this scaling has to do with the way parameters are stored in the chain which I don't understand very much. But it seems to me that there is something in the inner loop that could be in the outer loop? Basically all of the time is spent on getindex, not on any actual computation. Is that some inefficient dict lookup?

the time it takes to calculate logpdf of filldist(Normal(), n) with respect to x is going to scale with the length of x.

But is returned calculating logpdf? That's surprising to me. I would have though that every ~ statement doesn't have to be evaluated as such because the values are stored in the chain.

So I think at the very least returned has to take O(n) time.

If it did take O(n) time, it wouldn't be that bad. The problems come from whatever scales much worse than that

[penelopeysm]

But it seems to me that there is something in the inner loop that could be in the outer loop? Basically all of the time is spent on getindex, not on any actual computation. Is that some inefficient dict lookup?

Perhaps. I'm probably not going to investigate the performance on v0.36 (current release) or v0.37 (to be released tomorrow), since that entire code bit will be rewritten anyway in v0.38 (the PR quite literally already exists #984). But I will check the issue on v0.38.

But is returned calculating logpdf? That's surprising to me.

Yes, it does. Example:

@model function f()
    x ~ Normal()
    return DynamicPPL.getlogp(__varinfo__) < -1.0
end

If returned didn't calculate logp, then that would either error or return wrong values. What we might be able to do is to put in a switch that allows you to not calculate logp based on some knowledge of your model (for example if you don't call getlogp), but we can't globally disable logp accumulation in the implementation of returned because we don't know what kinds of models people are writing.

If it did take O(n) time, it wouldn't be that bad. The problems come from whatever scales much worse than that

Yes, definitely agreed. I just wanted to make sure we were on the same page and you weren't expecting, say, O(1) :)

[penelopeysm]

Okay, on #984 a model with 50000 parameters takes about 0.6 seconds per iteration on my laptop. Still a bit slow, but still a lot better than the current situation. I don't really know how to bring it down at this point in time, but some observations:

• Only ~ 10% of the time is on actually executing the model. In fact even the model execution on its own still looks like O(N^2), which is probably still not ideal, but that alone might at least be somewhat acceptable, because essentially it's 0.06 seconds per iteration, so if you have 10000 iterations it's 10 minutes.

• The remaining 90% of the time is spent constructing the dictionary of (varname => value) from the chain, which is honestly a bit ridiculous. But I'm not sure how to fix this unless we rewrite MCMCChains (something which I've consistently argued for).

[penelopeysm]

But I'm not sure how to fix this unless we rewrite MCMCChains (something which I've consistently argued for).

I tested the code in #984 together with my half-baked rewrite of MCMCChains.jl, and the performance is MUCH better. Here's a million parameters, and a thousand iterations, and the whole thing takes just one second:

julia> using Turing, DynamicPPL, FlexiChains

julia> @model function mymodel(n)
           x ~ filldist(Normal(), n)
           return x[1]
       end
mymodel (generic function with 2 methods)

julia> model = mymodel(1000000)
Model{typeof(mymodel), (:n,), (), (), Tuple{Int64}, Tuple{}, DefaultContext}(mymodel, (n = 1000000,), NamedTuple(), DefaultContext())

julia> fc = sample(model, Prior(), 1000; chain_type=FlexiChain{VarName})
Sampling 100%|████████████████████████████████████████████| Time: 0:00:14
FlexiChain (1000 iterations, 1 chain)

Parameter type   VarName
Parameters       x
Other keys       {:stats} logjoint, loglikelihood, logprior

julia> @time DynamicPPL.returned(model, fc)
  1.036771 seconds (1.18 M allocations: 15.029 GiB, 11.76% gc time, 5.52% compilation time)

(Code run using current main of FlexiChains 93eb7d0, Turing #2550 (5743ff7), and DynamicPPL current main 084f5af)

But I don't know when this will all be released, sorry :(

[tiemvanderdeure]

Thanks for investigating and great to know that this can be much faster. I understand weeding out the problems with MCMCChains is not easy - but I'm excited about your rewrite :)