TuringLang / TuringLang/DynamicPPL.jl

check_model=true (the default) silently consumes draws from the caller's RNG, breaking manual per-chain seed reproduction

Open
#1,466 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Julia
Stars
286
Forks
41
Avg merge
1d 11h
Merged PRs (30d)
34

Description

What happened

We build one MCMC chain at a time, in separate OS processes (a Snakemake-per-chain-job pattern, similar to how CmdStan is normally driven on a cluster), and want each standalone chain to reproduce bit-for-bit what chain i of a combined n-chain MCMCSerial run would have produced. We do this by replicating mcmcsample's own per-chain seeding (sample.jl, MCMCSerial method):

seeds = rand(rng, UInt, nchains)
# ...
Random.seed!(rng, seeds[i])

We pre-position rng (via Random.seed!(rng, rng_seed) then discarding i - 1 UInt draws) so that our own rand(rng, UInt, 1) call, inside a nchains=1 ensemble, lands on the same value seeds[i] would have been in the nchains=n run. We verified this positioning is exactly correct by direct inspection — the discarded-then-drawn value matches the target seed bit-for-bit.

Despite that, chain 1 always reproduced correctly and every later chain (2, 3, ...) diverged from the combined run — different sampled values from the first HMC step onward.

Root cause

Turing's AbstractMCMC.sample override (mcmc/abstractmcmc.jl) defaults check_model=true, and runs it before dispatching to mcmcsample:

check_model && Turing._check_model(model, spl)

Turing._check_modelDynamicPPL.check_model(model; ...)DynamicPPL.check_model(Random.default_rng(), model; ...), which evaluates the model once via DynamicPPL.init!!(rng, model, oavi, InitFromPrior(), ...) — sampling every site from its prior. That's a real draw from Random.default_rng(), the same global stream we had carefully positioned.

Because check_model runs once per sample(...) call (not once per chain), and our seed-positioning discard count depends on chain index, the check's own draws land at a different stream position for chain 1 (zero prior discard) than for chain 2+ (i - 1 draws already discarded) — so the number of draws the check consumes is the same, but where it starts consuming from differs, which shifts the subsequent internal seeds = rand(rng, UInt, nchains) draw to a different value than the combined run's equivalent.

Why this is worth flagging upstream

This isn't a bug in the sense of check_model doing something undocumented — it's documented, and there's a kwarg to disable it. But nothing signals that it draws from the passed/default RNG, which makes it a sharp edge for anyone trying to reason about or reproduce the exact RNG stream a sample(...) call consumes (checkpointing, resumable/split sampling, RNG-stream auditing). We only found it by bisecting with debug prints across several hours; a docstring note would have saved that.

Suggested fix (happy to open a PR if this is welcome)

  • Note in check_model's docstring (both the Turing._check_model/AbstractMCMC.sample kwarg and DynamicPPL.check_model itself) that it draws from the passed-in / default RNG and will shift any subsequent sampling stream derived from the same RNG object.
  • Optionally, have the check_model=true path use a RNG copy/independent stream by default (e.g., Random.default_rng() explicitly copied, or a locally-seeded throwaway RNG) rather than consuming from the same object subsequent sampling will use — this would make check_model's default behavior reproducibility-neutral, which seems like the least-surprising default regardless of our specific use case.

Reproduction

Happy to provide a minimal repro (small @model, compare sample(model, NUTS(), MCMCSerial(), N, 2)'s chain 2 against manually re-deriving and sampling chain 2 alone with check_model=false vs check_model=true) if useful — let me know and I'll trim our internal repro down to something upstream-shareable.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with mcmc/abstractmcmc.jl and the DynamicPPL.check_model entry point named in the issue. Review the check_model docstrings and document that model checking consumes draws from the passed or default RNG, including the effect on later sampling. Done means both relevant APIs clearly warn about this RNG behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
documentation
Issue type
Documentation
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
64/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.