What's the state of Automatic Differentiation in Julia January 2023?

General Usage
21

On my M1 and Julia 1.9 rc3 with Enzyme#main, I am getting a difference between Enzyme and Zygote gradients of

2-element Vector{Float64}:
 0.15407736749039058
 0.22841827115822566

by running the same code above.

Edit: also ] test Enzyme causes a segfault. Is the M1 not supported by Enzyme?

22

Like I said in the other thread (Use Enzyme in flux - #7 by wsmoses) the code Chris provided is not a good way to call it because of the type unstable closure capture of dudt, p, and st. These should be passed explicitly as Const parameters for performance and other reasons.

M1 indeed is not officially supported, partially because we don’t have a dev machine or CI for it. If you can isolate and post which test fails on GitHub we can try to look into it.

23

I can concur that it’s giving something odd on my M2.

24

Thanks William. I tried Windows and got the exact same difference in gradients. So it’s not an M1 thing. Enzyme tests pass on Windows. I suppose that example is not supposed to work at all?

25

On my Intel CPU, the example just throws errors.

julia> Enzyme.autodiff(Reverse, f, Duplicated(x, bx), Duplicated(y, ones(2)))
warning: didn't implement memmove, using memcpy as fallback which can result in errors
ERROR: Enzyme execution failed.
Enzyme: not yet implemented in reverse mode, jl_getfield
Stacktrace:
 [1] getindex
   @ ./tuple.jl:29
 [2] getindex
   @ ./tuple.jl:0

Stacktrace:
 [1] throwerr(cstr::Cstring)
   @ Enzyme.Compiler ~/.julia/packages/Enzyme/YBQJk/src/compiler.jl:2536
 [2] macro expansion
   @ ~/.julia/packages/Enzyme/YBQJk/src/compiler.jl:8646 [inlined]
 [3] enzyme_call
   @ ~/.julia/packages/Enzyme/YBQJk/src/compiler.jl:8338 [inlined]
 [4] CombinedAdjointThunk
   @ ~/.julia/packages/Enzyme/YBQJk/src/compiler.jl:8301 [inlined]
 [5] autodiff
   @ ~/.julia/packages/Enzyme/YBQJk/src/Enzyme.jl:205 [inlined]
 [6] autodiff
   @ ~/.julia/packages/Enzyme/YBQJk/src/Enzyme.jl:228 [inlined]
 [7] autodiff(::EnzymeCore.ReverseMode{false}, ::typeof(f), ::Duplicated{Vector{Float64}}, ::Duplicated{Vector{Float64}})
   @ Enzyme ~/.julia/packages/Enzyme/YBQJk/src/Enzyme.jl:214
 [8] top-level scope
   @ REPL[12]:1

This is common across Julia 1.8.5, 1.9-rc3, and master.
This was with Enzyme 0.11.1. I also tried 0.11.0 on 1.9-rc3.

26

Well @ChrisRackauckas did say he was using the unreleased main branch of Enzyme, not any of the registered versions.

27

:man_facepalming:
Now I get

julia> @show bx - bx2
bx - bx2 = [0.1540773674903897, 0.22841827115822566]
2-element Vector{Float64}:
 0.1540773674903897
 0.22841827115822566

julia> versioninfo()
Julia Version 1.9.0-rc3
Commit 1853b90328 (2023-04-26 15:51 UTC)
Platform Info:
  OS: Linux (x86_64-generic-linux)
  CPU: 28 × Intel(R) Core(TM) i9-9940X CPU @ 3.30GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-14.0.6 (ORCJIT, skylake-avx512)
28

@ChrisRackauckas we seem to all be getting the exact same non-zero difference on MacOS, Windows and Linux. Which OS did you use to run the above example :sweat_smile:?

29

For CI I’m happy to help you set up Cirrus on M1, if that helps

2 Likes
30

Peculiar. It was on Windows, and now I cannot recreate it. I do not know the manifest of the time, though I do know that it was in a clean REPL. It is quite mysterious to me :sweat_smile:. Maybe a difference in Julia version?

31

Probably not. I tested on Julia 1.9 rc3. Chris Elrod tests 3 Julia versions. It’s most likely dependencies but I am guessing :person_shrugging:

Anyways, I think Enzyme will be well served by a simpler AbstractDifferentiation API for non-mutating functions. Probably worth giving it another go soon-ish.

1 Like
32

Again, I think Enzyme is really promising, blazing fast when it works, and probably the best way forward. I really appreciate all the great work developers have put into it. It is just that, at the moment, it is not stable enough for those who are not prepared to invest extra effort in checking and debugging (again, for nontrivial examples).

For my particular issue, yes, the workaround works. But I ended up rewriting the function in a C-style approach: instead of higher order functions, just loop over stuff and increment a float. This works fine.

My other meta-gripe with Enzyme is that it is tricky to understand the error messages. Whenever something goes wrong, hundreds of lines of IR are dumped on the user. Again, I understand that mapping that back to Julia code is tricky, but in this form it is difficult to isolate errors and produce MWEs.

33

I think the global capture issue is causing wrong gradients though so my initial tests didn’t pass.

Edit: no segfaults though which is an improvement from last time I played with Enzyme for AbstractDifferentiation

3 Likes
34

I’ve landed some code that will automatically throw an error, saying to enable runtimeActivity if it detects a scenario where it is needed. Hopefully that should at least partially abate issues (all the activity bugs on the Enzyme.jl tracker hit it).

The longer error messages involve a combination of issues. The biggest one at the moment is an issue on Julia proper that we’ve been waiting for someone to look at since March =/ (Frame registration for foreign JITs · Issue #49056 · JuliaLang/julia · GitHub).

@ChrisRackauckas am not sure, but our unreleased main branch is an unreleased main branch for a reason. It usually has a lot more fixes, but it also means we can break it (like I did over the weekend to update the ABI, which awaited a jll bump). Not sure if that’s what you ran into (I never tested your code myself).

6 Likes
35

:man_facepalming: I had the bug report code in the same script, so it had runtime activity checking turned on from that part. The following is what works:

using Enzyme
Enzyme.API.runtimeActivity!(true)

x  = [2.0, 2.0]
bx = [0.0, 0.0]
y  = [0.0,0.0]

using ComponentArrays, Lux, Random

rng = Random.default_rng()
Random.seed!(rng,100)
dudt2 = Lux.Chain(x -> x.^3,
                  Lux.Dense(2, 50, tanh),
                  Lux.Dense(50, 2))
p, st = Lux.setup(rng, dudt2)

function f(x::Array{Float64}, y::Array{Float64})
    y .= dudt2(x, p, st)[1]
    return nothing
end

Enzyme.autodiff(Reverse, f, Duplicated(x, bx), Duplicated(y, ones(2)))

function f2(x::Array{Float64})
    dudt2(x, p, st)[1]
end

using Zygote
bx2 = Zygote.pullback(f2, x)[2](ones(2))[1]

@show bx - bx2

The difference is Enzyme.API.runtimeActivity!(true), so yeah @wsmoses chalk this one up to another runtimeActivity check.

I talked with Jeff and Gabriel about this today. We were talking about a v1.10 milestone.

1 Like
36

I’ve fixed main now to no longer require runtimeActivity in this case (and in most other cases throw an error when it may be needed).

5 Likes
37

In the past few days I have been trying to make Enzyme work on a medium-sized log posterior calculation. Ran into what turned out to be 2 (or 3?) separate issues, and got a lot of help from @wsmoses with isolating them as an MWE, at which point they were promptly fixed.

Just want to say thank you! And recommend Enzyme for everyone who is looking for a fast reverse-mode AD solution. Understandably the library is maturing and there are bugs, but they can be tracked down and the support from the authors is outstanding.

23 Likes
38

:+1: that’s what I was trying to convey. Is it perfect right now? No. But it’s definitely on the right path in terms of supporting the full language, and it’s moving pretty fast. Check your answers of course since it’s still early stages, but with v0.11 it’s really looking on track to where it needs to be.

11 Likes