TuringLang / TuringLang/DynamicPPL.jl
check_model=true (the default) silently consumes draws from the caller's RNG, breaking manual per-chain seed reproduction
Nobody has claimed this yet.
- 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_model → DynamicPPL.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 theTuring._check_model/AbstractMCMC.samplekwarg andDynamicPPL.check_modelitself) 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=truepath 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 makecheck_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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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