# MPI.jl + Zygote OOM on embarrassingly parallel tasks

**URL:** <https://discourse.julialang.org/t/mpi-jl-zygote-oom-on-embarrassingly-parallel-tasks/110622>\
**Category:** Julia at Scale\
**Created:** [February 23, 2024, 4:21am UTC](https://discourse.julialang.org/t/mpi-jl-zygote-oom-on-embarrassingly-parallel-tasks/110622 "2024-02-23T04:21:09Z")\
**Posts on this page:** 12\
**Page:** 1

<div class="post-metadata">

**Author:** ![richardr2926](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/richardr2926/32/207243_2.png) [@richardr2926](https://discourse.julialang.org/u/richardr2926)\
**Post date:** [February 23, 2024, 4:21am UTC](https://discourse.julialang.org/t/mpi-jl-zygote-oom-on-embarrassingly-parallel-tasks/110622/1 "2024-02-23T04:21:09Z")

</div>

We have simple NN:

```julia
using MPI
using Zygote
using CUDA
using Flux
using LinearAlgebra

MPI.Init()

comm = MPI.COMM_WORLD
rank = MPI.Comm_rank(comm)
size = MPI.Comm_size(comm)

nx, ny, nz, nt = parse.(Int, ARGS[1:4])
T = Float32

x = rand(T, 5, nx*ny*nz*nt)
y = rand(T, 1, nx*ny*nz*nt)

weights = Dict(
    :w1 => rand(T, 20, 5),
    :w2 => rand(T, 128, 20),
    :w3 => rand(T, 1, 128)
)

x = x |> gpu
y = y |> gpu
weights = Dict(k => gpu(v) for (k, v) in pairs(weights))

function forward(weights, x)
    w1, w2, w3 = weights[:w1], weights[:w2], weights[:w3]
    return norm(relu.(w3 * relu.(w2 * (w1 * x))) - y)
end

forward(weights, x)
gradient_weights = Zygote.gradient(weights -> forward(weights, x), weights)

MPI.Finalize()

```

We are using NVIDIA A100 SMX 80GB GPUs.

Max problem size before OOMing:

Doing `nx, ny, nz, nt = 128, 128, 64, 15` works on one - two GPUs

We can only do up to `nx, ny, nz, nt = 128, 128, 64, 9` on 4 GPUs

Why does this scale so poorly ?

---

<div class="post-metadata">

**Author:** ![carstenbauer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/carstenbauer/32/4981_2.png) [@carstenbauer](https://discourse.julialang.org/u/carstenbauer)\
**Post date:** [February 23, 2024, 6:59pm UTC](https://discourse.julialang.org/t/mpi-jl-zygote-oom-on-embarrassingly-parallel-tasks/110622/2 "2024-02-23T18:59:41Z")

</div>

Native question: how do you make sure that each MPI rank uses a different GPU? Or what’s the idea here?

---

<div class="post-metadata">

**Author:** ![simonbyrne](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/simonbyrne/32/19_2.png) [@simonbyrne](https://discourse.julialang.org/u/simonbyrne)\
**Post date:** [February 24, 2024, 10:20pm UTC](https://discourse.julialang.org/t/mpi-jl-zygote-oom-on-embarrassingly-parallel-tasks/110622/3 "2024-02-24T22:20:41Z")

</div>

How much CPU memory do you have? If you’re not using any sort of resource manager (e.g. slurm), the easiest option would be to use the [`--heap-size-hint`](https://julialang.org/blog/2023/04/julia-1.9-highlights/#memory_usage_hint_for_the_gc_with_--heap-size-hint) option to limit each process to a fraction of the total memory.

---

<div class="post-metadata">

**Author:** ![simonbyrne](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/simonbyrne/32/19_2.png) [@simonbyrne](https://discourse.julialang.org/u/simonbyrne)\
**Post date:** [February 24, 2024, 10:22pm UTC](https://discourse.julialang.org/t/mpi-jl-zygote-oom-on-embarrassingly-parallel-tasks/110622/4 "2024-02-24T22:22:18Z")

</div>

One other thing worth trying is to avoid using CPU memory at all, and use [`CUDA.rand`](https://cuda.juliagpu.org/stable/usage/array/#Random-numbers) instead.

---

<div class="post-metadata">

**Author:** ![richardr2926](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/richardr2926/32/207243_2.png) [@richardr2926](https://discourse.julialang.org/u/richardr2926)\
**Post date:** [February 29, 2024, 8:27pm UTC](https://discourse.julialang.org/t/mpi-jl-zygote-oom-on-embarrassingly-parallel-tasks/110622/5 "2024-02-29T20:27:19Z")

</div>

We use slurm. so something like:

`salloc --nodes=1 --constraint=gpu --gpus=4 --qos=interactive --time=00:20:00 --ntasks=4 --gpus-per-task=1 --gpu-bind=none`

and then:

`mpiexecjl --project=./ julia-1.8 scaling.jl 128 128 64 15`

---

<div class="post-metadata">

**Author:** ![richardr2926](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/richardr2926/32/207243_2.png) [@richardr2926](https://discourse.julialang.org/u/richardr2926)\
**Post date:** [February 29, 2024, 8:42pm UTC](https://discourse.julialang.org/t/mpi-jl-zygote-oom-on-embarrassingly-parallel-tasks/110622/6 "2024-02-29T20:42:01Z")

</div>

Hello,

We do use slurm and allocate constant memory per task.

---

<div class="post-metadata">

**Author:** ![richardr2926](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/richardr2926/32/207243_2.png) [@richardr2926](https://discourse.julialang.org/u/richardr2926)\
**Post date:** [February 29, 2024, 8:45pm UTC](https://discourse.julialang.org/t/mpi-jl-zygote-oom-on-embarrassingly-parallel-tasks/110622/7 "2024-02-29T20:45:53Z")

</div>

Hey,

I have tried this as well. The error denotes that the issue arises when trying to allocate too much GPU memory:

`ERROR: LoadError: Out of GPU memory trying to allocate 7.500 GiB Effective GPU memory usage: 99.98% (79.136 GiB/79.154 GiB) Memory pool usage: 9.023 GiB (16.531 GiB reserved)`

---

<div class="post-metadata">

**Author:** ![richardr2926](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/richardr2926/32/207243_2.png) [@richardr2926](https://discourse.julialang.org/u/richardr2926)\
**Post date:** [February 29, 2024, 8:58pm UTC](https://discourse.julialang.org/t/mpi-jl-zygote-oom-on-embarrassingly-parallel-tasks/110622/8 "2024-02-29T20:58:53Z")

</div>

I would also like to note that binding or not binding the gpus to the tasks has no effect on the outcome

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [February 29, 2024, 11:00pm UTC](https://discourse.julialang.org/t/mpi-jl-zygote-oom-on-embarrassingly-parallel-tasks/110622/9 "2024-02-29T23:00:38Z")

</div>

It looks like all 4 GPUs are attached to the same node. This means CUDA.jl will use GPU ID 0 by default unless you tell it not to. `Flux.gpu` (which has nothing to do with Zygote, so I don’t think Zygote actually matters here), calls `CUDA.cu` under the hood. This means each of your tasks is trying to allocate memory on the same GPU, thus causing OOMs when they probably shouldn’t be.

When not using MPI, the best way to assign different GPUs per process/task would be to follow [Multiple GPUs · CUDA.jl](https://cuda.juliagpu.org/stable/usage/multigpu/). I’ll let the HPC experts in this thread comment on whether there’s something special which is available when using MPI.

---

<div class="post-metadata">

**Author:** ![simonbyrne](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/simonbyrne/32/19_2.png) [@simonbyrne](https://discourse.julialang.org/u/simonbyrne)\
**Post date:** [March 1, 2024, 9:08pm UTC](https://discourse.julialang.org/t/mpi-jl-zygote-oom-on-embarrassingly-parallel-tasks/110622/10 "2024-03-01T21:08:46Z")

</div>

Try

```julia
CUDA.device!(MPI.Comm_rank(MPI.COMM_WORLD))

```

---

<div class="post-metadata">

**Author:** ![richardr2926](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/richardr2926/32/207243_2.png) [@richardr2926](https://discourse.julialang.org/u/richardr2926)\
**Post date:** [March 1, 2024, 11:27pm UTC](https://discourse.julialang.org/t/mpi-jl-zygote-oom-on-embarrassingly-parallel-tasks/110622/11 "2024-03-01T23:27:51Z")

</div>

Hello @ToucheSir, thank you. this seems to be the issue. even when binding one gpu to one task through slurm, I noticed that all tasks are mapped to one gpu.

Thank you !

---

<div class="post-metadata">

**Author:** ![richardr2926](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/richardr2926/32/207243_2.png) [@richardr2926](https://discourse.julialang.org/u/richardr2926)\
**Post date:** [March 1, 2024, 11:29pm UTC](https://discourse.julialang.org/t/mpi-jl-zygote-oom-on-embarrassingly-parallel-tasks/110622/12 "2024-03-01T23:29:16Z")

</div>

Hey @simonbyrne, thank you this seems to be the solution. I will also note that for this to work, I had to unbind the gpus so they are visible to all tasks on a node (`--gpu-bind=none`)