What is max_con in _bose_hubbard within the operator?

[usernamesjh]

Dear developers,

Hello! I encountered a difficult problem while using Netket, and I hope to get your help here. I have tried various methods, but the problem still persists, so I decided to seek advice from you.

At first, I saw max_con in numba.py, and it seems to determine the size of x_mels (a collection of different quantum states). Therefore, after we defined other special_edges, I increased max_con.


Then this will lead to a problem where the required memory increases rapidly.

We ran this program on the GPU and selected the sampler as 992. We found that the required memory reached 115G, which was too large to fit on the GPU, so it couldn't run. We had to continuously adjust the sampler to a smaller value until it reached 272 before it could be placed on the GPU.
After that, we tried a series of other samplers, ranging from 272 to 36. We found that as long as the program runs normally, the required video memory size is around 7750M. It seems that the size of the sampler has nothing to do with the video memory.

So I wonder, is my understanding of max_con correct? Why does the required video memory decrease as the sampler decreases when the program fails to run normally, and when the program can run, it only requires 7750MB of video memory regardless of the value of the sampler?
Are there any improvement methods?

Or, is the required memory when it fails to run normally the physical memory? But my mainframe has several terabytes of memory. I don't understand why such a problem occurs. Please feel free to give your advice!

Thank you very much for your attention and help. I am looking forward to your reply. I will try to provide more information if necessary. Have a nice day!

[PhilipVinc]

constant gpu memory usage is normal, because jax preallocates 75% of your memory all the time. You can read more about it in jax documentation.

The rest, I don't understand what exactly you are doing and you do not give a full stack trace.
The (long) stack traces that python prints are useful to understand what goes wrong. You can just paste them into gist.github.com and send a link. This would make it easier to help you. Those screenshots do not contain enough information.

Regardless, it seems to me you have an operator with several connected elements.
It's likely that the memory bottleneck you are encountering arises because you are trying to compute everything in a single pass instead of chunking your calculation in smaller blocks.

Did you try to specify vstate.chunk_size = some large multiple of your # of samples? For example, if you have 1024 samples, you can specify it to be 10240 or so. Then, reduce the chunk size until your calculation fits on the gpu.