Scalability

numpyro
1

Hi everyone. Great project; very excited to see HMC reaching scalability via techniques like HMCECS.

I’ve been playing around with a single layer neural network with a data set of shape ~ (500k, 300). From my naive understanding, sub-sampling SVI and HMC via HMCECS on a GPU should allow a data set like this to scale. However I see very large memory requirements even for a small hidden layer of h1 = 5. Additionally it takes a very long time for (say SVI) to begin.

I’d like to understand from others where the known bottlenecks are located. Is it XLA compilation or elsewhere? I have not started profiling but thought it best to get information from those who know this project and the related literature better than I before I begin.

Thanks!

_Chris

2

Hi @cfusting, for SVI, if the data is large, it is best to feed data in batches (see some discussions here). Memory requirement should be small because in no way the whole dataset is involved. For HMCECS, in the initialization stage, a large memory requirement is used to calculate some control variate factors - those calculations use the whole dataset. After that, at each sampling step, we need to draw new subsampled data from the whole dataset - this subsampling step happens in XLA so it takes some memory.

it takes a very long time for (say SVI) to begin

It should not be the case if you use the pattern

svi_state = svi.init(PRNGKey(0), batch_data)
for i in range(1000):
    batch_data = get_batch()  # it is better to do this in CPU
    svi_state = svi.update(svi_state, batch_data)
params = svi.get_params(svi_state)
3

Thank you!

I thought I could use a sub-sampling plate as in HMCECS, NUTS, etc. with SVI to avoid a full allocation of memory.

On that topic, what is the expected behavior of algorithms that do not support sub-sampling when the sub-sampling plate is used? For example in the HMCECS example HMCECS is compared with HMC; but should we expect similar posteriors for the later given we are taking sub-samples?

4

Did you mean covtype example? There we run HMC with subsample_size=None, i.e. no subsampling. HMC does not support subsampling so I guess you will get some errors if you run it with subsample_size!=None.

1 Like