RNN on GPU fails on first backward call

[jeremiedb]

In the following minimal RNN example, first call to Flux.train! fails with CUDNN_STATUS_BAD_PARAM. Same error raised either cell is GRU, RNN or LSTM.

using Flux
using Statistics

rnn = Chain(GRU(16, 8),
  Dense(8,1, σ),
  x -> reshape(x,:))

X = [rand(16,10) for i in 1:20]
Y = rand(10,20) ./ 10

rnn = rnn |> gpu
X = gpu(X)
Y = gpu(Y)

θ = Flux.params(rnn)
loss(x,y) = mean((Flux.stack(rnn.(X),2) .- y) .^ 2f0)
opt = ADAM(1e-3)
size(rnn[1].state)
Flux.reset!(rnn)
size(rnn[1].state)
Flux.train!(loss, θ, [(X,Y)], opt)
size(rnn[1].state)
loss(X,Y)

It can be observed that both prior and after reset!, rnn state is of size (8), while after a call to train! on GPU, state becomes the expected proper size (8,10). After each call to reset!, the CUDNN_STATUS_BAD_PARAM error pops out after first call to train!, but subsequent ones are fine as the state size stays (8,10). Can't confirm whether that state size is the root cause, but appears closely tied to the bug. Also, a call to loss(X,Y) results in a proper state dimension of (8,10). Running on CPU doesn't result in any error/warning.

Pkg info (same error also raised wih latest Zygote release):

  [4f1ea46c] AWSCore v0.6.9
  [1c724243] AWSS3 v0.6.10
  [fbb218c0] BSON v0.2.5
  [336ed68f] CSV v0.6.1
  [3a865a2d] CuArrays v2.0.1
  [a93c6f00] DataFrames v0.20.2
  [587475ba] Flux v0.10.4
  [28b8d3ca] GR v0.48.0
  [91a5bcdd] Plots v1.0.5
  [2913bbd2] StatsBase v0.32.2
  [e88e6eb3] Zygote v0.4.15 #master (https://github.com/FluxML/Zygote.jl.git)
  [10745b16] Statistics

CUDA: 10.1.168
CUDNN: 7.6.5

[DhairyaLGandhi]

Seems like we're sending something off in our parameters in the backwards pass

[maleadt]

Smaller repro:

using Flux
using Statistics
using Random

function main()
    Random.seed!(1)
    rnn = GRU(1, 1)

    X = [rand(1,2) for i in 1:2]
    Y = rand(2,2) ./ 10

    drnn = rnn |> gpu
    dX = gpu(X)
    dY = gpu(Y)

    θ = Flux.params(drnn)
    loss(x,y) = mean((Flux.stack(drnn.(dX),2) .- y) .^ 2f0)
    opt = ADAM(1e-3)

    try
        Flux.train!(loss, θ, [(dX,dY)], opt)
    catch ex
        @error "First training failed" exception=(ex, catch_backtrace())
    end

    Flux.train!(loss, θ, [(dX,dY)], opt)
    @info "Second training succeeded"
end

isinteractive() || main()

Looks like the dimensionality of the initial state is wrong. Initially, rnn.init is a 1-dimensional array [0], while after the first (failed) training the state becomes a 2-dimensional array with 2 elements. Indeed, fixing the initial state to be a 2-dimensional array of 2 zeros makes the first training work too:

drnn.state = CuArrays.zeros(Float32, (1,2))
Flux.train!(loss, θ, [(dX,dY)], opt)
# works at the first invocation

I'm not sure how/where that data is set and used, so I'll leave that to Flux experts 🙂

[DhairyaLGandhi]

The issue seems more to do with that our init breaks now with the shape that we have. The repro should probably add the call to Flux.reset!, since this breakage will happen every time we call the reset! function, when our state is set to zeros(8), instead of zeros(8,10)

[DhairyaLGandhi]

(Or zeros(1,2)), in the case of the repro.

[jeremiedb]

The CUDNNError: CUDNN_STATUS_BAD_PARAM (code 3) bug continues to show out on latest Flux and Zygote, but only for LSTM:

(Flux) pkg> st
Project Flux v0.11.1
Status `D:\Github\Flux.jl\Project.toml`
  [1520ce14] AbstractTrees v0.3.3
  [79e6a3ab] Adapt v2.3.0
  [052768ef] CUDA v1.3.3
  [944b1d66] CodecZlib v0.7.0
  [5ae59095] Colors v0.12.4
  [d9f16b24] Functors v0.1.0
  [e5e0dc1b] Juno v0.8.4
  [1914dd2f] MacroTools v0.5.6
  [872c559c] NNlib v0.7.5
  [189a3867] Reexport v0.2.0
  [2913bbd2] StatsBase v0.33.2
  [a5390f91] ZipFile v0.9.3
  [e88e6eb3] Zygote v0.5.9
  [8bb1440f] DelimitedFiles
  [37e2e46d] LinearAlgebra
  [44cfe95a] Pkg
  [de0858da] Printf
  [9a3f8284] Random
  [ea8e919c] SHA
  [10745b16] Statistics
  [8dfed614] Test

MRE:

rnn = Chain(LSTM(3, 4),
  Dense(4,1, σ),
  x -> reshape(x,:))

X = [rand(3,5) for i in 1:10]
Y = rand(5,10) ./ 10

rnn = rnn |> gpu
X = gpu(X)
Y = gpu(Y)

θ = Flux.params(rnn)
loss(x,y) = mean((Flux.stack(rnn.(x),2) .- y) .^ 2f0)
opt = ADAM(1e-3)
Flux.train!(loss, θ, [(X,Y)], opt)
Flux.train!(loss, θ, [(X,Y)], opt)
Flux.reset!(rnn)
Flux.train!(loss, θ, [(X,Y)], opt)

The error shows that the pullback step would be in error:

Flux.jl/src/cuda/curnn.jl

Line 85 in c5c35cc

m̄ = back(dy, dho, dco)

Which relates to the following CUDA.jl:
https://github.com/JuliaGPU/CUDA.jl/blob/31f67d6caf7fca566a9c7b36b850ce31df450ac6/lib/cudnn/rnn.jl#L195

I couldn't figure what's different between the RNN/GRU vs the LSTM that would cause a bug only in the later. I'm not clear about this line where the cell state c_ for LSTM isn't used: (dWi, dWh), db = CUDNN.backwardWeights(rnn, x, h_, y, reserve), is it normal?

[CarloLucibello]

@jeremiedb is this fixed by your PR?

[jeremiedb]

Yes it's now working in master. Closing the issue.