Error in `update!` for Metal arrays and Adam optimiser

[CarloLucibello]

The following lines using Descent run fine

julia> using Optimisers, Flux

julia> x = rand(Float32, 3) |> Flux.gpu
3-element MtlVector{Float32, Metal.MTL.MTLResourceStorageModePrivate}:
 0.6247813
 0.75746953
 0.9677551

julia> g = rand(Float32, 3) |> Flux.gpu
3-element MtlVector{Float32, Metal.MTL.MTLResourceStorageModePrivate}:
 0.22261995
 0.24903625
 0.07521629

julia> opt = Optimisers.setup(Optimisers.Descent(1e-3), x)
Leaf(Descent{Float64}(0.001), nothing)

julia> Optimisers.update!(opt, x, g) # OK
(Leaf(Descent{Float64}(0.001), nothing), Float32[0.6245587, 0.7572205, 0.96767986])

Now with Adam instead:

julia> opt = Optimisers.setup(Optimisers.Adam(1e-3), x)
Leaf(Adam{Float64}(0.001, (0.9, 0.999), 2.22045e-16), (Float32[0.0, 0.0, 0.0], Float32[0.0, 0.0, 0.0], (0.9, 0.999)))

julia> Optimisers.update!(opt, model, g)
ERROR: InvalidIRError: compiling MethodInstance for (::GPUArrays.var"#broadcast_kernel#26")(::Metal.mtlKernelContext, ::MtlDeviceVector{Float32, 1}, ::Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{1}, Tuple{Base.OneTo{Int64}}, typeof(+), Tuple{Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{1}, Nothing, typeof(*), Tuple{Float64, Base.Broadcast.Extruded{MtlDeviceVector{Float32, 1}, Tuple{Bool}, Tuple{Int64}}}}, Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{1}, Nothing, typeof(*), Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0}, Nothing, typeof(-), Tuple{Int64, Float64}}, Base.Broadcast.Extruded{MtlDeviceVector{Float32, 1}, Tuple{Bool}, Tuple{Int64}}}}}}, ::Int64) resulted in invalid LLVM IR
Reason: unsupported unsupported use of double value
Reason: unsupported unsupported use of double value
Reason: unsupported unsupported use of double value
Reason: unsupported unsupported use of double value
Reason: unsupported unsupported use of double value
Stacktrace:
  [1] Float64
    @ ./float.jl:261
  [2] convert
    @ ./number.jl:7
  [3] _promote
    @ ./promotion.jl:358
  [4] promote
    @ ./promotion.jl:381
  [5] *
    @ ./promotion.jl:411
  [6] _broadcast_getindex_evalf
    @ ./broadcast.jl:683
  [7] _broadcast_getindex
    @ ./broadcast.jl:656
  [8] _getindex
    @ ./broadcast.jl:679
  [9] _broadcast_getindex
    @ ./broadcast.jl:655
 [10] getindex
    @ ./broadcast.jl:610
 [11] broadcast_kernel
    @ ~/.julia/packages/GPUArrays/5XhED/src/host/broadcast.jl:59
Reason: unsupported unsupported use of double value
...
Stacktrace:
 [1] Float32
   @ ./float.jl:258
 [2] convert
   @ ./number.jl:7
 [3] setindex!
   @ ~/.julia/packages/Metal/qeZqc/src/device/array.jl:105
 [4] setindex!
   @ ~/.julia/packages/Metal/qeZqc/src/device/array.jl:118
 [5] broadcast_kernel
   @ ~/.julia/packages/GPUArrays/5XhED/src/host/broadcast.jl:59
Hint: catch this exception as `err` and call `code_typed(err; interactive = true)` to introspect the erronous code with Cthulhu.jl
Stacktrace:
  [1] check_ir(job::GPUCompiler.CompilerJob{GPUCompiler.MetalCompilerTarget, Metal.MetalCompilerParams}, args::LLVM.Module)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/YO8Uj/src/validation.jl:149
  [2] macro expansion
    @ ~/.julia/packages/GPUCompiler/YO8Uj/src/driver.jl:415 [inlined]
  [3] macro expansion
    @ ~/.julia/packages/TimerOutputs/RsWnF/src/TimerOutput.jl:253 [inlined]
  [4] macro expansion
    @ ~/.julia/packages/GPUCompiler/YO8Uj/src/driver.jl:414 [inlined]
  [5] emit_llvm(job::GPUCompiler.CompilerJob; libraries::Bool, toplevel::Bool, optimize::Bool, cleanup::Bool, only_entry::Bool, validate::Bool)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/YO8Uj/src/utils.jl:89
  [6] emit_llvm
    @ ~/.julia/packages/GPUCompiler/YO8Uj/src/utils.jl:83 [inlined]
  [7] codegen(output::Symbol, job::GPUCompiler.CompilerJob; libraries::Bool, toplevel::Bool, optimize::Bool, cleanup::Bool, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/YO8Uj/src/driver.jl:129
  [8] codegen
    @ ~/.julia/packages/GPUCompiler/YO8Uj/src/driver.jl:110 [inlined]
  [9] compile(target::Symbol, job::GPUCompiler.CompilerJob; libraries::Bool, toplevel::Bool, optimize::Bool, cleanup::Bool, strip::Bool, validate::Bool, only_entry::Bool)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/YO8Uj/src/driver.jl:106
 [10] compile
    @ ~/.julia/packages/GPUCompiler/YO8Uj/src/driver.jl:98 [inlined]
 [11] #51
    @ ~/.julia/packages/Metal/qeZqc/src/compiler/compilation.jl:57 [inlined]
 [12] JuliaContext(f::Metal.var"#51#52"{GPUCompiler.CompilerJob{GPUCompiler.MetalCompilerTarget, Metal.MetalCompilerParams}})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/YO8Uj/src/driver.jl:47
 [13] compile(job::GPUCompiler.CompilerJob)
    @ Metal ~/.julia/packages/Metal/qeZqc/src/compiler/compilation.jl:56
 [14] actual_compilation(cache::Dict{Any, Any}, src::Core.MethodInstance, world::UInt64, cfg::GPUCompiler.CompilerConfig{GPUCompiler.MetalCompilerTarget, Metal.MetalCompilerParams}, compiler::typeof(Metal.compile), linker::typeof(Metal.link))
    @ GPUCompiler ~/.julia/packages/GPUCompiler/YO8Uj/src/execution.jl:125
 [15] cached_compilation(cache::Dict{Any, Any}, src::Core.MethodInstance, cfg::GPUCompiler.CompilerConfig{GPUCompiler.MetalCompilerTarget, Metal.MetalCompilerParams}, compiler::Function, linker::Function)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/YO8Uj/src/execution.jl:103
 [16] macro expansion
    @ ~/.julia/packages/Metal/qeZqc/src/compiler/execution.jl:162 [inlined]
 [17] macro expansion
    @ ./lock.jl:267 [inlined]
 [18] mtlfunction(f::GPUArrays.var"#broadcast_kernel#26", tt::Type{Tuple{Metal.mtlKernelContext, MtlDeviceVector{Float32, 1}, Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{1}, Tuple{Base.OneTo{Int64}}, typeof(+), Tuple{Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{1}, Nothing, typeof(*), Tuple{Float64, Base.Broadcast.Extruded{MtlDeviceVector{Float32, 1}, Tuple{Bool}, Tuple{Int64}}}}, Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{1}, Nothing, typeof(*), Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0}, Nothing, typeof(-), Tuple{Int64, Float64}}, Base.Broadcast.Extruded{MtlDeviceVector{Float32, 1}, Tuple{Bool}, Tuple{Int64}}}}}}, Int64}}; name::Nothing, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Metal ~/.julia/packages/Metal/qeZqc/src/compiler/execution.jl:157
 [19] mtlfunction(f::GPUArrays.var"#broadcast_kernel#26", tt::Type{Tuple{Metal.mtlKernelContext, MtlDeviceVector{Float32, 1}, Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{1}, Tuple{Base.OneTo{Int64}}, typeof(+), Tuple{Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{1}, Nothing, typeof(*), Tuple{Float64, Base.Broadcast.Extruded{MtlDeviceVector{Float32, 1}, Tuple{Bool}, Tuple{Int64}}}}, Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{1}, Nothing, typeof(*), Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0}, Nothing, typeof(-), Tuple{Int64, Float64}}, Base.Broadcast.Extruded{MtlDeviceVector{Float32, 1}, Tuple{Bool}, Tuple{Int64}}}}}}, Int64}})
    @ Metal ~/.julia/packages/Metal/qeZqc/src/compiler/execution.jl:155
 [20] macro expansion
    @ ~/.julia/packages/Metal/qeZqc/src/compiler/execution.jl:77 [inlined]
 [21] #launch_heuristic#98
    @ ~/.julia/packages/Metal/qeZqc/src/gpuarrays.jl:14 [inlined]
 [22] launch_heuristic
    @ ~/.julia/packages/Metal/qeZqc/src/gpuarrays.jl:12 [inlined]
 [23] _copyto!
    @ ~/.julia/packages/GPUArrays/5XhED/src/host/broadcast.jl:65 [inlined]
 [24] materialize!
    @ ~/.julia/packages/GPUArrays/5XhED/src/host/broadcast.jl:41 [inlined]
 [25] materialize!
    @ ./broadcast.jl:881 [inlined]
 [26] macro expansion
    @ ~/.julia/packages/Optimisers/1x8gl/src/interface.jl:201 [inlined]
 [27] apply!(o::Optimisers.Adam{Float64}, state::Tuple{MtlVector{Float32, Metal.MTL.MTLResourceStorageModePrivate}, MtlVector{Float32, Metal.MTL.MTLResourceStorageModePrivate}, Tuple{Float64, Float64}}, x::MtlVector{Float32, Metal.MTL.MTLResourceStorageModePrivate}, dx::MtlVector{Float32, Metal.MTL.MTLResourceStorageModePrivate})
    @ Optimisers ~/.julia/packages/Optimisers/1x8gl/src/rules.jl:213
 [28] _update!(ℓ::Optimisers.Leaf{Optimisers.Adam{Float64}, Tuple{MtlVector{Float32, Metal.MTL.MTLResourceStorageModePrivate}, MtlVector{Float32, Metal.MTL.MTLResourceStorageModePrivate}, Tuple{Float64, Float64}}}, x::MtlVector{Float32, Metal.MTL.MTLResourceStorageModePrivate}; grads::IdDict{Optimisers.Leaf, Any}, params::IdDict{Any, Any})
    @ Optimisers ~/.julia/packages/Optimisers/1x8gl/src/interface.jl:92
 [29] _update!
    @ ~/.julia/packages/Optimisers/1x8gl/src/interface.jl:88 [inlined]
 [30] update!(::Optimisers.Leaf{Optimisers.Adam{Float64}, Tuple{MtlVector{Float32, Metal.MTL.MTLResourceStorageModePrivate}, MtlVector{Float32, Metal.MTL.MTLResourceStorageModePrivate}, Tuple{Float64, Float64}}}, ::MtlVector{Float32, Metal.MTL.MTLResourceStorageModePrivate}, ::MtlVector{Float32, Metal.MTL.MTLResourceStorageModePrivate})
    @ Optimisers ~/.julia/packages/Optimisers/1x8gl/src/interface.jl:73
 [31] top-level scope
    @ REPL[31]:1

[mcabbott]

Metal really doesn't like Float64 even as an intermediate step, not an array:

julia> using Metal

julia> x = MtlArray([1 2 3f0])
1×3 MtlMatrix{Float32, Metal.MTL.MTLResourceStorageModePrivate}:
 1.0  2.0  3.0

julia> Float32.(x .+ 4.0)
ERROR: InvalidIRError: compiling MethodInstance for (::GPUArrays.var"#broadcast_kernel#26")(::Metal.mtlKernelContext, ::MtlDeviceMatrix{Float32, 1}, ::Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{2}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}, Metal.var"#86#87"{Float32}, Tuple{Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{2}, Nothing, typeof(+), Tuple{Base.Broadcast.Extruded{MtlDeviceMatrix{Float32, 1}, Tuple{Bool, Bool}, Tuple{Int64, Int64}}, Float64}}}}, ::Int64) resulted in invalid LLVM IR
Reason: unsupported unsupported use of double value
Reason: unsupported unsupported use of double value

I guess the equivalent for CUDA is slow but we seldom care for Optimisers.

Seems very likely to work with opt = Optimisers.setup(Optimisers.Adam(1f-3), x) instead, but I've forgotten the magic incantation to change what gpu does. (Can't we automate this, so that loading any one GPU package just works?)

[CarloLucibello]

It turns out that I was using a Float64 learning rate, with Float32 works instead

julia> opt = Optimisers.setup(MyAdam(1f-3), x) |> Flux.gpu
Leaf(MyAdam{Float32}(0.001, (0.9, 0.999), 1.19209f-7), (Float32[0.0, 0.0, 0.0], Float32[0.0, 0.0, 0.0], (0.9, 0.999)))

julia> Optimisers.update!(opt, x, g)
(Leaf(MyAdam{Float32}(0.001, (0.9, 0.999), 1.19209f-7), (Float32[0.0368521, 0.0404, 0.0267676], Float32[0.000135806, 0.000163214, 7.16493f-5], (0.81, 0.998001))), Float32[0.7669642, 0.055360284, 0.8887157])

Can we make this more robust?

[mcabbott]

Xref also #119, Adam(0) fails, and #120, Adam(0.01) has a different epsilon to Adam(0.01f0).

Would probably be ideal to follow the eltype of the array (or the state arrays). That would add some complexity to each rule. But if we do that, then the rule need not have a type parameter at all, we could just store all learning rates as Float64 say.

[mcabbott]

Maybe this was closed prematurely, trying after #151:

julia> using Optimisers, Flux

julia> x = rand(Float32, 3) |> Flux.gpu
3-element MtlVector{Float32, Metal.MTL.MTLResourceStorageModePrivate}:
 0.42997134
 0.18900502
 0.7357338

julia> g = rand(Float32, 3) |> Flux.gpu;

julia> opt = Optimisers.setup(Optimisers.Descent(1e-3), x);

julia> Optimisers.update!(opt, x, g) # OK, as above
(Leaf(Descent(0.001), nothing), Float32[0.74836177, 0.57043684, 0.4284932])

julia> opt_adam = Optimisers.setup(Flux.Adam(), x)  # Flux.Adam() always made Float64
Leaf(Adam(0.001, (0.9, 0.999), 1.0e-8), (Float32[0.0, 0.0, 0.0], Float32[0.0, 0.0, 0.0], (0.9, 0.999)))

julia> Optimisers.update!(opt_adam, x, g)  # above there's a typo, model isn't defined
ERROR: InvalidIRError: compiling MethodInstance for (::GPUArrays.var"#broadcast_kernel#26")(::Metal.mtlKernelContext, ::MtlDeviceVector{Float32, 1}, ::Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{1}, Tuple{Base.OneTo{Int64}}, typeof(-), Tuple{Base.Broadcast.Extruded{MtlDeviceVector{Float32, 1}, Tuple{Bool}, Tuple{Int64}}, Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{1}, Tuple{Base.OneTo{Int64}}, typeof(*), Tuple{Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{1}, Nothing, typeof(/), Tuple{Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{1}, Nothing, typeof(/), Tuple{Base.Broadcast.Extruded{MtlDeviceVector{Float32, 1}, Tuple{Bool}, Tuple{Int64}}, Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0}, Nothing, typeof(-), Tuple{Int64, Float64}}}}, Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{1}, Nothing, typeof(+), Tuple{Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{1}, Nothing, typeof(sqrt), Tuple{Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{1}, Nothing, typeof(/), Tuple{Base.Broadcast.Extruded{MtlDeviceVector{Float32, 1}, Tuple{Bool}, Tuple{Int64}}, Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0}, Nothing, typeof(-), Tuple{Int64, Float64}}}}}}, Float32}}}}, Float32}}}}, ::Int64) resulted in invalid LLVM IR
Reason: unsupported unsupported use of double value
Reason: unsupported unsupported use of double value

 [23] materialize!(dest::Any, bc::Base.Broadcast.Broadcasted{<:Any})
    @ Base.Broadcast ./broadcast.jl:876 [inlined]
 [24] subtract!(x::MtlVector{…}, x̄::Base.Broadcast.Broadcasted{…})
    @ Optimisers ~/.julia/packages/Optimisers/TxzMn/src/interface.jl:103
 [25] _update!(ℓ::Optimisers.Leaf{…}, x::MtlVector{…}; grads::IdDict{…}, params::IdDict{…})
    @ Optimisers ~/.julia/packages/Optimisers/TxzMn/src/interface.jl:97
 [26] update!(::Optimisers.Leaf{…}, ::MtlVector{…}, ::MtlVector{…})
    @ Optimisers ~/.julia/packages/Optimisers/TxzMn/src/interface.jl:77

(jl_OTs5qt) pkg> st
Status `/private/var/folders/yq/4p2zwd614y59gszh7y9ypyhh0000gn/T/jl_OTs5qt/Project.toml`
  [587475ba] Flux v0.14.3 `https://github.com/FluxML/Flux.jl.git#master`
  [3bd65402] Optimisers v0.3.0

The problem isn't subtract!, it's that Adam still contains some Float64, which leaks into its lazy modified gradient:

Optimisers.jl/src/rules.jl

Lines 209 to 213 in 6a4f948

	init(o::Adam, x::AbstractArray) = (zero(x), zero(x), o.beta)
	function apply!(o::Adam, state, x::AbstractArray{T}, dx) where T
	η, β, ϵ = T(o.eta), T.(o.beta), T(o.epsilon)
	mt, vt, βt = state

julia> lazyg = Optimisers.apply!(Optimisers.Adam(), opt_adam.state, x, g)[2]
Base.Broadcast.Broadcasted{Metal.MtlArrayStyle{1}}(*, (Base.Broadcast.Broadcasted(/, (Base.Broadcast.Broadcasted(/, (Float32[0.17961349, 0.2150245, 0.1685136], Base.Broadcast.Broadcasted(-, (1, 0.8099999785423279)))), Base.Broadcast.Broadcasted(+, (Base.Broadcast.Broadcasted(sqrt, (Base.Broadcast.Broadcasted(/, (Float32[0.0009599409, 0.0013757595, 0.0008449606], Base.Broadcast.Broadcasted(-, (1, 0.9980010128617287)))),)), 1.0f-8)))), 0.001f0))

julia> x .= x .- lazyg
ERROR: InvalidIRError: compiling MethodInstance for (::GPUArrays.var"#broadcast_kernel#26")(::Metal.mtlKernelContext, ::MtlDeviceVector{Float32, 1}, ::Base.Broadcast.Broadcasted{Metal.MtlArrayStyle