Numpyro: Looking for example of mini-batch training with SVI

[vitkl]

Hi numpyro team!

I am interested in doing mini-batch training with both local and global variables. Ideally using Pytorch style with DataLoader to avoid copying all data to GPU and lax.scan for speed. Any related examples would be appreciated!

I have this in mind but not sure if it could work:

# from jax tutorials
def numpy_collate(batch):
    if isinstance(batch[0], jnp.ndarray):
        return jnp.stack(batch)
    elif isinstance(batch[0], (tuple, list)):
        transposed = zip(*batch)
        return [numpy_collate(samples) for samples in transposed]
    else:
        return jnp.array(batch)

class NumpyLoader(DataLoader):
    def __init__(self, dataset, batch_size=1,
                 shuffle=False, sampler=None,
                 batch_sampler=None, num_workers=0,
                 pin_memory=False, drop_last=False,
                 timeout=0, worker_init_fn=None):
        super(self.__class__, self).__init__(dataset,
                                             batch_size=batch_size,
                                             shuffle=shuffle,
                                             sampler=sampler,
                                             batch_sampler=batch_sampler,
                                             num_workers=num_workers,
                                             collate_fn=numpy_collate,
                                             pin_memory=pin_memory,
                                             drop_last=drop_last,
                                             timeout=timeout,
                                             worker_init_fn=worker_init_fn)

loader = DataLoader(dataset, batch_size=256)
state, losses = lax.scan(lambda state_1, batch: self.svi[name].update(state_1, batch),
                                         init_state, loader)

The example here looks rather convoluted because it is not clear what are the requirements for train_fetch https://github.com/pyro-ppl/numpyro/blob/master/examples/vae.py#L89.

[fehiepsi]

to avoid copying all data to GPU and lax.scan for speed

I don't know how to use lax.scan and avoid copying all data to GPU. I think you can raise this question upstream to JAX. I just guess that it is not possible.

About examples, you can jump around flax/haiku repository to find a pattern that is most suited for your problem. In my opinion, if you use data loader, then you don't need to use lax.scan. Just jit the svi.update method, you are good to go. Just curious, what are speed differences between using lax.scan and not using it? (I thought the reason to use GPU is to speed up 1 single expensive update step, whose time should dominate the advantages of lax.scan)

[vitkl]

I see. Thanks for explaining!

With regard to speedup, lax.scan https://github.com/vitkl/cell2location_numpyro/blob/main/cell2location_numpyro/models/pyro_model.py#L234-L245 for my problem takes roughly 9 minutes on GPU whereas wrapping svi.update in jit (https://github.com/vitkl/cell2location_numpyro/blob/main/cell2location_numpyro/models/pyro_model.py#L247-L255, https://github.com/vitkl/cell2location_numpyro/blob/main/cell2location_numpyro/models/pyro_model.py#L301-L309) takes 2-3 hours on GPU. Maybe I am doing something else incorrectly.
I will try svi.update in jit with minibatch data.

[fehiepsi]

I was seeing you redefined body_fn in each step, so jit will not hit the cache. You can do something like

jit_step_update = jit(self.step_update, static_argnums=(0, 3))
for e in epochs_iterator:
    self.state[name], loss = jit_step_update(self.svi[name],
                                             self.state[name],
                                             self.x_data,
                                             self.extra_data)

[vitkl]

You were right! Thanks for pointing this out.
Training with svi.update in jit (https://github.com/vitkl/cell2location_numpyro/blob/main/cell2location_numpyro/models/pyro_model.py#L251-L260) is still slower than lax.scan - 28 min vs 9 min on my problem.

[fehiepsi]

Interesting! I have an impression for a long time that using jit(update) alone only adds dispatching cost, which moves the loss from GPU to CPU, is very small comparing to time to perform an actual svi.update. Assume that you did SVI with 100000 steps, then according to your last comment, the additional cost takes 10ms per svi step, which is quite large in my opinion. I am not sure what is the reasoning behind that cost. Just curious, what happens when you removed those tqdm stuffs?

[vitkl]

That's right, I am doing 30k iterations (full data 6486 obs * 6303 var, not minibatch). Removing tqdm does not change anything.

I omitted one important detail previously, I am adding the following flags:

os.environ["JAX_PLATFORM_NAME"] = "gpu"
os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"]="platform" 
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"]="false" # without this all GPU memory is reserved (e.g. 32 GB)

It looks like os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"]="platform" has most substantial effects:

1. Removing this flag equalised the speed difference between svi.update in jit and lax.scan - to 8 min for both.
2. Removing it also increased the GPU memory utilisation for svi.update in jit from 4.4 to 9.2 GB. I also see that the memory utilisation fluctuates between 2.4 and 4.4 GB when this flag is set - potentially indicating that some data is moved in and out of memory (which could explain lower speed).
3. Removing it also increased the GPU memory utilisation for lax.scan from 4.4 to 9.2 GB

Requiring 9.2 GB for this relatively small dataset is quite limiting. Also, for the same model, same data and the same number of iterations, pymc3 takes 14 min, 3 GB of GPU memory and 30GB of RAM (numpyro with lax.scan takes 9.6GB RAM).

[fehiepsi]

Interesting! Thanks for your info. I just learn that there are those settings in JAX. Looking like in jax docs, they said platform is very slow and is not recommended for general use. :D

Looking like this issue is clearer now. To avoid copying all data to GPU, you need to jit body function instead of using scan and this solution is not as bad as it seems initially (2-3 hours). With platform (28 min), it is 2x slower than pymc3 (14 min), while without platform (9 min), it has a similar performance as scan and is 1.5x faster than pymc3. I guess the fastest solution is to not use platform allocator and set XLA_PYTHON_CLIENT_MEM_FRACTION=.XX as in jax docs to decide how much memory to use in your problem.

I hope that we can find a way to share your observations because not many NumPyro/Pyro users are familiar with JAX and especially its settings. One way is to add links to your repository in the README file. Or to add a note section to SVI.run. If you can contribute a tutorial, that would be great! ;)

[vitkl]

I will try to summarise the note for SVI.run next week.

We will implement the minibatch training in our repo within the next 2-3 weeks and then it can be added as example somewhere. So far I did tests on all data rather than minibatch. Our package aims to provide general training code for many similar models so it won't be a minimal example of minibatch training - but we can try to distil the relevant code and add somewhere as an example. Would that be useful?

[fehiepsi]

Please just stick with what is convenient for you. It would be nice if you share some of the best practices in your repository README (or any other form). Then we can add link in numpyro README pointing to your repository, to illustrate numpyro usage on GPU in real application. I believe either way would be useful for many users. When someone asks, we can point to those best practices. :)

Pls let me know if you need further help. I'm looking forward to seeing your example.

[fehiepsi]

I guess the questions have been clarified. Please feel free to reopen the issue or use forum if there are follow-up questions.