Potential memory leak for cuTENSOR contractions with minimal example

[lamberev]

I've possibly encountered a memory leak in cuTENSOR-based contractions. The memory allocated in each call does not seem to be fully freed even after explicit calls to the garbage collector. The same, or at least a similar amount of memory appears to accumulate on the both the CPU and GPU (seen with calls of CUDA.reclaim() + CUDA.used_memory()). I've had no luck playing around with where input/output/intermediate tensors are stored via CUDAAllocator.

I've constructed and attached a minimal working example of the issue. It may be worth mentioning that this leak was found during transition to TensorOperations v5.2.0 with cuTENSOR v2 from TensorOperations v3.2.5 & cuTENSOR v1.

Minimal Working Example
Output

[lkdvos]

Thanks for reporting this, and providing a MWE!

May I ask what kind of system you are running this on?

I'm not entirely sure how reliable the results from Sys.free_memory() are, and there are many things that can be affecting this.
From what I understand, the results of this function would be affected by other processes running on your machine as well, and many different architectures have systems in place to cache recently used data which may be reported as non-free, but would quickly be freed if you are actually needing additional memory.
For example, I found this thread on its use.

I might be slightly misunderstanding what you said, but am I correct in assuming that you notice that after every iteration, CUDA.used_memory() goes up and the explicit GC calls don't resolve this? It might be an idea to check if for example you run this with larger arrays, do you actually run out of memory?

I have to admit I'm not entirely sure how cuTENSOR and CUDA in general handle their allocation strategies either. I know they have a dedicated pool, and we definitely do insert CUDA.unsafe_free calls on unreachable temporary objects to lower the GC pressure, but from what I understood this was more of a hint than an actual forced free. Additionally, there are some objects that are allocated during the contractions, for describing the different tensors etc that we are not explicitly freeing either, and while I wouldn't expect this to be this much, this should also definitely show up in these numbers.

I am definitely not an expert in most of this though, so it is most definitely possible that I am mistaken about some of this.

[lamberev]

Thanks for getting back so soon. This was Rocky Linux 8.10 (Green Obsidian).

It looks like the memory usage is not superficial. I've managed to trigger an OOM error on a ~180 GB compute node with that MWE after running for 2.5 hours. The original script which led me to this makes hundreds of thousands calls to a functionally equivalent version of that contract function in the MWE and crashes after a similar amount of time. Sys.free_memory does seem to faithfully reflect what's observed through top here AFAICT but good to know that it's not always reliable.

Regarding using larger arrays, I have tried that and what I found is interesting. Varying the dimension of each tensor in that MWE reproduces the same leak. I.e. I get the same ~15 MB leak per iteration no matter the tensor sizes. It seems to be only dependent on the number of tensors considered in each iteration so I suspect it's doing something like saving unnecessary pointers to GPU memory (doubling the number of tensors allocated in each iteration doubles the leak to ~30 MB/iteration).

Sorry, I take it back regarding the CUDA.used_memory() showing buildup, I see no leak on the GPU in that MWE. I should have confirmed before mentioning so please disregard. Regardless, the CPU leak is still present.

[lamberev]

Update: I was able to reproduce the leak on my previous system (CentOS v7).

[lkdvos]

Thanks for the additional information and extra check. I'll try and find some time later today (but more likely tomorrow) to investigate what is going on, but in case you have some time yourself, the first thing I would try is to check and see if using cuTENSOR.jl directly has the same behavior. In that case I am definitely out of my debt and we should try and open an issue there to get this resolved, since that is definitely a large amount of memory for each iteration.

The reason that is my first guess is mainly because I tried following their implementation as closely as possible, and in principle I shouldn't be doing any manual memory management here (I think, not entirely sure since there are some calls to unsafe_wrap), so the problem really might just be with cuTENSOR itself.

[lamberev]

I believe this reimplementation below is equivalent but let me know if I've missed something. (No leak observed when contractions done explicitly via low-level cuTENSOR operations.)

cuTENSOR MWE
cuTENSOR MWE Output

[lkdvos]

Just wanted to let you know that I can reproduce the issue and I'm investigating. Thanks again for the updated MWE, this helps out enormously!

[lkdvos]

So I've gone through and I have absolutely no idea what is going on, but if something is wrong it might actually be a problem with cuTENSOR itself. From what I can tell, the reason you don't get the leaks in the second example is because you are reusing the contraction plans, and these seem to actually cause the problem.

Let me reiterate how grateful I am for the MWEs, this sped up the process for me by a substantial amount. I'll keep an eye on this and the linked issue to see if we can find a solution.

[lkdvos]

Looks like the issue will be partially resolved once JuliaGPU/CUDA.jl#2794 lands.

Just as a reminder to myself here: this same strategy has to be adapted here: https://github.com/Jutho/TensorOperations.jl/blob/541f2f4daff7ac4e72d1e24a2a5259a5db81a1c0/ext/TensorOperationscuTENSORExt.jl#L266