# Flux LayerNorm slower than pytorch?

**URL:** <https://discourse.julialang.org/t/flux-layernorm-slower-than-pytorch/78084>\
**Category:** Machine Learning\
**Tags:** question, flux, pytorch\
**Created:** [March 18, 2022, 1:26pm UTC](https://discourse.julialang.org/t/flux-layernorm-slower-than-pytorch/78084 "2022-03-18T13:26:38Z")\
**Posts on this page:** 11\
**Page:** 1

<div class="post-metadata">

**Author:** ![gpucce](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gpucce/32/21995_2.png) [@gpucce](https://discourse.julialang.org/u/gpucce)\
**Post date:** [March 18, 2022, 1:26pm UTC](https://discourse.julialang.org/t/flux-layernorm-slower-than-pytorch/78084/1 "2022-03-18T13:26:38Z")

</div>

I closed a similar topic I opened about one hour ago by mistake, here I try again with clearer example, the issue is that the same LayerNorm layer in pytorch and Flux has large difference in performance and I don’t know if it is expected or not.

If anybody could clarify this or tell me where to look for an answer I would be grateful here are two minimal working examples and the resulting benchmarks (both done in Pluto in case it affects the results):

In pytorch:

```julia
using PythonCall

torch = pyimport("torch")
np = pyimport("numpy")

l = torch.nn.LayerNorm(768)
v = torch.tensor(
    np.array(
        rand(Float32, 2, 128, 768)
    )
)
	
@benchmark l(v) 

```

BenchmarkTools.Trial: 10000 samples with 1 evaluation.  
Range (min … max): 97.483 μs … 3.871 ms ┊ GC (min … max): 0.00% … 0.00%  
Time (median): 135.234 μs ┊ GC (median): 0.00%  
Time (mean ± σ): 138.385 μs ± 42.160 μs ┊ GC (mean ± σ): 0.00% ± 0.00%

Memory estimate: 16 bytes, allocs estimate: 1.

In Flux:

```julia
using Flux

jl = Flux.LayerNorm(768)
jv = rand(Float32, 768, 128, 2)

@benchmark jl(jv)

```

BenchmarkTools.Trial: 10000 samples with 1 evaluation.  
Range (min … max): 286.678 μs … 6.144 ms ┊ GC (min … max): 0.00% … 91.80%  
Time (median): 322.035 μs ┊ GC (median): 0.00%  
Time (mean ± σ): 341.300 μs ± 324.158 μs ┊ GC (mean ± σ): 5.36% ± 5.33%

Memory estimate: 2.25 MiB, allocs estimate: 24.

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [March 18, 2022, 3:14pm UTC](https://discourse.julialang.org/t/flux-layernorm-slower-than-pytorch/78084/2 "2022-03-18T15:14:12Z")

</div>

Flux’s LayerNorm just calls [https://github.com/FluxML/Flux.jl/blob/v0.12.9/src/layers/stateless.jl#L36](https://github.com/FluxML/Flux.jl/blob/v0.12.9/src/layers/stateless.jl#L36), so if you have any ideas for optimizing it please do file a PR 🙂

---

<div class="post-metadata">

**Author:** ![gpucce](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gpucce/32/21995_2.png) [@gpucce](https://discourse.julialang.org/u/gpucce)\
**Post date:** [March 18, 2022, 3:22pm UTC](https://discourse.julialang.org/t/flux-layernorm-slower-than-pytorch/78084/3 "2022-03-18T15:22:43Z")

</div>

Ok thanks! I didn’t go throught the code because I thought in no way I could understand it but I gues I could have. Then I gues that’s just how things are!!

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [March 18, 2022, 3:23pm UTC](https://discourse.julialang.org/t/flux-layernorm-slower-than-pytorch/78084/4 "2022-03-18T15:23:23Z")

</div>

As the comment there says, `σ = std(x, dims=dims, corrected=false)` could use the mean, since [this rule](https://github.com/FluxML/Zygote.jl/blob/master/src/lib/array.jl#L375-L378) accepts it (as does [this one](https://github.com/JuliaDiff/ChainRules.jl/blob/main/src/rulesets/Statistics/statistics.jl#L42-L56).)

---

<div class="post-metadata">

**Author:** ![gpucce](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gpucce/32/21995_2.png) [@gpucce](https://discourse.julialang.org/u/gpucce)\
**Post date:** [March 18, 2022, 3:35pm UTC](https://discourse.julialang.org/t/flux-layernorm-slower-than-pytorch/78084/5 "2022-03-18T15:35:45Z")

</div>

Out of curiosity I tried but at least just replacing the Flux example above with

```julia
function normalise(x::AbstractArray; dims=ndims(x), ϵ=Flux.ofeltype(x, 1e-5))
    μ = mean(x, dims=dims)
    σ = std(x, dims=dims, mean=μ, corrected=false) 
    return (x .- μ) ./ (σ .+ ϵ)
end

# jl = Flux.LayerNorm(768)
jv = rand(Float32, 768, 128, 2)

@benchmark normalise(jv)

```

Gives very similar results, in fact a bit worse, would it make a difference if I integrate the different normalize within a struct using `@functor`?

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [March 18, 2022, 3:44pm UTC](https://discourse.julialang.org/t/flux-layernorm-slower-than-pytorch/78084/6 "2022-03-18T15:44:06Z")

</div>

No, I don’t think so. But the right comparison for the `mean` change is probably this, a small improvement:

```julia
julia> @btime normalise($jv; dims=1:1);
  min 240.000 μs, mean 322.445 μs (14 allocations, 770.50 KiB)

julia> @btime Flux.normalise($jv; dims=1:1);
  min 265.708 μs, mean 358.654 μs (20 allocations, 771.72 KiB)

```

Note that there are other steps in applying the layer, `@less jl(jv)`. There might be room to optimise these further.

---

<div class="post-metadata">

**Author:** ![gpucce](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gpucce/32/21995_2.png) [@gpucce](https://discourse.julialang.org/u/gpucce)\
**Post date:** [March 18, 2022, 3:50pm UTC](https://discourse.julialang.org/t/flux-layernorm-slower-than-pytorch/78084/7 "2022-03-18T15:50:31Z")

</div>

Which package is @less from? Sorry I found it

---

<div class="post-metadata">

**Author:** ![gpucce](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gpucce/32/21995_2.png) [@gpucce](https://discourse.julialang.org/u/gpucce)\
**Post date:** [March 18, 2022, 3:54pm UTC](https://discourse.julialang.org/t/flux-layernorm-slower-than-pytorch/78084/8 "2022-03-18T15:54:49Z")

</div>

I did find `less` but I don’t understand what you meant indicating it, like is there smth in the output of @less that you think would be useful?

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [March 18, 2022, 8:09pm UTC](https://discourse.julialang.org/t/flux-layernorm-slower-than-pytorch/78084/9 "2022-03-18T20:09:53Z")

</div>

`@less` just shows you the source code and location of the method being called. In this case, [https://github.com/FluxML/Flux.jl/blob/master/src/layers/normalise.jl#L175-L178](https://github.com/FluxML/Flux.jl/blob/master/src/layers/normalise.jl#L175-L178). I think Michael’s point is that more is going on than just `normalise`, as evidenced by [https://github.com/FluxML/Flux.jl/blob/master/src/layers/normalise.jl#L177](https://github.com/FluxML/Flux.jl/blob/master/src/layers/normalise.jl#L177). Given you mentioned above that just `normalise` is slower than the entire PyTorch LayerNorm forward pass, that function is still very much a culprit.

---

<div class="post-metadata">

**Author:** ![gpucce](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gpucce/32/21995_2.png) [@gpucce](https://discourse.julialang.org/u/gpucce)\
**Post date:** [March 18, 2022, 8:29pm UTC](https://discourse.julialang.org/t/flux-layernorm-slower-than-pytorch/78084/10 "2022-03-18T20:29:29Z")

</div>

Thanks so much for the explanation, I would have never understood the point. I will look into the code more closely since it is so straightforward to read, though without fully understanding all that is going on I guess.

If I can ask you what would you say is a good starting point in the code to understand the flux design?

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [March 18, 2022, 10:46pm UTC](https://discourse.julialang.org/t/flux-layernorm-slower-than-pytorch/78084/11 "2022-03-18T22:46:38Z")

</div>

Do you mean in general or for this specific layer? For LayerNorm specifically, the implementation is not that different from what you’d write by hand using Base Julia or e.g. Numpy. `.diag` is extracted into its own [layer type](https://github.com/FluxML/Flux.jl/blob/master/src/layers/basic.jl#L174) because that operation is useful outside of layer normalization. Ideally we’d like to avoid having to write a fully custom kernel like [https://github.com/pytorch/pytorch/blob/ef066f0832eab3192a0610a1ecf955239c26de0b/aten/src/ATen/native/cpu/layer\_norm\_kernel.cpp](https://github.com/pytorch/pytorch/blob/ef066f0832eab3192a0610a1ecf955239c26de0b/aten/src/ATen/native/cpu/layer_norm_kernel.cpp), but if that’s not possible then such a kernel would be a welcome addition to [https://github.com/FluxML/NNlib.jl](https://github.com/FluxML/NNlib.jl).