zarr store and GPU sampling

[Wmuntean]

I've noticed that when using zarr store and sampling a large model on a GPU via JAX, the trace is still in active memory as well as getting written out to the zarr store. Unfortunately, my system RAM runs out before sampling is complete (or just after sampling is complete and nutpie is doing post processing on the trace). I thought the expected behavior of setting zarr store was to provide some relief to the system memory requirements with large models. For reference, the same model sampled in PyMC with a numpyro backend is much more memory efficient but less performant despite numpyro using multiple GPUs. I know this issue lacks detail so let me know whether you need additional details or need a working example.

[aseyboldt]

That is certainly not supposed to happen, thank you for reporting this.

I don't know for sure if that is the problem you are running into, but it seems I misunderstood something about the LimitStore in the rust object_store crate. I thought it would limit the number of queued writes, which it doesn't (and can't with its API...).

This means that if the sampler is faster than we can write the trace, the trace will still be buffered in memory. For instance, with this the memory usage currently grows quite large:

import pymc as pm
import nutpie
import numpy as np

with pm.Model() as model:
    x = pm.Normal("x", shape=1000)
    pm.Deterministic("y", x[:, None] + x[None, :])

compiled = nutpie.compile_pymc_model(model)

!rm -rf test-trace.zarr
!mkdir test-trace.zarr

zarr = nutpie.zarr_store.LocalStore("test-trace.zarr")

tr = nutpie.sample(compiled, zarr_store=zarr, chains=4)

If that is the issue, I think I've got a fix, where the sampler actually waits for the trace backend as was intended from the start, but I'll have to do a bit more testing with it.

[Wmuntean]

Thanks for looking into this. I think you nailed it, but let me also put together a concrete example because I think the memory issue might span beyond the zarr store. In my exploratory analyses, using numpyro as the backend stresses system ram to about ~10GB memory and whereas nutpie pushes system ram north of ~90GB.

[aseyboldt]

If you have an example where nutpie uses 10x the memory compared to numpyro, I'd definitely like to know about it. :-)

[aseyboldt]

@Wmuntean 0.16.4 is being built now, and should fix this issue.

I tried to reproduce the memory problem you mentioned, but with the models I tried nutpie only used the expected amount of memory. If you come across a model where that's not the case, I'd welcome another issue!