Skip to content

Avoid reevaluating model where unnecessary#2759

Merged
penelopeysm merged 2 commits intomainfrom
py/discard
Jan 28, 2026
Merged

Avoid reevaluating model where unnecessary#2759
penelopeysm merged 2 commits intomainfrom
py/discard

Conversation

@penelopeysm
Copy link
Member

@penelopeysm penelopeysm commented Jan 28, 2026

Closes #2639

Quoting from the changelog:

Avoid reevaluating the model on MCMC iterations where the transition is not saved to the chain (e.g. in initial burn-in, or when using thinning).
Also avoid each component sampler of Gibbs unnecessarily evaluating the model once per iteration.

This makes use of the discard_sample keyword argument to step / step_warmup TuringLang/AbstractMCMC.jl#194

Note that this is still type stable as the value of the discard_sample kwarg can be constant propagated (I've checked this with code_warntype etc).

@github-actions
Copy link
Contributor

Turing.jl documentation for PR #2759 is available at:
https://TuringLang.github.io/Turing.jl/previews/PR2759/

@penelopeysm
Copy link
Member Author

penelopeysm commented Jan 28, 2026

Assuming this passes CI, @joelkandiah, you might find this a little bit useful. It doesn't help with compilation, but here's a demonstration:

using Turing

modelevals = Ref(0)

@model function f()
    modelevals[] += 1
    a1 ~ Normal()
    a2 ~ Normal()
    a3 ~ Normal()
    a4 ~ Normal()
    a5 ~ Normal()
    a6 ~ Normal()
    a7 ~ Normal()
    a8 ~ Normal()
end

modelevals[] = 0
@time sample(f(), Gibbs(
        @varname(a1) => MH(),
        @varname(a2) => MH(),
        @varname(a3) => MH(),
        @varname(a4) => MH(),
        @varname(a5) => MH(),
        @varname(a6) => MH(),
        @varname(a7) => MH(),
        @varname(a8) => MH()
    ),
    1000;
    chain_type=Any, progress=false
);
@show modelevals[]

# this PR: 0.274875 seconds (2.26 M allocations: 169.478 MiB, 4.95% gc time)
# modelevals[] = 17003
#
# main: 0.354257 seconds (2.88 M allocations: 208.250 MiB, 4.81% gc time)
# modelevals[] = 25003

If you're using discard_initial and/or thinning, the benefits should scale accordingly. For example, adding thinning=10 to the above gives:

# With thinning=10
#
# this PR: 2.419929 seconds (21.36 M allocations: 1.582 GiB, 4.45% gc time)
# modelevals[] = 160859
#
# main: 3.395153 seconds (28.70 M allocations: 2.029 GiB, 4.05% gc time)
# modelevals[] = 249778

@codecov

This comment was marked as resolved.

@penelopeysm
Copy link
Member Author

penelopeysm commented Jan 28, 2026

(I definitely think it should in theory be possible to cut the model evaluations down to 8000, or maybe 9000 (each component sampler goes once and Gibbs itself has to do it once), but that really requires deeper refactoring that would follow on from #2756. The extra 8000 evaluations are coming from this setparams_varinfo!! stuff that complicates matters.)

@penelopeysm penelopeysm merged commit e487ed0 into main Jan 28, 2026
23 of 30 checks passed
@penelopeysm penelopeysm deleted the py/discard branch January 28, 2026 17:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Gibbs wastefully re-evaluates model once per component sampler per iteration

1 participant