Significant compile time latency in Flux with default optimization

[Alexander-Barth]

I am trying to optimize a super-resolution GAN, and I see very significant compile time.
This issue is discussed here:
https://discourse.julialang.org/t/significant-compile-time-latency-in-flux-with-a-gan/68518

@ToucheSir was able to reduce this issue to the code below (thanks a lot!).

I am using Julia 1.6.3 on Linux and these packages:

(@v1.6) pkg> st Zygote ZygoteRules Flux
      Status `~/.julia/environments/v1.6/Project.toml`
  [587475ba] Flux v0.12.8 `https://github.com/FluxML/Flux.jl.git#master`
  [e88e6eb3] Zygote v0.6.30
  [700de1a5] ZygoteRules v0.2.2

It appears that this issue is sensitive to the optimization level and to the julia version used.

with julia 1.6.3 (default optimization):

 2.129568 seconds (6.12 M allocations: 352.760 MiB, 5.21% gc time, 99.96% compilation time)
  0.000323 seconds (1.80 k allocations: 138.938 KiB)
2740.221683 seconds (62.97 M allocations: 3.596 GiB, 0.05% gc time, 0.44% compilation time)
  0.004320 seconds (12.04 k allocations: 2.761 MiB)

Note the timing of the first call the loss_grad(lr_images, ps).

In julia 1.5.3 I get the following timing:

  1.969913 seconds (8.92 M allocations: 458.885 MiB, 6.31% gc time)
  0.000265 seconds (1.30 k allocations: 102.344 KiB)
 25.426727 seconds (57.88 M allocations: 2.927 GiB, 3.93% gc time)
  0.005641 seconds (14.71 k allocations: 1.047 MiB)

with "julia 1.6.3 -O1":

 1.389465 seconds (6.12 M allocations: 352.764 MiB, 25.61% gc time, 99.94% compilation time)
  0.000403 seconds (2.10 k allocations: 143.625 KiB)
 28.301099 seconds (62.99 M allocations: 3.596 GiB, 3.96% gc time, 26.40% compilation time)
  0.004889 seconds (12.94 k allocations: 2.775 MiB)

using Flux

channels = 4

function resblock(channels)
    return SkipConnection(Chain(
        Conv((3, 3), channels => channels, pad=1),
        Conv((3, 3), channels => channels, pad=1),
    ), +)
end

model = Chain(
    SkipConnection(
        Chain(
            resblock(channels),
            resblock(channels),
            resblock(channels),
            resblock(channels),
            resblock(channels),

            resblock(channels),
            resblock(channels),
            resblock(channels),
            resblock(channels),
            resblock(channels),

            resblock(channels),
            resblock(channels),
            resblock(channels),
            resblock(channels),
            resblock(channels),
        ),
    +),
    AdaptiveMeanPool((1, 1))
)

display(model)
println()
@show typeof(model) 

loss(x) = sum(model(x))

lr_images = randn(Float32, 2, 2, channels, 1)
@time loss(lr_images)
@time loss(lr_images) 
    
loss_grad(x, ps) = gradient(() -> loss(x), ps)
    
ps = Flux.params(model)
@time loss_grad(lr_images, ps)
@time loss_grad(lr_images, ps)

This issue might be related
#1119

[mcabbott]

With Julia master I get this:

julia> @time loss_grad(lr_images, ps);
 63.622937 seconds (67.18 M allocations: 3.529 GiB, 1.54% gc time, 99.91% compilation time)

julia> @time loss_grad(lr_images, ps);
  0.030501 seconds (13.07 k allocations: 2.931 MiB)

julia> @btime loss_grad($lr_images, $ps);
  min 33.048 ms, mean 33.933 ms (12887 allocations, 2.91 MiB. GC mean 0.94%)

julia> model.layers[1].layers.layers |> length  # 40s compile time with 10 layers, no huge jump
15

julia> VERSION
v"1.8.0-DEV.1063"

If I replace the long tuple with a vector, very crudely:

struct VChain
  layers
  VChain(xs...) = new(collect(xs))
end
Flux.@forward VChain.layers Base.getindex, Base.length, Base.first, Base.last,
  Base.iterate, Base.lastindex, Base.keys

Flux.functor(::Type{<:VChain}, c) = c.layers, ls -> VChain(ls...)

(c::VChain)(x) = applychain(c.layers, x)
applychain(fs, x) = isempty(fs) ? x : applychain(fs[2:end], first(fs)(x))

... then I get better times. Which is perhaps a similar story to the -O1 times above.

julia> @time loss_grad(lr_images, ps);
 14.622136 seconds (66.20 M allocations: 3.480 GiB, 5.22% gc time, 99.71% compilation time)

julia> @time loss_grad(lr_images, ps);
  0.032955 seconds (13.47 k allocations: 1.664 MiB)

julia> @btime loss_grad($lr_images, $ps);
  min 36.019 ms, mean 36.910 ms (13381 allocations, 1.66 MiB. GC mean 0.39%)

[ToucheSir]

Interesting that it allocates less as well. I wonder if switching from applychain to foldl now that the latter has been optimized would help. Another thing to test would be @nograding DenseConvDims and ConvDims. I tried ChainRulesCore.@non_differentiable with no luck, but there's a good chance that was user error. Trying to create a MWE without Flux or NNlib is tough, because it basically requires copying over definitions from both wholesale.

[mcabbott]

Oh right, foldl is a good idea. Instead of applychain's tuple-peeling. Quick attempt:

using Flux
struct FChain{T}
  layers::T
  FChain(xs...) = new{typeof(xs)}(xs)
  FChain(xs::AbstractVector) = new{typeof(xs)}(xs)  # allows FChain(Any[Layer, Layer, ...]), type-unstable
end
Flux.@forward FChain.layers Base.getindex, Base.length, Base.first, Base.last,
  Base.iterate, Base.lastindex, Base.keys

Flux.functor(::Type{<:FChain}, c) = c.layers, ls -> FChain(ls...)
Flux.functor(::Type{<:FChain{<:AbstractVector}}, c) = c.layers, ls -> FChain(ls)  # maybe?

# (c::FChain)(x) = foldl((y,f) -> f(y), c.layers; init=x) # NO, this forgets the gradient for x
(c::FChain)(x) = foldl((y,f) -> f(y), (x, c.layers...))
(c::FChain{<:AbstractVector})(x) = foldl((y,f) -> f(y), vcat([x], c.layers)))

julia> @time loss_grad(lr_images, ps);
 23.586338 seconds (68.82 M allocations: 3.611 GiB, 3.41% gc time, 99.85% compilation time)

julia> @time loss_grad(lr_images, ps);
  0.026522 seconds (13.36 k allocations: 2.719 MiB)

julia> model.layers[1].layers.layers |> length
15

julia> @btime loss_grad($lr_images, $ps);
  min 32.178 ms, mean 32.809 ms (13357 allocations, 2.72 MiB. GC mean 0.75%)

That's on a cold start. By accident, running this example after some unrelated checks (m = FChain(Dense(2,3), Dense(3,4)); g = gradient(() -> sum(m([1,2])), params(m)); g[m.layers[1].weight]) I got a much faster time:

 13.804523 seconds (23.11 M allocations: 1.178 GiB, 1.70% gc time, 99.72% compilation time)

So perhaps 10s of this is generic startup (and perhaps that could be helped by precompilation?) and not the effects which scale. Comparable warm start for the original Chain:

 52.581168 seconds (21.92 M allocations: 1.122 GiB, 0.59% gc time, 99.88% compilation time)

[ToucheSir]

There are precompilation statements for explicit params, but not for the version of pullback that uses implicit params: https://github.com/FluxML/Zygote.jl/blob/master/src/precompile.jl. If/how much of a difference this would make I'm not sure.

[DhairyaLGandhi]

Most of the code is shared between the two modes, and the cached precompilation would be cached anyway so I don't know how much of a difference it would make.

[Alexander-Barth]

Thank you all for insights. I am impressed how easy it is to implemented an alternative Chain structure. The proposed foldl approach works also well in my original SRGAN implementation.

[mcabbott]

The foldl approach might be worth adopting in general. My only reservation is that I'm not certain it will handle 2nd derivatives, it's meant to, but this isn't tested.

A bigger change would be to allow, or even automatically switch to, a deliberately type-unstable approach. This could be nicer (and prob. faster) than my crude VChain above. But some people want very small Chains to be fast; an automatic switch would probably need some heuristic to decide when (by counting sizes on construction?) and sounds like a big project.

[ToucheSir]

@Alexander-Barth if you have them, I'd love to see how the times look now on 1.6/1.7.

If we can think of some easy tests to ensure foldl is 2nd derivative safe, I'd vote for switching to it asap. For the type stability tradeoff, we might be able to get away with simply allowing Chain.layers to be non-tuple container types. Then folks not terribly worried about small chain performance could opt to pass a Vector{Any} instead.

[mcabbott]

True, allowing Chain([Layer, Layer, Layer]) might not be so strange. I believe foldl should handle that too, have added a method to FChain above. Although it would be worth running the race against VChain.

[Alexander-Barth]

@ToucheSir Of course, here are the times for FChain (with a simplified SRGAN model)

Julia 1.7.0-rc3 (first and 2nd call of pullback)

 33.978640 seconds (76.75 M allocations: 3.952 GiB, 4.16% gc time, 99.75% compilation time)
  0.655267 seconds (49.21 k allocations: 3.040 MiB, 98.58% compilation time)

Julia 1.6.3:

 32.954155 seconds (71.80 M allocations: 4.020 GiB, 4.61% gc time, 66.53% compilation time)
  0.579941 seconds (51.56 k allocations: 3.501 MiB, 98.33% compilation time)

So the timing are very similar. Surprisingly `VChain' is a bit faster:

Julia 1.7.0-rc3 (first and 2nd call of pullback)

24.264638 seconds (64.14 M allocations: 3.350 GiB, 5.62% gc time, 99.71% compilation time)                                                                           
  0.041324 seconds (48.69 k allocations: 2.564 MiB, 78.66% compilation time)       

Julia 1.6.3:

24.329194 seconds (61.18 M allocations: 3.471 GiB, 5.11% gc time, 64.15% compilation time)
  0.042680 seconds (51.36 k allocations: 3.016 MiB, 78.65% compilation time)

[mcabbott]

Thanks for checking. Interesting that the 2nd run with FChain is 10x slower than with VChain. Do they settle down (or @btime) to similar numbers?

[Alexander-Barth]

Indeed with @btime, the runtime of both approaches are very similar:

	julia 1.6.3	julia 1.7.0-rc3
FChain	7.922 ms (8874 allocations: 1.13 MiB)	7.736 ms (8571 allocations: 1009.80 KiB)
VChain	7.931 ms (8675 allocations: 687.95 KiB)	7.696 ms (8023 allocations: 548.80 KiB)

[ToucheSir]

Is foldl with init not working a Zygote or ChainRules thing? I imagine getting that ironed out would result in a far more competitive FChain.

[mcabbott]

I think it's both, really -- by convention neither package believes keyword arguments should have gradients. This is perhaps the most convincing example I've seen where the gradient would be desired. Cc @oxinabox for design thoughts?

I edited the code above to make it the first element foldl((y,f) -> f(y), (x, c.layers...)). I believe that this should be approximately free when c.layers isa Tuple but haven't checked.

[oxinabox]

In theory, I think we can write rules for gradients w.r.t keyword arguments by writing rules for:

rrule(::RuleConfig, ::Core.kwftype(foldl), kwargs::NamedTuple, ::typeof(foldl), op, iter)

But I have never tested that.
It may be that the kwarg handling code that is there now doesn't allow that to work because it redirects to

rrule(::RuleConfig, ::typeof(foldl), op, iter; kwargs...)

I think it might work if the version that takes kwargs without gradients isn't defined though.