Potential Memory Leakage

[dannychen0830]

Hi,

I've encountered a memory problem when using Netket for drawing samples while using optimizers from different packages. I attached a working code to hopefully provide some concreteness. In short, I'm using a gradient-free optimizer (COBYLA) to optimize the parameters of an RBM with respect to a cost function defined using the samples generated using Netket packages. However, when running this for graph sizes above 50 with parallel MPI processes, the memory used grew rapidly and became problematic. I'm hoping to get some help regarding this issue.

Thanks a lot!
Danny
memory_test.py.zip

[PhilipVinc]

Does changing your initializer to

# initializer for passing in the parameters wanted for the neural network
def custom_init(x, dtype=jnp.float_):
    def init(key, shape, dtype=dtype):
        a = blank(x)
        return a
    return nk.jax.HashablePartial(init)

and leaving everything else exactly the same fix your issue?

[PhilipVinc]

To expand a bit on my answer: NetKet compiles most of it's methods calls with jax.jit. The main methods that get compiled are 1) initialisation of the NN, 2) sampling and 3) evaluation. Every one gets recompiled if you change the shape (hilbert size or number of samples) and/or if the hash of the model changes.

Depending on the models, every compiled instance of the code takes up from little to a lot of space (symmetries take up a lot of space for a bunch of reasons in the compiled code). For a simple RBM it should not take up much space but if you recompile thousands of times this can add up.

You can see if Jax/Netket are recompiling methods by setting the environment variable export export JAX_LOG_COMPILES=1 before running python. You should essentially try to reduce this to the minimum.

In this particular case, you can check what is the hash of the model by adding a print statement to your evaluate function

        # create the Boltzman machine with custom parameters
        model = nk.models.RBM(alpha=1, kernel_init=custom_init(kernel), hidden_bias_init=custom_init(bias), visible_bias_init=custom_init(vis_bias))
        print("-----------------------")
        print(hash(model))
        print("-----------------------")
        # create variational state
        vs = nk.vqs.MCState(sampler=self.sampler, model=model, n_samples=4000)

You will see that this changes every time you call evaluate, triggering recompilation (And significantly slowing down your code).

I go to great lengths in netket to ensure that identical objects have the same hash, but that is not easy.
The simplest solution would be to create the model outside the evaluate function and just reuse it all the time.

The alternative is to make sure the hash does not change, which is what I proposed above. (The hash changes because you are creating a new init function inside of custom_init every time you call custom_init, and even though they are identical, for python those are different objects and have different hashes unless you use HashablePartial that avoids that.)

--

Another issue: creating an MCState is a very expensive thing and it is not optimised at all. Unless your model is huge/you take many samples, I'd change the code to create one MCState for every worker and then simply change the parameters every time you call evaluate. You can flatten the weights by using nk.jax.tree_ravel. If you ever use a much more complex models you should jit it, but for simple one it does not matter.

See this updated gist
https://gist.github.com/PhilipVinc/9aeadfedee0e2763be718042dea65de6

[PhilipVinc]

@dannychen0830 did you get a chance to investigate this further?

[dannychen0830]

Yes, sorry for the delay. Thanks a lot for this detailed reply. The memory usage significantly improved and it is not a problem anymore.