Apparent memory leak when using Distributed?

[jarbus]

Package Version

Status ~/.julia/environments/v1.8/Project.toml [587475ba] Flux v0.13.8 https://github.com/FluxML/Flux.jl.git#master

Julia Version

julia version 1.8.2

OS / Environment

OS: Arch Linux x86_64
Kernel: 6.0.6-arch1-1

Describe the bug

Flux appears to use an increasing amount of memory when it's reconstruction function and Adam optimizer is used with Distributed.jl.

Steps to Reproduce

Warning: This will quickly eat up your ram
I've identified four conditions must be met in order for this leak to occur:

1. Must be using Distributed
2. Must be using an optimizer
3. Must be running update! on a worker
4. Must include reconstruction on a worker

using Distributed
addprocs(1)
@everywhere begin
  using Flux
  opt = Adam()
  theta, re = Chain(Dense(1000 => 1000, tanh)) |> Flux.destructure
end

println("Beginning loop")

for i in 1:1000
  println(i)

  # Note: bug doesn't occur if this remotecall is on thread 1
  fetch(remotecall(2) do
    # Note: bug doesn't occur if you remove below line
    re(theta)
    1
  end)

  @everywhere begin
    grad = theta * 0.01
    # Note: bug doesn't occur if you switch out update! with this line
    #theta .+= grad  
    Flux.Optimise.update!(opt, theta, grad)
  end
end

Expected Results

This bit of code should run to completion, using a relatively consistent amount of memory during each iteration of the loop, as nothing is saved during each loop.

Observed Results

Memory usage quickly increases indefinitely

Relevant log output

No response

[ToucheSir]

Using a slightly modified version of the code to show what's going on:

Modified code:

using Distributed
addprocs(1)
@everywhere begin
  using Flux
  opt = Adam()
  theta, re = Chain(Dense(1 => 1, tanh)) |> Flux.destructure
  @show objectid(theta)
end

println("\nBeginning loop\n")

for i in 1:5
  println(i)

  @show objectid(theta)
  # Note: bug doesn't occur if this remotecall is on thread 1
  fetch(remotecall(2) do
    # Note: bug doesn't occur if you remove below line
    theta1 = theta
    @show objectid(theta1)
  end)

  @everywhere begin
    theta2 = theta
    @show objectid(theta2)
    opt_param_oids = [objectid(k) for k in keys(opt.state)]
    @show opt_param_oids
    grad = theta2 * 0.01
    # Note: bug doesn't occur if you switch out update! with this line
    #theta .+= grad
    Flux.Optimise.update!(opt, theta2, grad)
  end
end

Output:

1
objectid(theta) = 0x9993d0b4b67d6070
      From worker 2:	objectid(theta1) = 0xd58cd9fdecd282b0
objectid(theta2) = 0x9993d0b4b67d6070
opt_param_oids = UInt64[]
      From worker 2:	objectid(theta2) = 0xd58cd9fdecd282b0
      From worker 2:	opt_param_oids = UInt64[]
2
objectid(theta) = 0x9993d0b4b67d6070
      From worker 2:	objectid(theta1) = 0xbb76db5e563635a9
objectid(theta2) = 0x9993d0b4b67d6070
opt_param_oids = UInt64[0x9993d0b4b67d6070]
      From worker 2:	objectid(theta2) = 0xbb76db5e563635a9
      From worker 2:	opt_param_oids = UInt64[0xd58cd9fdecd282b0]
3
objectid(theta) = 0x9993d0b4b67d6070
      From worker 2:	objectid(theta1) = 0x60003cba1e4994be
objectid(theta2) = 0x9993d0b4b67d6070
opt_param_oids = UInt64[0x9993d0b4b67d6070]
      From worker 2:	objectid(theta2) = 0x60003cba1e4994be
      From worker 2:	opt_param_oids = UInt64[0xd58cd9fdecd282b0, 0xbb76db5e563635a9]
4
objectid(theta) = 0x9993d0b4b67d6070
      From worker 2:	objectid(theta1) = 0x9a29f970928150b6
objectid(theta2) = 0x9993d0b4b67d6070
opt_param_oids = UInt64[0x9993d0b4b67d6070]
      From worker 2:	objectid(theta2) = 0x9a29f970928150b6
      From worker 2:	opt_param_oids = UInt64[0xd58cd9fdecd282b0, 0xbb76db5e563635a9, 0x60003cba1e4994be]
5
objectid(theta) = 0x9993d0b4b67d6070
      From worker 2:	objectid(theta1) = 0x00319adddf1c0a67
objectid(theta2) = 0x9993d0b4b67d6070
opt_param_oids = UInt64[0x9993d0b4b67d6070]
      From worker 2:	objectid(theta2) = 0x00319adddf1c0a67
      From worker 2:	opt_param_oids = UInt64[0xd58cd9fdecd282b0, 0x9a29f970928150b6, 0xbb76db5e563635a9, 0x60003cba1e4994be]

The culprit here is that the identity (i.e. pointer, the values are still ==) of the captured theta in worker 2 changes from iteration to iteration. Because Flux's optimizers use IdDicts internally to keep track of state for previously seen parameters, worker 2 will constantly allocate new state because it thinks it's seeing a new theta on every loop.

I don't understand Distributed well enough to know why theta is moved again to worker 2 on every iteration instead of using the variable defined in the first @everywhere block (do remotecall captures unconditionally cause data movement?), but thankfully there is a solution. https://github.com/FluxML/Optimisers.jl does not have this problem and supports all the optimizers Flux currently does, so you can use it instead. You'll also be future-proofing your code because Flux will switch to using Optimisers.jl instead of the current IdDict-based optimizers in the next major version (it already imports it now, so you don't have to worry about the overhead of another dependency).

[jarbus]

Thank you very much for your time and in-depth response, @ToucheSir. This is incredibly helpful, and I wouldn't have been able to diagnose this issue with the knowledge of Julia I had. It's good to know that this won't be an issue in future versions of Flux.

I'm going to play around with Distributed and see if I can avoid re-sending theta in the first place, as the whole point of the algorithm I'm implementing is that it doesn't require sending model parameters between processes, just their fitness, essentially eliminating IPC overhead. From what I can tell, it's difficult to detect if there is implicit data transfer between processes in Distributed when both processes have a variable with the same name, but that's a problem for a different forum.

[jarbus]

Fixed theta getting sent by referring to it as a global in a named function, without passing it as a parameter:

fetch(remotecall(2) do
    re(theta)
    1
end)

changes to

function fitness()
re(theta)
1
end
fetch(remotecall(2) do
  fitness()
end)