CUDNN cache locking prevents finalizers resulting in OOMs

[mashu]

Describe the bug

Have been working on model using 5 CNN layers and Dense, but same model that worked before after updating CUDA started giving out error

Out of GPU memory trying to allocate 653.794 MiB
Effective GPU memory usage: 99.32% (7.734 GiB/7.787 GiB)
Memory pool usage: 5.630 GiB (6.344 GiB reserved)

Revetring back to CUDA@3.8.0 which I luckily had in git log fixed the problem.

To reproduce

I used https://github.com/FluxML/model-zoo/tree/master/vision/conv_mnist
Updated CUDA and Flux to latest versions and changed a bit model to make it stupidly bigger so number of parameters is more similar to my model. See attached diff
LaNet.txt

Manifest.toml

I've tested in the following order

 CUDA@3.8.5 #1 memory out
 CUDA@3.8.4 #2 memory out #6 memory out again
 CUDA@3.8.0 #3 works
 CUDA@3.8.2 #4 works
 CUDA@3.8.3 #5 works

So something bad happens between 3.8.3 and 3.8.4.

Expected behavior

I expected the same code that worked before to still work with CUDA@3.8.4+

Version info

Details on Julia:

Julia Version 1.7.2
Commit bf53498635 (2022-02-06 15:21 UTC)
Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-12.0.1 (ORCJIT, skylake)
Environment:
  JULIA_NUM_THREADS = 6
  JULIA_EDITOR = atom  -a

Details on CUDA:

CUDA toolkit 11.6, artifact installation
NVIDIA driver 470.103.1, for CUDA 11.4
CUDA driver 11.4

Libraries: 
- CUBLAS: 11.8.1
- CURAND: 10.2.9
- CUFFT: 10.7.0
- CUSOLVER: 11.3.2
- CUSPARSE: 11.7.1
- CUPTI: 16.0.0
- NVML: 11.0.0+470.103.1
- CUDNN: 8.30.2 (for CUDA 11.5.0)
  Downloaded artifact: CUTENSOR
- CUTENSOR: 1.4.0 (for CUDA 11.5.0)

Toolchain:
- Julia: 1.7.2
- LLVM: 12.0.1
- PTX ISA support: 3.2, 4.0, 4.1, 4.2, 4.3, 5.0, 6.0, 6.1, 6.3, 6.4, 6.5, 7.0
- Device capability support: sm_35, sm_37, sm_50, sm_52, sm_53, sm_60, sm_61, sm_62, sm_70, sm_72, sm_75, sm_80

1 device:
  0: NVIDIA GeForce RTX 2070 with Max-Q Design (sm_75, 4.775 GiB / 7.787 GiB available)

Additional context

With broken CUDA@3.8.4 using

    GC.gc(true);
    CUDA.reclaim();

at the end of each mini-batch also helped to keep memory usage down to 21%, when this is not done in all cases nvtop shows 100% memory usage, but CUDA 3.8.0->3.8.3 is not crashing with out of memory error, while 3.8.4 and 3.8.5 just crash.

[maleadt]

Since it works on 3.8 and not on 3.8.4, between which there shouldn't have been any significant changes, could you try bisecting to see where the issue was introduced?

[fabricerosay]

If this can help, it seems related https://discourse.julialang.org/t/oom-when-using-flux-and-loops/79468/5?u=fabrice_rosay

[ToucheSir]

That issue may be in part due to invalidation of the cuDNN conv algorithm descriptor cache in CUDA.jl. Note how said cache is keyed by input, so if the input size changes then a somewhat compute + memory-intensive algorithm search will have to happen again. It seems to me that the constraints on this cache could probably be loosened (e.g. use input size + eltype + strides instead of a descriptor, ignore the batch dimension, etc) but unsure of any potential negative side effects.

[fabricerosay]

I don't know enough about CUDNN but I made the change in my algo( a kind of alphazero) to use a constant batch size and the gpu usage went much higher from between 20 to 40 and a strange behaviour where sometimes it dropped under 1% to above 70% and steady, which largely compensate the fact that you do useless computation... This also could be related with the addition of lock in cudnn to make it compatible with threading, which happened in 3.8.4, making this behaviour even more pronouced. (sorry for the English:). As another hint pointing to this direction, when i enable multithreading, i encounter a lot of dead lock: sometimes it runs smoothly for hours and sometimes it won't even start.

[ToucheSir]

Yeah, in that case I'd be pretty confident it is the repeated algorithm selection that is the culprit (the lock creates basically no overhead, IIUC). You can confirm this by running an nsight profile (see the docs for this package on how to do that) and checking the CUDA API calls. If there are repeated blocks of kernels with names containing "fft", "conv", "winograd" etc for each new input size, then that's indicative of algorithm selection. If instead you only see such a block at the beginning and somewhat regular calls to routines with "conv" in the name afterwards, then no re-tuning is happening.

[fabricerosay]

I once tried to use nsight profile which ended in a major failure. But as I just got tested COVID positif I will have some time to try again:)
Between in my experiments only the batching dimension is changing, do you think it would be easy to force the same algo to be used ?

[ToucheSir]

Unfortunately not without modifying either CUDA or NNlib(CUDA) + Flux. There would need to be some mechanism for passing in p here instead of unconditionally running an algorithm search (even if it hits the cache). I count ~4 layers of functions between there and Flux's Conv layer, so you'd either have to thread it through the call stack or come up with some non-local mechanism for relaying the option like how math mode is currently done. This would probably tie in with #938 as well.

[fabricerosay]

If I understand correctly, there's a whole machiney in Flux and nnlib to decide which algo to perform, but then it mostly call cudnn.
As most of the time is spent on inference in alphazero, I could make a small script that given a flux model call directly cuddn to perform the convolution, if needed using first the heavy machinery to find the best algorithm on average batch size and then switching back to flux for training. In most sessions the network have the same toppology.

Anyway training is very stable now. We had this issue from the begining suspecting some leaks or gc, getting crazy... this could be very valuable for alphazero.jl, which suffer from the same problem.

I will try to go back to 3.9 and multithreading and see.

[ToucheSir]

If I understand correctly, there's a whole machiney in Flux and nnlib to decide which algo to perform, but then it mostly call cudnn.

It's actually the opposite: the machinery exists only in CUDA.jl, so Flux and NNlib have no idea which algorithm is chosen. That's why it's hard to manually specify algorithm types, because the API boundary between the lowest-level Julia library and the next highest level provides no knobs to tune them.

[fabricerosay]

It works flawlessly with CUDA 3.9 and multithreading provided batch size is constant.

[tclements]

I just ran into a similar OOM problem when doing convolution on graph features with GraphNeuralNetworks.jl JuliaGraphs/GraphNeuralNetworks.jl#150. In my case, enforcing a constant batch size is tricky because the batch size is a function of the number of nodes in each graph, which varies graph to graph.

Is there a way around having to do the conv algo search when batch size varies? All other input dimensions remain the same during training/inference.

[fabricerosay]

I tried to hack directly into cudnn, there you can force an algorithm. For example i tried CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM = 0 always and memory consumption dropped. . Problem is it seems some algorithm do not work for some convolutions and if you cancel the call to algo search you have to guess memmory consumption.
Hacking directly into cudnn you can print which algo is used so one very ugly workaround would be to force the googd algorithm at each step depending on the structure of your network.
This would allow to bypass the optimization step. If speed is not an issue it could work.
You have to change p.algo there with something like CUDA.CUDNN.cudnnConvolutionFwdAlgo_t(0)

CUDA.jl/lib/cudnn/convolution.jl

Line 105 in 146ad00

cudnnConvolutionForward(handle(), alpha, xDesc, x, wDesc, w, convDesc, p.algo, workspace, sizeof(workspace), beta, yDesc, y)

For me it worked dividing memory usage by 2. I don't know yet how it will perform on a longer run. I have to wait till tomorrow morning.
All in all it is very ugly

[tclements]

Thanks, that's really helpful! Probably a bit too low-level for what I'm trying to do with varying convolution sizes.

@ToucheSir if there's a way to not trigger the algorithm search if only the batch size dimension varies, that would be awesome!

[ToucheSir]

@tclements per my comment at #1461 (comment), someone would have to figure out the minimum subset of properties required to safely cache conv algorithms without relying on descriptor identity.