sample_posterior 00Ms during posterior reconstruction on configs that fit pre-#1465( <= v1.5.2)

[TylerDurdenUI]

Summary

Since PR #1465 ("Optimize MCMC performance by decoupling sampling and reconstruction graphs", merged 2026-02-21, released in v1.5.3), sample_posterior runs a separate "reconstruction" step after NUTS completes. This step runs the full model forward over
all draws at once (n_chains × n_keep) via tf.vectorized_map, materializing a single very large tensor. On geo-level models this 00Ms on hardware that comfortably fit the same model in v1.3.0 / pre-#1465.

Environment

• Meridian 1.7.0 (stock, pip 'google-meridian==1.7.0') - also reproduced conceptually vs 1.3.0.
• TensorFlow 2.19.1, Python 3.12, NVIDIA L40S (~43 GB VRAM), set_memory_growth(True).

Repro / observed

• A geo-level model at production scale: on the order of ~10¹-10² geos and ~10² weekly time points, with a few dozen paid media channels (+ RF + organic) and a few dozen controls.

• Posterior sampling with serial chain batches (n_chains passed as a list) and a moderate n_keep, giving a total of ~400 draws (total_ chains x n_keep).
• NUTS completes successfully (tens of minutes). The crash happens afterward, in the reconstruction step:

tensorflow.python. framework.errors_impl.ResourceExhaustedError:
OOM when allocating tensor with shape IN_draws, P, G, T, Cl and type float
... device:GPU:0 by allocator GPU_0_bfc
File "
.../meridian/model/posterior_sampler.py", line 888, in _call
reconstructed_items = self._reconstruct_posteriors (...)
File "
.../meridian/model/posterior_sampler.py", line 711, in _reconstruct_posteriors return full_dist_unpinned. sample(value=values, seed=...)

where the leading dim N_draws = total_chains × n_keep (~400 here) and G, T, C are the geo / time / channel dimensions. This single allocation reached ~20 GiB on our hardware (a ~40 GB-class GPU), exhausting the device.

Root cause

• v1.3.0 had no _reconstruct_posteriors. NUTS sampled the full joint distribution and
mcmc.all_states already contained deterministics; post-sampling was just array reshaping - no large forward-pass tensor ever existed.
• PR Optimize MCMC performance by decoupling sampling and reconstruction graphs #1465 added a yield_deterministics flag: NUTS samples latents-only, then reconstruction rebuilds deterministic (ROI / Hill / adstock / contributions) over all draws simultaneously.
This trades faster sampling for a large one-shot memory spike whose peak scales with total draws.

Why the existing OOM mitigations don't cover it

• The documented fix (pass n_chains as a list for serial chains; see developers.google.com/meridian/docs/post-modeling/model-debugging#gpu-oom-error) reduces NUTS peak memory. Reconstruction happens after all chain batches are concatenated, so serial chains don't help it.
• The friendly MCMCOOMError try/except wraps only the NUTS call (posterior_sampler.py ~863-868). Reconstruction OOM escapes it and surfaces as a raw ResourceExhaustedError* with no clear guidance.
• `posterior_thinning (added 1.7.0) runs after reconstruction - does not reduce its memory.
• JAX backend reduces NUTS memory, not reconstruction.

Potential fixes

1. Batch/chunk the reconstruction forward pass over draws (e.g. iterate over blocks of
 n_keep and concatenate results) so peak memory is controlled and no longer scales with total draws.
2. At minimum, wrap reconstruction OOM in the same OOM handling as NUTS and emit an
MCMCOOMError-style message pointing to n_keep as the lever, plus a docs note that reconstruction memory ∝ n_chains × n_keep.

Evidence of linear scaling

• 400 total draws → tensor [400, ...] = ~20 GiB → 00M.
• 200 total draws (halved via n_keep) → tensor [200, ...] = ~10 GiB → runs end-to-end, no 00M (full fit + reconstruction + logging close). Model metrics essentially unchanged vs the pre-Optimize MCMC performance by decoupling sampling and reconstruction graphs #1465 baseline, confirming reconstruction memory scales linearly with
total_chains × n_keep and that the OOM is reconstruction's intrinsic peak - not HMC-residual memory or fragmentation.

Notes for when reporting this

• No existing upstream issue covers reconstruction-phase 00M (checked: GPU Memory Exhaustion During Model Fitting on NVIDIA T4 #1666 is NUTS-phase on a T4; Possible Memory Leak, Colab Crashes #534 is host-RAM growth). This appears unreported.