Benchmarking some very simple Flux models
- 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
Assessment
This issue has not been assessed yet.