EnzymeAD / EnzymeAD/Enzyme.jl

Benchmarking some very simple Flux models

Open
#2,069 6 comments 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
586
Forks
108
Avg merge
1d 5h
Merged PRs (30d)
44

Description

On some extremely simple Flux models, Enzyme seems to be slower than Zygote for me. What's going wrong here?
```julia
julia> using Flux, Enzyme, Test, BenchmarkTools

julia> mlp = Chain(Flux.flatten, Dense(28^2 => 32, tanh), Dense(32 => 10));

julia> img = rand32(28, 28, 1, 128);

julia> @inferred mlp(img); # type-stable

julia> Flux.gradient((m,x) -> sum(abs2, m(x)), mlp, img)[1].layers[2].bias[1:3]
3-element Vector{Float32}:
-15.980308
6.2900686
-79.44746

julia> Enzyme.gradient(Reverse, (m,x) -> sum(abs2, m(x)), mlp, img)[1].layers[2].bias[1:3]
3-element Vector{Float32}:
-15.980312
6.2900686
-79.44745

julia> @btime $mlp($img);
min 10.958 μs, mean 14.119 μs (6 allocations, 43.09 KiB)

julia> @btime Flux.gradient((m,x) -> sum(abs2, m(x)), $mlp, $img);
min 38.250 μs, mean 67.356 μs (86 allocations, 596.27 KiB)

julia> @btime Enzyme.gradient(Reverse, (m,x) -> sum(abs2, m(x)), $mlp, $img);
min 75.125 μs, mean 119.919 μs (55 allocations, 579.61 KiB)

# a slightly bigger model

julia> lenet = Chain( # from the model zoo
Conv((5, 5), 1=>6, relu),
MaxPool((2, 2)),
Conv((5, 5), 6=>16, relu),
MaxPool((2, 2)),
Flux.flatten,
Dense(256 => 120, relu),
Dense(120 => 84, relu),
Dense(84 => 10),
);

julia> @inferred lenet(img); # type-stable

julia> Flux.gradient((m,x) -> sum(abs2, m(x)), lenet, img)[1].layers[1].bias
6-element Vector{Float32}:
10.119315
0.0
...

julia> Enzyme.gradient(Reverse, (m,x) -> sum(abs2, m(x)), lenet, img)[1].layers[1].bias
6-element Vector{Float32}:
10.119322
0.0
...

julia> @btime $lenet($img);
min 655.583 μs, mean 1.107 ms (160 allocations, 5.60 MiB)

julia> @btime Flux.gradient((m,x) -> sum(abs2, m(x)), $lenet, $img);
min 4.979 ms, mean 6.300 ms (558 allocations, 14.18 MiB)

julia> @btime Enzyme.gradient(Reverse, (m,x) -> sum(abs2, m(x)), $lenet, $img);
min 8.347 ms, mean 9.752 ms (538 allocations, 15.42 MiB)

# tweak Enzyme to see if details matter...

julia> tmp_loss(m,x) = sum(abs2, m(x)); # give it a name

julia> @btime Enzyme.gradient(Reverse, tmp_loss, $lenet, $img);
min 8.260 ms, mean 9.766 ms (538 allocations, 15.42 MiB)

julia> @btime Enzyme.gradient(Reverse, tmp_loss, $lenet, Const($img));
min 8.030 ms, mean 9.235 ms (479 allocations, 14.75 MiB)

julia> @btime Enzyme.autodiff(Reverse, tmp_loss, Active, $(Duplicated(lenet, deepcopy(lenet))), Const($img));
min 7.642 ms, mean 8.638 ms (359 allocations, 14.57 MiB)
```
Versions:
```julia
(jl_w98UzC) pkg> st Enzyme
Status `/private/var/folders/yq/4p2zwd614y59gszh7y9ypyhh0000gn/T/jl_w98UzC/Project.toml`
[7da242da] Enzyme v0.13.14

julia> versioninfo()
Julia Version 1.10.4
Commit 48d4fd48430 (2024-06-04 10:41 UTC)
Build Info:
Official https://julialang.org/ release
Platform Info:
OS: macOS (arm64-apple-darwin22.4.0)
CPU: 11 × Apple M3 Pro
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-15.0.7 (ORCJIT, apple-m1)
Threads: 4 default, 0 interactive, 2 GC (on 5 virtual cores)
Environment:
JULIA_NUM_THREADS = 4
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.