# Flux Transformer Out of Memory

**URL:** <https://discourse.julialang.org/t/flux-transformer-out-of-memory/95922>\
**Category:** Machine Learning\
**Created:** [March 11, 2023, 4:38pm UTC](https://discourse.julialang.org/t/flux-transformer-out-of-memory/95922 "2023-03-11T16:38:55Z")\
**Posts on this page:** 6\
**Page:** 2

<div class="post-metadata">

**Author:** ![chengchingwen](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chengchingwen/32/8390_2.png) [@chengchingwen](https://discourse.julialang.org/u/chengchingwen)\
**Post date:** [March 12, 2023, 10:24pm UTC](https://discourse.julialang.org/t/flux-transformer-out-of-memory/95922/21 "2023-03-12T22:24:13Z")

</div>

> [@Gadersd](#):
>
> but the backpropagation itself uses around 80%

So you mean PyTorch only use 3% of your gpu (so that’s 180MB) during backpropagation?

---

<div class="post-metadata">

**Author:** ![Gadersd](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gadersd/32/47603_2.png) [@Gadersd](https://discourse.julialang.org/u/Gadersd)\
**Post date:** [March 12, 2023, 10:39pm UTC](https://discourse.julialang.org/t/flux-transformer-out-of-memory/95922/22 "2023-03-12T22:39:12Z")

</div>

> [@chengchingwen](#):
>
> So you mean PyTorch only use 3% of your gpu (so that’s 180MB) during backpropagation?

PyTorch uses that amount to store the model and all its gradients. I am not sure exactly how much is used during the calculations as some of it may be freed. PyTorch is about 4x faster though.

---

<div class="post-metadata">

**Author:** ![jpsamaroo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jpsamaroo/32/46804_2.png) [@jpsamaroo](https://discourse.julialang.org/u/jpsamaroo)\
**Post date:** [March 12, 2023, 10:40pm UTC](https://discourse.julialang.org/t/flux-transformer-out-of-memory/95922/23 "2023-03-12T22:40:12Z")

</div>

Something to consider is that 80% is not necessarily the amount that backpropagation needs; instead, it’s more like how much memory the GC felt like holding onto. Manual usage of `CUDA.unsafe_free!` would probably bring this number down significantly, but doing so in Flux directly isn’t really feasible in most cases (because who knows if the array might be accessed somewhere else).

I think it would be helpful to first figure out where these allocations are happening, and then determine if any of those allocations can be quickly `unsafe_free!`'d after usage.

---

<div class="post-metadata">

**Author:** ![chengchingwen](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chengchingwen/32/8390_2.png) [@chengchingwen](https://discourse.julialang.org/u/chengchingwen)\
**Post date:** [March 12, 2023, 10:45pm UTC](https://discourse.julialang.org/t/flux-transformer-out-of-memory/95922/24 "2023-03-12T22:45:31Z")

</div>

Unfortunately, currently this cannot be done because the ChainRules is design to be used for gradient, jacobian, and also higher order AD. The memory management become tricky with those assumption.

---

<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:** [March 13, 2023, 1:44am UTC](https://discourse.julialang.org/t/flux-transformer-out-of-memory/95922/25 "2023-03-13T01:44:13Z")

</div>

Not sure if it’s relevant in this specific context, but I’m wondering if @mcabbott WIP improvement to ChainRules’s CUDA memory management couldn"t be useful here: [Flux.jl/cuda.jl at 494d351c15bec709af7049ccc6a2ffec4580f682 · mcabbott/Flux.jl · GitHub](https://github.com/mcabbott/Flux.jl/blob/494d351c15bec709af7049ccc6a2ffec4580f682/src/cuda/cuda.jl#L13-L38)

The avbove `rrule` for Chain did brought significant improvement to the memory mgmt of Resnet models (more than doubled max batch size capacity IIRC).

---

<div class="post-metadata">

**Author:** ![Gadersd](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gadersd/32/47603_2.png) [@Gadersd](https://discourse.julialang.org/u/Gadersd)\
**Post date:** [March 13, 2023, 2:36pm UTC](https://discourse.julialang.org/t/flux-transformer-out-of-memory/95922/26 "2023-03-13T14:36:13Z")

</div>

I am not sure if this is relevant but `CUDA.memory_status()` shows `No memory pool is in use.nothing`. My understanding is that a memory pool is normally used by default and I never turned it off.

Also, I just did a more thorough benchmark of the PyTorch implementation and it never goes over 2.6 GB of GPU usage. With a batch size of 64 the PyTorch version never goes over 4.4 GB. The Flux version crashes with a batch size of 58 so that would be a memory usage of around 6 GB. It seems that Flux uses about 36% more memory than PyTorch in this case. Is this relative memory usage difference common or is Flux usually closer to PyTorch is terms of performance/memory usage?

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