Reduce allocations in Dropout #1791
Conversation
| end | ||
|
|
||
| @adjoint function dropout(x, p; dims=:, active::Bool=true) | ||
| active || return x, Δ -> (Δ, nothing) | ||
| y = dropout_mask(x, p, dims=dims) | ||
| return x .* y, Δ -> (Δ .* y, nothing) | ||
| y = rand!(similar(x, _dropout_shape(x, dims))) |
There was a problem hiding this comment.
Do we need to replace dropout_mask here?
There was a problem hiding this comment.
It would probably be easier to just remove the call to dropout_kernel in dropout_mask and keep the first line here. That also gives a reason for dropout_mask's continued existence (I can see it coming in handy in the future if we ever think of more efficient ways to generate or store the mask depending on input type).
Edit: a slimmed down dropout_mask could also be used by
Flux.jl/src/layers/normalise.jl
Line 114 in 1242c20
| y = dropout_mask(x, p, dims=dims) | ||
| return x .* y, Δ -> (Δ .* y, nothing) | ||
| y = rand!(similar(x, _dropout_shape(x, dims))) | ||
| return x .* _dropout_kernel.(y, p, 1-p), Δ -> (Δ .* _dropout_kernel.(y, p, 1-p), nothing) |
There was a problem hiding this comment.
This makes me wonder if _dropout_kernel should subsume the pointwise mul as well.
There was a problem hiding this comment.
That, I believe, would be equivalent to the change here (but perhaps with neater packaging).
There was a problem hiding this comment.
I believe it would also save a multiplication per element (assuming _dropout_kernel(x, y::T, p, q) where {T} = y > p ? T(x / q) : T(0) or some such)
There was a problem hiding this comment.
I think it'd be equivalent in the end. Maybe check the generated code to verify.
Co-authored-by: Brian Chen <[email protected]>
| @@ -31,7 +31,7 @@ The [`Dropout`](@ref) layer is what you should use in most scenarios. | |||
| function dropout(x, p; dims=:, active::Bool=true) | |||
| active || return x | |||
| y = rand!(similar(x, _dropout_shape(x, dims))) | |||
| @. y = x * _dropout_kernel(y, p, 1-p) | |||
| x .* _dropout_kernel.(y, p, 1-p) | |||
There was a problem hiding this comment.
This allocates a new vector rather than reusing y. I tried this variant and it produced the same lowered code as the original.
There was a problem hiding this comment.
I see, it's a compelling change, but I don't think it works when dims is set, because then the size of y is actually smaller than x. The current code relies on broadcasting to inflate size-1 dims in the mask to the equivalent full dim size in x.
There was a problem hiding this comment.
Something here breaks type-stability, which isn't visible on the benchmarks of large arrays:
x = randn(3,4)
function dropout_pr(x, p; dims=:, active::Bool=true)
active || return x
y = rand!(similar(x, Flux._dropout_shape(x, dims)))
x .* Flux._dropout_kernel.(y, p, 1-p)
end
@btime Flux.dropout($x, 0.5; dims=1) # 60.359 ns
@btime dropout_pr($x, 0.5; dims=1) # 619.457 ns
@code_warntype dropout_pr(x, 0.5; dims=1); # Body::AnyAlso, it is odd that the calculation of 1-p is pulled out of the kernel, but the more expensive 1/q is not. IMO this should be written _dropout_kernel(y, p, invq) = ifelse(y > p, invq, zero(invq)), although in examples I tried the compiler does seem to figure this out. But explicit is better than magic.
There was a problem hiding this comment.
Seems to infer fine for me, perhaps Random.rand! wasn't imported beforehand? I get 100.5ns for dropout_pr and 108.5ns for Flux.dropout.
Riffing of an earlier comment, I wonder if x should also be an arg to dropout_kernel. Local benchmarking didn't show much of a difference, but as long as it doesn't hurt codegen it could help to eliminate some redundancy.
There was a problem hiding this comment.
OK, I can't reproduce this on a restart, not sure what was wrong, sorry.
I doubt it matters much whether you write x .* _dropout_kernel.(y, p, invq) or _dropout_kernel.(x, y, p, invq), but not opposed. Your hope is roughly that it'll compile one broadcast for forwards & back, instead of two?
Pulling out the division and avoiding a branch seems like a good idea, although likewise I can't prove it causes issues.
There was a problem hiding this comment.
Cool. It may not work at all, not sure. It's also possible that this should be y .= rand.(Float32) .> p.
There was a problem hiding this comment.
I get:
julia> cx = cu(randn(100, 100));
julia> @btime CUDA.@sync pr_y($cx, dims=1);
min 12.529 μs, mean 14.025 μs (5 allocations, 160 bytes)
julia> @btime CUDA.@sync bool_y($cx, 0.5, dims=1);
min 14.835 μs, mean 16.531 μs (12 allocations, 640 bytes)
julia> @btime CUDA.@sync bool_y32($cx, 0.5f0, dims=1);
min 14.647 μs, mean 16.188 μs (13 allocations, 656 bytes)
julia> CUDA.@time pr_y(cx, dims=1); # these times very noisy
0.010914 seconds (165 CPU allocations: 8.750 KiB) (1 GPU allocation: 400 bytes, 0.26% memmgmt time)
julia> CUDA.@time bool_y32(cx, 0.5f0, dims=1);
0.006785 seconds (71 CPU allocations: 3.625 KiB) (1 GPU allocation: 100 bytes, 0.39% memmgmt time)
There was a problem hiding this comment.
function pr_y(x, p::Real; dims=:)
y = rand!(similar(x, Flux._dropout_shape(x, dims)))
y .= y .> p
end
julia> gx = cu(x);
julia> CUDA.@sync @btime bool_y($gx, 0.5);
5.213 μs (28 allocations: 1.41 KiB)
julia> CUDA.@sync @btime bool_y($gx, 0.5, dims=1);
5.212 μs (26 allocations: 1.36 KiB)
julia> CUDA.@sync @btime pr_y($gx, 0.5);
9.604 μs (30 allocations: 1.75 KiB)
julia> CUDA.@sync @btime pr_y($gx, 0.5, dims=1);
8.112 μs (26 allocations: 1.64 KiB)But confusingly:
julia> @btime bool_y($x, 0.5, dims=1);
207.288 ns (1 allocation: 144 bytes)
julia> @btime pr_y($x, 0.5, dims=1);
121.337 ns (1 allocation: 496 bytes)There was a problem hiding this comment.
What does CUDA.@sync @btime do? It seems like that would sync once after the benchmark has run, but perhaps it's not like that? I am never sure about GPU benchmarks.
The CPU result is sturprising. Note that your pr_y is different to mine, it makes a second pass over y, and broadcasts back to the same array in-place, so it might hit JuliaLang/julia#43153 . I was assuming that, if you materialise an array of random numbers, you should still fuse the .> p loop with the x one. These should if anything make it slower, though.
In the GPU case, this 2nd pass (over a small array, 50x2?) might mean one more kernel launch, and it's possible this is the entire timing here, 2 vs 1?
There was a problem hiding this comment.
The @sync @btime was just crossed wires on my end, these should be accurate:
julia> @btime CUDA.@sync pr_y($gx, 0.5, dims=1);
13.444 μs (26 allocations: 1.64 KiB)
julia> @btime CUDA.@sync pr_y($gx, 0.5);
14.847 μs (30 allocations: 1.75 KiB)
julia> @btime CUDA.@sync bool_y($gx, 0.5, dims=1);
10.059 μs (26 allocations: 1.36 KiB)
julia> @btime CUDA.@sync bool_y($gx, 0.5);
10.148 μs (28 allocations: 1.41 KiB)
# pr_randonly(x; dims=:) = rand!(similar(x, Flux._dropout_shape(x, dims)))
julia> @btime CUDA.@sync pr_randonly($gx, dims=1);
8.575 μs (4 allocations: 128 bytes)
julia> @btime CUDA.@sync pr_randonly($gx);
8.107 μs (4 allocations: 128 bytes)AFAICT from @device_code_llvm, the GPU path doesn't include a similar aliasing check.
Edit: another factor to consider is that rand! may be quite a bit faster on CPU with 1.7+ because of the new SIMD-friendly Xoshiro implementation.
| y = dropout_mask(x, p, dims=dims) | ||
| return x .* y, Δ -> (Δ .* y, nothing) | ||
| y = rand!(similar(x, _dropout_shape(x, dims))) | ||
| return x .* _dropout_kernel.(y, p, 1-p), Δ -> (Δ .* _dropout_kernel.(y, p, 1-p), nothing) | ||
| end | ||
|
|
||
| function dropout_mask(x, p; dims=:) |
There was a problem hiding this comment.
Note BTW that this re-use of y to save memory suffers from JuliaLang/julia#43153 . Applying the (possible) fix from there saves 30% or so:
julia> x = randn(Float32, 100, 1000);
julia> @btime Flux.dropout_mask($x, 0.5; dims=:);
min 70.791 μs, mean 129.631 μs (7 allocations, 390.80 KiB. GC mean 24.65%)
julia> @eval Base.Broadcast @inline function copyto!(dest::AbstractArray, bc::Broadcasted{Nothing})
axes(dest) == axes(bc) || throwdm(axes(dest), axes(bc))
# Performance optimization: broadcast!(identity, dest, A) is equivalent to copyto!(dest, A) if indices match
if bc.f === identity && bc.args isa Tuple{AbstractArray} # only a single input argument to broadcast!
A = bc.args[1]
if axes(dest) == axes(A)
return copyto!(dest, A)
end
end
bc′ = preprocess(dest, bc)
# Performance may vary depending on whether `@inbounds` is placed outside the
# for loop or not. (cf. https://github.com/JuliaLang/julia/issues/38086)
@simd ivdep for I in eachindex(dest)
@inbounds dest[I] = bc′[I]
end
return dest
end
copyto! (generic function with 126 methods)
julia> @btime Flux.dropout_mask($x, 0.5; dims=:);
min 55.750 μs, mean 102.479 μs (7 allocations, 390.80 KiB. GC mean 24.71%)
That's another reason to avoid this, in favour of the fusion proposed here.
There was a problem hiding this comment.
Also, I think deleting an internal function like this should be fine. If anyone was overloading this for some reason, better they find out sooner than later.
There was a problem hiding this comment.
Would this correctly not trigger for GPU arrays? The type of dest seems pretty broad.
There was a problem hiding this comment.
Not sure this issue exists for GPU, nor whether it calls the method which I pirate here.
There was a problem hiding this comment.
IDK, but the definition appears specific enough to not cause major problems: https://github.com/JuliaGPU/GPUArrays.jl/blob/master/src/host/broadcast.jl#L50
|
I think the biggest outstanding question here is how to handle the case when Another consideration is that the changed code path(s) are generally only active when taking gradients. Which means that most users will only hit the adjoint (which we haven't figured out a more memory-efficient path) and not the main function. One possible direction here is to create a bool mask (saving ~4x the memory). If we can show that doing that and calling the kernel part twice is not slower than what's currently on master, this should be a shoe-in :) |
The current implementation of
Dropoutconstructs an intermediate vector before doing a matrix multiply, but these operations could be fused. This PR explicitly fuses the dropout masking, which results in a reduction in allocations and a modest speedup on CPU and GPUs for both training and inference.Note: A helper function,
dropout_maskwas used in the previous operation (it was basically integrated directly into the dropout implementation in this PR). I can't tell if other packages are using this, so I didn't remove it. Perhaps it should be marked deprecated and then removed later?Some benchmarks, first on my not great computer/GPU and then on a better one that was lent to me:
An example for gradient calculation.
m1has the currentDropoutimplementation, andm2has the updated one: