Slow ∇softmax! compared with generic version.

[jumerckx]

The specialized backward-pass for softmax takes a lot longer than the generic implementation from NNlib.jl.
The effect seems especially pronounced when the batch-dimension is larger.

Here's the code to reproduce this issue.

Below are the results of my benchmarks:

my CUDA versioninfo:

CUDA toolkit 11.4.1, artifact installation
CUDA driver 11.2.0
NVIDIA driver 461.9.0

Libraries: 
- CUBLAS: 11.5.4
- CURAND: 10.2.5
- CUFFT: 10.5.1
- CUSOLVER: 11.2.0
- CUSPARSE: 11.6.0
- CUPTI: 14.0.0
- NVML: 11.0.0+461.9
- CUDNN: 8.20.2 (for CUDA 11.4.0)
- CUTENSOR: 1.3.0 (for CUDA 11.2.0)

Toolchain:
- Julia: 1.6.1
- LLVM: 11.0.1
- PTX ISA support: 3.2, 4.0, 4.1, 4.2, 4.3, 5.0, 6.0, 6.1, 6.3, 6.4, 6.5, 7.0
- Device capability support: sm_35, sm_37, sm_50, sm_52, sm_53, sm_60, sm_61, sm_62, sm_70, sm_72, sm_75, sm_80

1 device:
  0: GeForce RTX 2070 with Max-Q Design (sm_75, 5.873 GiB / 8.000 GiB available)

[DhairyaLGandhi]

We could try with forcing using the CUDNN kernel directly as a debugging step, sidestepping the case of contiguous dimensions.

[DhairyaLGandhi]

Would you be able to post the CUDA.versioninfo() as well

[jumerckx]

We could try with forcing using the CUDNN kernel directly as a debugging step, sidestepping the case of contiguous dimensions.

Indeed, just using _∇softmax! directly and thereby avoiding CUDNN runs in ~50µs instead of 2ms

Would you be able to post the CUDA.versioninfo() as well

Updated in OP.

[DhairyaLGandhi]

What happens if we use the CUDNN kernel instead?

I'd love a profile of the cudnn run, that should tell us what's going on

[jumerckx]

I can't run Nvidia Nsight profiling on my machine but I timed all the lines which does confirm it is the kernel evaluation that causes the slowdown.

# ... after having ran code from OP ...
using NNlibCUDA
s = NNlibCUDA.softmaxdims(b, 1) # 1.030 μs (3 allocations: 112 bytes)
xDesc = NNlibCUDA.cudnnTensorDescriptor(reshape(b,s)) # 1.500 μs (5 allocations: 256 bytes)
alpha, beta = NNlibCUDA.scalingParameter(Float32,1), NNlibCUDA.scalingParameter(Float32,0) # 1.010 μs (4 allocations: 96 bytes)
NNlibCUDA.cudnnSoftmaxBackward(NNlibCUDA.handle(), NNlibCUDA.softmaxalgo(), NNlibCUDA.CUDNN_SOFTMAX_MODE_CHANNEL, alpha, xDesc, y_b, xDesc, Δ_b, beta, xDesc, out_b) # 1.998 ms (98 allocations: 9.55 KiB)

[ToucheSir]

What is the softmax algorithm being selected, and what happens if you force it to use the fast one? https://github.com/FluxML/NNlibCUDA.jl/blob/5eed446f379183e59bad583790fb7491b507743d/src/cudnn/softmax.jl#L66

[jumerckx]

By default it's CUDNN_SOFTMAX_ACCURATE but using CUDA.math_mode!(CUDA.FAST_MATH)
to use CUDNN_SOFTMAX_ACCURATE doesn't lead to a discernible difference in timing

[ToucheSir]

Here's a benchmark for PyTorch.

--------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                            Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls  
--------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                   aten::softmax         5.26%       5.573ms        24.98%      26.467ms      26.467us       4.671ms         4.16%      25.506ms      25.506us          1000  
                  aten::_softmax        16.06%      17.020ms        19.72%      20.895ms      20.895us      20.836ms        18.56%      20.836ms      20.836us          1000  
                aten::empty_like         4.22%       4.467ms        11.29%      11.959ms       3.986us       0.000us         0.00%       0.000us       0.000us          3000  
                     aten::empty         8.82%       9.345ms         8.82%       9.345ms       2.336us       0.000us         0.00%       0.000us       0.000us          4000  
                       aten::sum        15.01%      15.900ms        17.23%      18.255ms      18.255us      27.518ms        24.51%      27.518ms      27.518us          1000  
                aten::as_strided         1.96%       2.081ms         1.96%       2.081ms       1.040us       0.000us         0.00%       0.000us       0.000us          2000  
                 aten::ones_like         4.94%       5.230ms        14.20%      15.042ms      15.042us       7.787ms         6.94%      13.961ms      13.961us          1000  
             aten::empty_strided         2.15%       2.283ms         2.15%       2.283ms       2.283us       0.000us         0.00%       0.000us       0.000us          1000  
                     aten::fill_         5.70%       6.045ms         5.70%       6.045ms       6.045us       6.173ms         5.50%       6.173ms       6.173us          1000  
                    SumBackward0         4.35%       4.611ms         9.20%       9.745ms       9.745us       9.016ms         8.03%       9.016ms       9.016us          1000  
                    aten::expand         3.70%       3.922ms         4.84%       5.134ms       5.134us       0.000us         0.00%       0.000us       0.000us          1000  
                 SoftmaxBackward         5.65%       5.989ms        34.40%      36.453ms      36.453us       5.843ms         5.21%      36.254ms      36.254us          1000  
    aten::_softmax_backward_data        12.96%      13.734ms        28.75%      30.464ms      30.464us      16.071ms        14.32%      30.411ms      30.411us          1000  
                       aten::mul         9.21%       9.762ms        11.71%      12.413ms      12.413us      14.340ms        12.77%      14.340ms      14.340us          1000  
--------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 105.961ms
Self CUDA time total: 112.255ms

Benchmark code:

import torch
import torch.utils.benchmark as benchmark

x = torch.rand(32, 10, 256, requires_grad=True).cuda()
with torch.autograd.profiler.profile(use_cuda=True) as prof:
    for _ in range(1000):
        sx = x.softmax(-1).sum()
        torch.autograd.grad([sx], [x])

print(prof.key_averages().table())

As you can see, the times are very similar to the "backup" implementation at https://github.com/FluxML/NNlibCUDA.jl/blob/06ccd9f5b0fa6d3bfc9c9d52dbf865a78d76a576/src/cudnn/softmax.jl#L33. I was curious to see how they managed to make cudnnSoftmaxBackward run faster, and to my surprise they don't seem to use it at all! Looking back at JuliaGPU/CUDA.jl#599, it seems only the forward pass was thoroughly benchmarked and not the backwards function. If we can show that the native Julia version beats out the CuDNN one for a reasonably large number of sizes, dims, eltypes and math modes, I'd vote for using it by default and removing the CuDNN path.

[CarloLucibello]

it's so weird CUDNN is so slow. Anyways, yes, if @jumerckx can compare performance on a few more size and results are the same we can remove the cudnn path

[DhairyaLGandhi]

We don't need to remove the cudnn paths, but make Julia default if this seems accurate enough for our use-cases.

[maleadt]

Looks like it really is CUDNN being slow:

[jumerckx]

it's so weird CUDNN is so slow. Anyways, yes, if @jumerckx can compare performance on a few more size and results are the same we can remove the cudnn path

I've ran these timings across a bunch of sizes and haven't encountered one config in which CuDNN beat generic Julia (do please check my benchmarking code to assure I'm not doing anything stupid as this is the first time I do benchmarking like this).
One thing that's really striking is the difference between enlarging the first or second dimension of the input. When the second dimension is large, CuDNN scales particularly bad (which is also what I ran into for the OP):
Note the logarithmic axes.

I've updated the github gist with my benchmarking code but here's the (incomplete) essence:

# case for growing second dimension:
timings = []
for i in 1:12
  sz = (32, 2^i)
  
  x = CUDA.rand(sz...)
  dy = CUDA.rand(size(x)...)
  y = softmax(x)
  dx = similar(x)
  @show i
  time_generic = @belapsed CUDA.@sync NNlibCUDA._∇softmax!($dx, $dy, $x, $y; dims=1)
  
  s = NNlibCUDA.softmaxdims(x, 1)
  xDesc = NNlibCUDA.cudnnTensorDescriptor(reshape(x,s))
  alpha, beta = NNlibCUDA.scalingParameter(Float32,1), NNlibCUDA.scalingParameter(Float32,0)
  time_cudnn = @belapsed CUDA.@sync NNlibCUDA.cudnnSoftmaxBackward($NNlibCUDA.handle(), $NNlibCUDA.softmaxalgo(), $NNlibCUDA.CUDNN_SOFTMAX_MODE_CHANNEL, $alpha, $xDesc, $y, $xDesc, $dy, $beta, $xDesc, $dx)
  push!(timings, (sz, time_generic, time_cudnn))
end

[ToucheSir]

Those results look pretty conclusive, thanks @jumerckx! Does anyone want to do the honours?