Skip to content

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

Description

@karimn

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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions