search

FluxML / Optimisers.jl

Optimisers.jl defines many standard optimisers and utilities for learning loops.
https://fluxml.ai/Optimisers.jl
MIT License
72 stars 20 forks source link

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

Closed CarloLucibello closed 1 year ago

CarloLucibello commented 1 year ago

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 commented 1 year ago

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 commented 1 year ago

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 commented 1 year ago

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 commented 1 year ago

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:

https://github.com/FluxML/Optimisers.jl/blob/6a4f9480f6b3792f700526e4848d8e07dce93819/src/rules.jl#L209-L213

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