Chaining gpu and cpu models

[lungd]

Is the following a bug or not supposed to work?

using Flux
using CUDA
m = Chain(Dense(1,1)|>gpu, cpu, Dense(1,1))
p, re = Flux.destructure(m)
p
# 4-element CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}:    ##  <-- params of 2nd Dense not on GPU
#  -0.3475855
#   0.0
#  1.0547627
#  0.0

re(p)([1]|>gpu)
ERROR: ArgumentError: cannot take the CPU address of a CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}
Stacktrace:
  [1] unsafe_convert(#unused#::Type{Ptr{Float32}}, x::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
    @ CUDA ~/.julia/packages/CUDA/Uurn4/src/array.jl:315
  [2] gemv!(trans::Char, alpha::Float32, A::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, X::Vector{Float32}, beta::Float32, Y::Vector{Float32})
    @ LinearAlgebra.BLAS ~/julia-1.7.2/share/julia/stdlib/v1.7/LinearAlgebra/src/blas.jl:667
  [3] gemv!(y::Vector{Float32}, tA::Char, A::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, x::Vector{Float32}, α::Bool, β::Bool)
    @ LinearAlgebra ~/julia-1.7.2/share/julia/stdlib/v1.7/LinearAlgebra/src/matmul.jl:544
  [4] mul!
    @ ~/julia-1.7.2/share/julia/stdlib/v1.7/LinearAlgebra/src/matmul.jl:66 [inlined]
  [5] mul!
    @ ~/julia-1.7.2/share/julia/stdlib/v1.7/LinearAlgebra/src/matmul.jl:275 [inlined]
  [6] *
    @ ~/julia-1.7.2/share/julia/stdlib/v1.7/LinearAlgebra/src/matmul.jl:47 [inlined]
  [7] (::Dense{typeof(identity), CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}})(x::Vector{Float32})
    @ Flux ~/.julia/packages/Flux/18YZE/src/layers/basic.jl:160
  [8] macro expansion
    @ ~/.julia/packages/Flux/18YZE/src/layers/basic.jl:53 [inlined]
  [9] applychain(layers::Tuple{Dense{typeof(identity), CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, typeof(cpu), Dense{typeof(identity), CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}, x::CUDA.CuArray{Int64, 1, CUDA.Mem.DeviceBuffer})
    @ Flux ~/.julia/packages/Flux/18YZE/src/layers/basic.jl:53
 [10] (::Chain{Tuple{Dense{typeof(identity), CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, typeof(cpu), Dense{typeof(identity), CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}})(x::CUDA.CuArray{Int64, 1, CUDA.Mem.DeviceBuffer})
    @ Flux ~/.julia/packages/Flux/18YZE/src/layers/basic.jl:51
 [11] top-level scope
    @ REPL[11]:1
 [12] top-level scope
    @ ~/.julia/packages/CUDA/Uurn4/src/initialization.jl:52

I think the parameter vector will always get typed with the type of the first layer, due to map ?

In my example I have a chain with a nested ODE layer and I think it would be better to keep it on the CPU (not sure though).

Tested with Flux 0.12.10 and 0.13

[DhairyaLGandhi]

What's happening here is that destructure should error on trying to vcat cpu and GPU arrays. Else, the second function in the model puts the output from the first layer on the cpu but the weights on the second layer are now on the GPU.

[lungd]

destructure should error

That would be nice.
So you are saying simply chaining layers with weights on CPU and GPU is not supposed to work, right?
I guess something like above would be possible in general but not without massive performance hit (Any[])?

[ToucheSir]

It is possible in general, but the limiting assumption of destructure is that every parameter in the model can fit in a homogeneously-typed vector. Evidently that does not hold if some parameters are in cpu memory and others in gpu memory. Can you describe a little more about your use case and how it makes use of both mixed-device params + destructure?

[mcabbott]

There's no special handling of this at present, nor tests. destructure itself will work or fail when reduce(hcat, arrays) does (i.e. when the first is a CuArray), but reconstruction and gradients may or may not.

Making one CuArray doesn't seem crazy when the whole model is there. But what should happen in mixed cases, ideally?

[lungd]

Sorry for the late reply.

@ToucheSir I am playing with a customizable encoder-decoder model wrapper which, in my use case, maps images to actions. A CNN encoder maps RGB images to some latent representation which then gets interpreted as some states of an ODE. The odemodel is not a common NODE with a MLP (or another architecture which makes use of many matrix multiplications) approximating the derivatives of the hidden states. It's rather a system of conductance-based neuron models with scalar operations. Part of the ODE solution then gets mapped using a decoder.

Here's some pseudo-like-code

struct EncDec enc; odemodel; dec; Ps; REs; end

function EncDec(enc, odemodel, dec, ...)
  p_enc, re_enc = Flux.destructure(enc)                                     # p_enc will be a CuArray
  p_odemodel, re_odemodel = Flux.destructure(odemodel)       # Float32[], or SparseArrays
  p_dec, re_dec = Flux.destructure(dec)                                     # Float32[] but may be CuArray
end

trainable(m::EncDec) = (p_enc = m.p_enc, ...)

function (m::EncDec)(x)
  p_enc, p_odemodel, p_dec = m.p_enc, ...
  z = m.re_enc(p_enc)(x) |> cpu
  dudt(du,u,p_odemodel,t) = du .= m.re_odemodel(p_odemodel)(...)
  u0 = vcat(z[1], ...)
  sol = solve EnsembleProblem (trajectories=batchsize) with PresetTimeCallbacks to update integrator.u[1:size(z,1)] = zt(t)
  out = m.re_dec(p_dec)(sol)   # may be inside dudt with additional states for the output
end

encdec = EncDec(Chain() |> gpu, ..., Chain() |> cpu, ...)

As you can see, I am currently using Flux/Zygote with implicit parameters. With ps = Flux.params(encdec) there is no problem because, if I am not completely wrong, the output is just a list of pointers pointing to the actual parameter which can by typed differently.

I still need to investigate at what size it would make sense to solve the ODE on GPU. For the size I've tried the performance was worse with u0 as a CuArray. One reason for that is probably that the input is a sequence of images and the largest batchsize I can use without running out of GPU memory is too small.

p_encdec, re_encdec = Flux.destructure(encdec) did work and returned the correct number of parameters, so I expected re_encdec(p_encdec) to work as well. But it failed and now I am looking for a good approach to solve this issue.

I could destructure enc, odemodel, and dec separately, define (m::EncDec)(x, p_enc, p_odemodel, p_dec) and compute gradients with gradient((p1,p2,p3) -> encdec(x, p1, p2, p3), p_enc, p_odemodel, p_dec)
Or maybe I should move parameters onto CPU when needed, i.e. re_odemodel(cpu(p_encdec[p_start:p_end]))?

TLDR;
Maybe I am thinking in a wrong way, but let's assume there are good reasons for defining models with heterogeneously-typed parameters, i.e. unrelated to CuArrays, some of the weights may be defined as sparse arrays.
How to handle such an use case?

[ToucheSir]

As you can see, I am currently using Flux/Zygote with implicit parameters. With ps = Flux.params(encdec) there is no problem because, if I am not completely wrong, the output is just a list of pointers pointing to the actual parameter which can by typed differently.
Maybe I am thinking in a wrong way, but let's assume there are good reasons for defining models with heterogeneously-typed parameters, i.e. unrelated to CuArrays, some of the weights may be defined as sparse arrays.

These are not a problem with the explicit parameter format used by Optimisers.jl either. The trouble comes when you try to combine it with destructure.

I could destructure enc, odemodel, and dec separately, define (m::EncDec)(x, p_enc, p_odemodel, p_dec) and compute gradients with gradient((p1,p2,p3) -> encdec(x, p1, p2, p3), p_enc, p_odemodel, p_dec)

One thing to consider is putting your ODE model parameters in a ComponentArray. Then there is no need for destructuring or restructuring at all. Even if that's not possible, you should be able to get away with not destructuring enc or dec, since at no point does the code above rely on their parameters being in a flat vector.

[lungd]

I will give the ComponentArrays a try. I can see how this could solve some of my issues. Thanks for the pointer!
Yes, you are right. There is no need to destructure enc and dec. I should have put them inside the ODE to show an example where that would have been necessary.

Thank you for your comments. I think I know now what I have to do.

I guess it was just me not using some packages properly and you can close this issue and sorry for the noise.

[ToucheSir]

No worries, this is something that has come up pre-explicit params and I think it's great that we've brainstormed a solution here. I'm sure this issue will be useful for someone in future. Feel free to reopen if you still run into AD+GPU issues with the proposed approach :)