wsmoses
commented
3 months ago Can you attach The whole error log
On Thu, Apr 18, 2024 at 6:57 PM ExpandingMan @.***> wrote:
This issue is to provide a minimal example of neural network training with Lux to hopefully make it easier for developers to work toward making it viable. It probably isn't news to anyone here that this example fails, but it was indicated to me on slack that it would still be helpful to have this issue for reference.
In this example, we train a neural network with zero hidden layers to approximate the polynomial $x^2 - 2x$. It is trivial to generalize this to deeper neural networks but probably not useful for this demonstration. This roughly follows the Lux tutorial here https://lux.csail.mit.edu/dev/tutorials/beginner/2_PolynomialFitting, but I have stripped out the opaque Lux training stuff so that it's more clear what's going on. I expect this example should be simpler to diagnose than the equivalent with Flux, as the explicit parameterization of Lux makes it easier to reason about, but I also expect that if this example were working the analogous Flux example surely would. Indeed, I think this example is a good proxy for a huge number of common use cases.
using LinearAlgebra, Random, Statistics, Optimisersusing CUDAusing Lux, LuxCUDAimport Zygote, Enzyme
confirmed this works for enzyme on gpuconst dev = gpu_device()
function makedata(rng::AbstractRNG) X = reshape(collect(range(-2.0f0, 2.0f0, 128)), (1, 128)) y = evalpoly.(X, ((0, -2, 1),)) .+ randn(rng, Float32, (1, 128)) .* 0.1f0 (X, y)end function loss(model, θ, ψ, (X, y)) (ŷ, ψ) = Lux.apply(model, X, θ, ψ) mean(abs2, ŷ .- y)end function gradloss_zygote(model, θ, ψ, (X, y)) (∇ℓ,) = Zygote.gradient(θ) do ϑ loss(model, ϑ, ψ, (X, y)) end ∇ℓend function gradloss_enzyme(model, θ, ψ, (X, y)) ℓ = ϑ -> begin loss(model, ϑ, ψ, (X, y)) end Enzyme.gradient(Enzyme.Reverse, ℓ, θ)end function main(rng=Random.Xoshiro(999), model=Chain(Dense(1=>16, gelu), Dense(16=>1)), (X, y)=makedata(rng) |> dev; nepochs=300, ) (θ, ψ) = Lux.setup(rng, model) |> dev
opts = Optimisers.setup(Adam(0.01f0), θ) for j ∈ 1:nepochs ∇ℓ = gradloss_enzyme(model, θ, ψ, (X, y)) (opts, θ) = Optimisers.update!(opts, θ, ∇ℓ) end (ŷ, _) = Lux.apply(model, X, θ, ψ) (y, ŷ)endNote that
- This works with both Zygote and Enzyme if dev = cpu_device (i.e. no GPU is involved at all).
- This works using gradloss_zygote using either cpu_device or gpu_device.
- This fails rather spectacularly using gradloss_enzyme and gpu_device.
The error output is so verbose that I won't try to reproduce it all here (it goes nuts and starts dumping LLVM IR), I expect others to be able to reproduce a same or similar error, but the stack trace is
Stacktrace: [1] julia_error(cstr::Cstring, val::Ptr{…}, errtype::Enzyme.API.ErrorType, data::Ptr{…}, data2::Ptr{…}, B::Ptr{…}) @ Enzyme.Compiler ~/.julia/packages/Enzyme/MIIMf/src/compiler.jl:1684 [2] EnzymeCreatePrimalAndGradient(logic::Enzyme.Logic, todiff::LLVM.Function, retType::Enzyme.API.CDIFFE_TYPE, constant_args::Vector{…}, TA::Enzyme.TypeAnalysis, returnValue::Bool, dretUsed::Bool, mode::Enzyme.API.CDerivativeMode, width::Int64, additionalArg::Ptr{…}, forceAnonymousTape::Bool, typeInfo::Enzyme.FnTypeInfo, uncacheable_args::Vector{…}, augmented::Ptr{…}, atomicAdd::Bool) @ Enzyme.API ~/.julia/packages/Enzyme/MIIMf/src/api.jl:154 [3] enzyme!(job::GPUCompiler.CompilerJob{…}, mod::LLVM.Module, primalf::LLVM.Function, TT::Type, mode::Enzyme.API.CDerivativeMode, width::Int64, parallel::Bool, actualRetType::Type, wrap::Bool, modifiedBetween::Tuple{…}, returnPrimal::Bool, expectedTapeType::Type, loweredArgs::Set{…}, boxedArgs::Set{…}) @ Enzyme.Compiler ~/.julia/packages/Enzyme/MIIMf/src/compiler.jl:3109 [4] codegen(output::Symbol, job::GPUCompiler.CompilerJob{…}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, toplevel::Bool, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing) @ Enzyme.Compiler ~/.julia/packages/Enzyme/MIIMf/src/compiler.jl:4964 [5] codegen @ ~/.julia/packages/Enzyme/MIIMf/src/compiler.jl:4391 [inlined] [6] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, postopt::Bool) @ Enzyme.Compiler ~/.julia/packages/Enzyme/MIIMf/src/compiler.jl:5646 [7] _thunk @ ~/.julia/packages/Enzyme/MIIMf/src/compiler.jl:5646 [inlined] [8] cached_compilation @ ~/.julia/packages/Enzyme/MIIMf/src/compiler.jl:5680 [inlined] [9] (::Enzyme.Compiler.var"#532#533"{…})(ctx::LLVM.Context) @ Enzyme.Compiler ~/.julia/packages/Enzyme/MIIMf/src/compiler.jl:5746 [10] JuliaContext(f::Enzyme.Compiler.var"#532#533"{…}; @.{}) @ GPUCompiler ~/.julia/packages/GPUCompiler/kqxyC/src/driver.jl:52 [11] JuliaContext(f::Function) @ GPUCompiler ~/.julia/packages/GPUCompiler/kqxyC/src/driver.jl:42 [12] #s1926#531 @ ~/.julia/packages/Enzyme/MIIMf/src/compiler.jl:5698 [inlined] [13] @ Enzyme.Compiler ./none:0 [14] (::Core.GeneratedFunctionStub)(::UInt64, ::LineNumberNode, ::Any, ::Vararg{Any}) @ Core ./boot.jl:602 [15] autodiff @ ~/.julia/packages/Enzyme/MIIMf/src/Enzyme.jl:270 [inlined] [16] autodiff @ ~/.julia/packages/Enzyme/MIIMf/src/Enzyme.jl:287 [inlined] [17] gradient @ ~/.julia/packages/Enzyme/MIIMf/src/Enzyme.jl:938 [inlined] [18] gradloss_enzyme(model::Chain{…}, @.{…}, @.{…}, ::Tuple{…}) @ Main ~/src/autodiff/zygote_enzyme_minimal.jl:31 [19] main(rng::Xoshiro, @.{…}, Nothing}, ::Tuple{CuArray{…}, CuArray{…}}; nepochs::Int64) @ Main ~/src/autodiff/zygote_enzyme_minimal.jl:44 [20] main(rng::Xoshiro, @.***{…}, Nothing}, ::Tuple{CuArray{…}, CuArray{…}}) @ Main ~/src/autodiff/zygote_enzyme_minimal.jl:34 [21] top-level scope @ REPL[2]:1 [22] top-level scope @ ~/.julia/packages/CUDA/fGE8R/src/initialization.jl:206
— Reply to this email directly, view it on GitHub https://github.com/EnzymeAD/Enzyme.jl/issues/1392, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAJTUXCTOIDMKDPI3IBBXE3Y6BFVZAVCNFSM6AAAAABGOEXZQSVHI2DSMVQWIX3LMV43ASLTON2WKOZSGI2TCNZSGA2DANI . You are receiving this because you are subscribed to this thread.Message ID: @.***>
ExpandingMan
This issue is to provide a minimal example of neural network training with Lux to hopefully make it easier for developers to work toward making it viable. It probably isn't news to anyone here that this example fails, but it was indicated to me on slack that it would still be helpful to have this issue for reference.
In this example, we train a neural network with zero hidden layers to approximate the polynomial $x^2 - 2x$. It is trivial to generalize this to deeper neural networks but probably not useful for this demonstration. This roughly follows the Lux tutorial here, but I have stripped out the opaque Lux training stuff so that it's more clear what's going on. I expect this example should be simpler to diagnose than the equivalent with Flux, as the explicit parameterization of Lux makes it easier to reason about, but I also expect that if this example were working the analogous Flux example surely would. Indeed, I think this example is a good proxy for a huge number of common use cases.
Note that
dev = cpu_device(i.e. no GPU is involved at all).gradloss_zygoteusing eithercpu_deviceorgpu_device.gradloss_enzymeandgpu_device.The error output is so verbose that I won't try to reproduce it all here (it goes nuts and starts dumping LLVM IR), I expect others to be able to reproduce a same or similar error, but the stack trace is