Out-of-memory crashes after sampling completes with high `tune` despite `save_warmup = False`

[mike-lawrence]

I have a model that returns samples ok with low values of tune, but when I bump to high values, python crashes just as the sampling completes, with my OS (Ubuntu) logs showing an OOM event. I don't need the warmup samples, so I tried setting save_warmup = False, but no dice.

Is it possible that the dropping of the warmup data is occurring unnecessarily late, ex. after conversion from arrow to InferenceData? That's the only thing I can think of to explain why higher values for tune cause crashes despite equal values for draws.

[aseyboldt]

Sorry for the slow reply.
You are right, the warmup draws are only discarded at the very end, so not storing them does currently not reduce the maximum memory requirement.
I'm currently working on changing how the trace is stored, and hopefully that should solve this issue.

[aseyboldt]

Version 0.16 includes a zarr backend, so that you don't have to keep the whole trace in memory. It also lowers the memory requirement without the zarr backend a little when creating the xarray datastructure.
I think this should solve the underlying issue here, that sometimes we run out of memory for large traces.

I'll keep the issue open, because we still keep the tuning draws around until the very end even if save_warmup = False.