google / google/meridian

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

Open
#1,709 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.5k
Forks
294
Avg merge
1d 21h
Merged PRs (30d)
47

Description

### 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 #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-#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: #1666 is NUTS-phase on a T4; #534 is host-RAM growth). This appears unreported.

Contributor guide

Open the contributing guide

Research direction

Start in meridian/model/posterior_sampler.py at _reconstruct_posteriors around line 711 and its call around line 888; inspect how reconstruction receives all draws after NUTS. Compare the existing OOM handling around lines 863-868, then reproduce with different n_keep values. Done means reconstruction no longer fails unexpectedly at the reported draw counts, with appropriate coverage for the chosen handling.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, tensorflow
Domain
machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.