# Flux runs out of memory

**URL:** <https://discourse.julialang.org/t/flux-runs-out-of-memory/73046>\
**Category:** Machine Learning\
**Tags:** memory-allocation, flux\
**Created:** [December 13, 2021, 7:17pm UTC](https://discourse.julialang.org/t/flux-runs-out-of-memory/73046 "2021-12-13T19:17:48Z")\
**Posts on this page:** 6\
**Page:** 2

<div class="post-metadata">

**Author:** ![jeremiedb](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jeremiedb/32/29150_2.png) [@jeremiedb](https://discourse.julialang.org/u/jeremiedb)\
**Post date:** [January 14, 2023, 6:27am UTC](https://discourse.julialang.org/t/flux-runs-out-of-memory/73046/21 "2023-01-14T06:27:40Z")

</div>

Also just to confirm, had you added the following block:

```julia
using ChainRulesCore
import ChainRulesCore: rrule

function ChainRulesCore.rrule(cfg::RuleConfig, c::Chain, x::AbstractArray)
    duo = accumulate(c.layers; init=(x, nothing)) do (input, _), layer
        out, back = rrule_via_ad(cfg, layer, input)
    end
    outs = map(first, duo)
    backs = map(last, duo)
    function un_chain(dout)
        multi = accumulate(reverse(backs); init=(nothing, dout)) do (_, delta), back
            dlayer, din = back(delta)
        end
        layergrads =
            foreach(CUDA.unsafe_free!, outs)
        foreach(CUDA.unsafe_free!, map(last, multi[1:end-1]))
        return (Tangent{Chain}(; layers=reverse(map(first, multi))), last(multi[end]))
    end
    outs[end], un_chain
end
# Could restrict this to x::CuArray... for testing instead write NaN into non-CuArrays, piratically:
CUDA.unsafe_free!(x::Array) = fill!(x, NaN)
CUDA.unsafe_free!(x::Flux.Zygote.Fill) = nothing

```

This has been the major helper that @mcabbott provided which roughly doubled the batch size capacity without having to disable CUDA pool which had significant adverse effect on performance.

---

<div class="post-metadata">

**Author:** ![DrChainsaw](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/drchainsaw/32/8497_2.png) [@DrChainsaw](https://discourse.julialang.org/u/DrChainsaw)\
**Post date:** [January 14, 2023, 9:07am UTC](https://discourse.julialang.org/t/flux-runs-out-of-memory/73046/22 "2023-01-14T09:07:25Z")

</div>

I sometimes wonder if one could take this even further and provide a callback function which gets the gradient and the thing the gradient is for instead of returning them.

One could then immediately apply any optimizer update and then free parameter gradients. Probably alot of things which does not work out with it (non-linear optimizers and parameters reoccuring?) though so it might be difficult to make generic.

---

<div class="post-metadata">

**Author:** ![LucasMSpereira](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lucasmspereira/32/26393_2.png) [@LucasMSpereira](https://discourse.julialang.org/u/LucasMSpereira)\
**Post date:** [January 14, 2023, 12:02pm UTC](https://discourse.julialang.org/t/flux-runs-out-of-memory/73046/23 "2023-01-14T12:02:25Z")

</div>

Now I tried including this block. Same result.

---

<div class="post-metadata">

**Author:** ![itan1](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/itan1/32/44584_2.png) [@itan1](https://discourse.julialang.org/u/itan1)\
**Post date:** [May 31, 2023, 7:38am UTC](https://discourse.julialang.org/t/flux-runs-out-of-memory/73046/24 "2023-05-31T07:38:56Z")

</div>

Hey everyone, I am running into the same problem.

I am training a residual U-Net for 3D image segmentation with FastAI.jl on GCloud with a 16GB T4 GPU but keep getting out of memory problems on the GPU. After searching online I made sure to set `JULIA_CUDA_MEMORY_POOL` to “none” and added a callback after every epoch that runs `GC.gc(true)` and `CUDA.reclaim()` . I think the problem is a memory leak, as it only occurs after ~30 epochs. I can also see the GPU utility dropping before the crash (see dashboard screenshot below). When I decrease the input image size it happens later. When I decrease the model size it happens earlier (this I really do not understand?). I posted the question [here on StackOverflow](https://stackoverflow.com/questions/76339447/stepwise-decrease-gpu-utility-followed-by-out-of-memory-error). Does anyone have an idea what could be the problem or what I can try to figure it out?

**Some context info:**

- I am using Julia 1.9.0 with FastAI.jl 0.5.1, Flux.jl 0.13.16, CUDA.jl 4.2.0
- The VM is Ubuntu 22.04 x86\_64 with CUDA toolkit 12.1, NVIDIA driver 530.30.02, and an NVIDIA Tesla T4 GPU with 16GB RAM
- The model is a residual U-Net with approximately 9.5 million parameters, the input data are 3D Float32 images with size (96, 96, 96) and I am using a batch size of 2.

**Some things I’ve tried:**

- I can reproduce the behaviour reliably, it happens every time after the same amount of epochs
- If I decrease the input image size, it still happens but later (epoch 60)
- If I decrease the model size it happens earlier (this I especially don’t understand)
- I’ve set `JULIA_CUDA_MEMORY_POOL` to `none` and added a callback after each epoch that executes `GC.gc(true)` and `CUDA.reclaim()`
- I’ve upgraded the GPU to an NVIDIA L4 with 24GB RAM. With this GPU I can train until epoch 125 but then the same thing happens (gradual decrease of GPU utility and finally an OOM error). I can work with this for now, simply restarting the training after every crash, but it is not ideal 😕

This is a screenshot of the dashboard with VM and GPU metrics, you can see the GPU utility dropping in several steps just before the crash:

 ![OOM](https://global.discourse-cdn.com/julialang/original/3X/7/8/788fd68b1a1521e957b035a5a0a4ceec960c4fa2.png)

---

<div class="post-metadata">

**Author:** ![DrChainsaw](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/drchainsaw/32/8497_2.png) [@DrChainsaw](https://discourse.julialang.org/u/DrChainsaw)\
**Post date:** [May 31, 2023, 3:22pm UTC](https://discourse.julialang.org/t/flux-runs-out-of-memory/73046/25 "2023-05-31T15:22:09Z")

</div>

Here is one possible reason which matches your symptoms quite well: [CUDA memory leak for Flux.Optimizer · Issue #148 · FluxML/FluxTraining.jl · GitHub](https://github.com/FluxML/FluxTraining.jl/issues/148)

---

<div class="post-metadata">

**Author:** ![itan1](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/itan1/32/44584_2.png) [@itan1](https://discourse.julialang.org/u/itan1)\
**Post date:** [June 1, 2023, 6:35am UTC](https://discourse.julialang.org/t/flux-runs-out-of-memory/73046/26 "2023-06-01T06:35:35Z")

</div>

Thank you very much, that was it! I replaced my `Flux.Nesterov` optimiser with `Optimisers.Nesterov` from Optimisers.jl. This new training run is already training longer than I managed before and I do not see those symptoms above anymore.

[Previous page](https://discourse.julialang.org/t/flux-runs-out-of-memory/73046.md?page=1)