Repository: pyro-ppl/numpyro

Issue:
Reducing GPU memory usage
I'm opening this issue following the discussion on the forum: https://forum.pyro.ai/t/reducing-mcmc-memory-usage/5639/6.

The problem is, not-in-place array copying that happens in `mcmc.run` after the actual sampling might result in an out-of-memory exception even though the sampling itself was successful. First of all, it would be nice if this could be avoided and the arrays could be transferred to CPU before any not-in-place operations.

More generally, the GPU memory can be controlled buy sampling sequentially using `post_warmup_state` and transferring each batch of samples to CPU before running the next one. However, this doesn't seem to work as expected, and the consequent batches require more memory than the first one (see the output for the code below).
```
mcmc_samples = [None] * (n_samples // 1000)
# set up MCMC
self.mcmc = MCMC(kernel, num_warmup=n_warmup, num_samples=1000, num_chains=n_chains)
for i in range((n_samples) // 1000):
    print(f"Batch {i+1}")
    # run MCMC for 1000 samples
    self.mcmc.run(jax.random.PRNGKey(0), self.spliced, self.unspliced)
    # store samples transferred to CPU
    mcmc_samples[i] = jax.device_put(self.mcmc.get_samples(), jax.devices("cpu")[0])
    # reset the mcmc before running the next batch
    self.mcmc.post_warmup_state = self.mcmc.last_state
```
the code above results in:
```
Running MCMC in batches of 1000 samples, 2 batches in total.
First batch will include 1000 warmup samples.
Batch 1
sample: 100%|██████████| 2000/2000 [11:18<00:00,  2.95it/s, 1023 steps of size 5.13e-06. acc. prob=0.85]
Batch 2
sample: 100%|██████████| 1000/1000 [05:48<00:00,  2.87it/s, 1023 steps of size 5.13e-06. acc. prob=0.85]
2023-11-24 14:43:23.854505: W external/tsl/tsl/framework/bfc_allocator.cc:485] Allocator (GPU_0_bfc) ran out of memory trying to allocate 2.56GiB (rounded to 2750440192)requested by op 
```

To summarise,
1. Could not-in-place operations at the end of sampling be optionally transferred to CPU?
2. How should one sample sequentially so that memory usage is not increased in the process?


Discussion:
Could you try replacing those [two lines](https://github.com/pyro-ppl/numpyro/blob/master/numpyro/infer/mcmc.py#L652-L653) by
```
        self._states = jax.device_get(states)
        self._states_flat = jax.device_get(states_flat)
```
If it works, then we can introduce a method named `transfer_states_to_host()` to perform those `device_get` operator.
Yep, it works!
Neither requires extra GPU memory after sampling nor leads to memory increase throughout sequential samples.
Thanks! Do you want to make a PR to add a helper doing such work? :)

Resolving patch:
diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py
--- a/numpyro/infer/mcmc.py
+++ b/numpyro/infer/mcmc.py
@@ -9,7 +9,7 @@
 
 import numpy as np
 
-from jax import jit, lax, local_device_count, pmap, random, vmap
+from jax import jit, lax, local_device_count, pmap, random, vmap, device_get
 import jax.numpy as jnp
 from jax.tree_util import tree_flatten, tree_map
 
@@ -721,6 +721,13 @@ def print_summary(self, prob=0.9, exclude_deterministic=True):
                 "Number of divergences: {}".format(jnp.sum(extra_fields["diverging"]))
             )
 
+    def transfer_states_to_host(self):
+        """
+        Reduce the memory footprint of collected samples by transfering them to the host device. 
+        """
+        self._states = device_get(self._states)
+        self._states_flat = device_get(self._states_flat) 
+
     def __getstate__(self):
         state = self.__dict__.copy()
         state["_cache"] = {}


Test patch: