Memory Leak when using CUDA (and Flux)

[fps]

I have no idea whether this is a Pluto.jl, Flux.jl or a CUDA.jl issue, so I had to choose one of them to report it to first.

Repeatedly calling the simplest convolutional Flux model on the GPU leaks host memory. Not GPU memory, but host memory as can be seen from the below screencast.

I found a previous report where turning off the distributed feature of pluto fixed the issue. I could live with that but sadly workspace_use_distributed=false just breaks Pluto completely (importing packages fails) so I could not test that hypothesis. I will probably make a separate bug report for that.

Notebook:

This notebook reproduces the issue on my systems:

• NixOS, X86_64. NVIDIA RTX 3080 gpu
• Windows 10, X86_64, NVIDIA Geforce GTX1050ti gpu

https://gist.github.com/fps/1a091a3ea8aab061e25ab7f11b87abb9

Video:

Medien1.mp4

[fonsp]

You can figure out if this is a Pluto issue by running the notebook with and without Pluto.

You can probably run the notebook.jl directly in the REPL with include. And run Pluto.activate_notebook_environment before to get the same package versions.

[fps]

Oh! I have forgot to mention this: Running the same code directly in the REPL does not leak memory. I have moved my training code back to the REPL after the initial prototyping in Pluto for this reason. But I will confirm for this reproducer notebook just to be sure..

[fps]

You can figure out if this is a Pluto issue by running the notebook with and without Pluto.

You can probably run the notebook.jl directly in the REPL with include. And run Pluto.activate_notebook_environment before to get the same package versions.

Yeah, confirmed: Running it in the REPL does not leak memory. It can still be a CUDA.jl issue instead I guess. Interactions between different packages can be tricky :)

[fps]

Hmm, OK, I tried running Pluto with workspace_use_distributed=false after follwing the tips from #3184 (namely running pluto in an environment with the required libs installed). And alas, it did not help with the memory leak.

[fps]

Hmm, it seems to be the cuDNN library. If I inspect the variables in the workspace with varinfo(@__MODULE__(), imported=true) I can see the cuDNN module increasing in size (see the below screenshot and note the module size of cuDNN.) The increase in size seems to correlate 1 to 1 with the host memory usage of the julia process. Also it seems the varinfo call gets slower and slower the bigger the cuDNN module becomes..

[fps]

Hmmm, 5.6 million log messages in cuDNN.log_messages:

Credit to: https://discourse.julialang.org/t/memory-usage-increasing-with-each-epoch/121798/2 and #2913 for making me look at that..

To "reclaim" the memory just call popfirst!(cuDNN.log_messages) the appropriate number of times or use some other method to empty that array.

[fps]

And it seems all these messages are completely benign. So there is no need to store them. I wonder how to change the cuDNN log level and why they pile up in the module only in Pluto.jl but not on the REPL.

[fps]

The CUDNN documentation talks about environment variables: https://docs.nvidia.com/deeplearning/cudnn/backend/v9.1.0/reference/troubleshooting.html

I tried to explicitly set CUDNN_LOGLEVEL_DBG=0 and CUDNN_LOGDEST_DBG=NULL for the julia process. They seem to get ignored though.

EDIT: The environment variables are ignored because the cuDNN package explicitly sets up a callback to get messages.

[fps]

This seems to work around it (run in a cell after importing cuDNN):

cuDNN.cudnnSetCallback(UInt(0), C_NULL, C_NULL)

I still don't quite get why this would only happen in Pluto.jl. The initialization code [1] in cuDNN checks the following conditions before setting the callback:

if !precompiling && (isdebug(:init, cuDNN) || Base.JLOptions().debug_level >= 2)

No idea how to check those conditions from the outside to dig deeper into it.

[1] https://github.com/JuliaGPU/CUDA.jl/blob/12caaf5d4d31a26f00039770ad07fa65cc75bee1/lib/cudnn/src/cuDNN.jl#L160

[fps]

OK, isdebug in the end checks for

Base.CoreLogging.current_logger_for_env(Base.CoreLogging.Debug, :init, CUDA) != nothing

which is true because: Base.CoreLogging.current_logger_for_env(Base.CoreLogging.Debug, :init, CUDA) gives PlutoRunner.PlutoCellLogger.... I wonder if the check in cuDNN shouldn't rather be:

if !precompiling && (isdebug(:init, cuDNN) && Base.JLOptions().debug_level >= 2)

but I am not familiar with Julia's logging/debug eco system at all.

[fps]

Part of the story seems to have been on the cuda.jl side: JuliaGPU/CUDA.jl#2750

We still pay a performance penalty due to debug logging being enabled in the first place. And also even with this fix on the cuda.jl side we miss lots of these messages, i.e. pluto's debug logger is still a bit fishy.

[fps]

OK, the CUDA.jl side seems to be mostly sorted out:

JuliaGPU/CUDA.jl#2750

and related:

JuliaGPU/CUDA.jl#2754

With that second PR the issue is gone completely. I'll open up an issue about the performance hit because of Pluto's debug logger.