Error while using GPU on UNET with FLUX in Julia - julia

Question

I'm using Julia for UNet training with GPU. (Runs without issues on CPU).
The error is at NNlib's function conv_im2col()
function conv_im2col!(
y::AbstractArray{T,5}, x::AbstractArray{T,5},
w::AbstractArray{T,5}, cdims::DenseConvDims;
col::AbstractArray{T,3}=similar(x, im2col_dims(cdims)),
alpha::T=T(1), beta::T=T(0)) where {T}
check_dims(size(x), size(w), size(y), cdims)
# COL * W -> Y
# [M x K] * [K x N] -> [M x N]
#
# M: output spatial resolution
# N: output channels
# K: size of input "patch" (kernel size and input channels combined)
#
# In english, we're grabbing each input patch and laying them out along
# the M dimension in `col`, so that the GEMM call below multiplies each
# kernel (which is kernel_h * kernel_w * channels_in elments long) is
# dotproducted with that input patch, effectively computing a convolution
# in a somewhat memory-wasteful but easily-computed way (since we already
# have an extremely highly-optimized GEMM call available in BLAS).
M = prod(output_size(cdims))
N = channels_out(cdims)
K = prod(kernel_size(cdims))*channels_in(cdims)
#threads for batch_idx in 1:size(x,5)
# col_slice is a thread-local workspace
col_slice = view(col, :, :, threadid())
im2col!(col_slice, view(x, :, :, :, :, batch_idx), cdims)
GC.#preserve col_slice w y begin
col_ptr = pointer(col_slice)
w_ptr = pointer(w)
y_ptr = pointer(y, (batch_idx - 1)*M*N + 1)
gemm!(Val(false), Val(false), M, N, K, alpha, col_ptr, w_ptr, beta, y_ptr)
end
and the detailed error is:
[ Info: training... [ Info: Using Device gpu Epoch Minibatch Loss
Test Dice Coef
---------------------------------------------- 1ERROR: LoadError: TaskFailedException
nested task error: TaskFailedException
Stacktrace:
[1] wait
# ./task.jl:322 [inlined]
[2] threading_run(func::Function)
# Base.Threads ./threadingconstructs.jl:34
[3] macro expansion
# ./threadingconstructs.jl:93 [inlined]
[4] conv_im2col!(y::SubArray{Float32, 5, Array{Float32, 5}, Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}},
Base.Slice{Base.OneTo{Int64}}, UnitRange{Int64},
Base.Slice{Base.OneTo{Int64}}}, false}, x::SubArray{Float32, 5,
Array{Float32, 5}, Tuple{Base.Slice{Base.OneTo{Int64}},
Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}},
UnitRange{Int64}, Base.Slice{Base.OneTo{Int64}}}, false},
w::CuArray{Float32, 5, CUDA.Mem.DeviceBuffer}, cdims::DenseConvDims{3,
3, 3, 6, 3}; col::Array{Float32, 3}, alpha::Float32, beta::Float32)
# NNlib ~/.julia/packages/NNlib/0QnJJ/src/impl/conv_im2col.jl:47
[5] conv_im2col! # ~/.julia/packages/NNlib/0QnJJ/src/impl/conv_im2col.jl:28 [inlined]
[6] (::NNlib.var"#262#266"{Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, DenseConvDims{3, 3, 3, 6, 3},
SubArray{Float32, 5, Array{Float32, 5},
Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}},
Base.Slice{Base.OneTo{Int64}}, UnitRange{Int64},
Base.Slice{Base.OneTo{Int64}}}, false}, CuArray{Float32, 5,
CUDA.Mem.DeviceBuffer}, SubArray{Float32, 5, Array{Float32, 5},
Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}},
Base.Slice{Base.OneTo{Int64}}, UnitRange{Int64},
Base.Slice{Base.OneTo{Int64}}}, false}})()
# NNlib ./threadingconstructs.jl:169
nested task error: MethodError: no method matching gemm!(::Val{false}, ::Val{false}, ::Int64, ::Int64, ::Int64,::Float32, ::Ptr{Float32}, ::CuPtr{Float32}, ::Float32, ::Ptr{Float32})
Closest candidates are:
gemm!(::Val, ::Val, ::Int64, ::Int64, ::Int64, ::Float32, ::Ptr{Float32}, ::Ptr{Float32}, ::Float32, ::Ptr{Float32}) at
.julia/packages/NNlib/0QnJJ/src/gemm.jl:29
gemm!(::Val, ::Val, ::Int64, ::Int64, ::Int64, ::Float64, ::Ptr{Float64}, ::Ptr{Float64}, ::Float64, ::Ptr{Float64}) at
.julia/packages/NNlib/0QnJJ/src/gemm.jl:29
gemm!(::Val, ::Val, ::Int64, ::Int64, ::Int64, ::ComplexF64, ::Ptr{ComplexF64}, ::Ptr{ComplexF64}, ::ComplexF64, ::Ptr{ComplexF64})
at .julia/packages/NNlib/0QnJJ/src/gemm.jl:29
...
Stacktrace:
[1] macro expansion
# ~/.julia/packages/NNlib/0QnJJ/src/impl/conv_im2col.jl:56 [inlined]
[2] (::NNlib.var"#909#threadsfor_fun#504"{Array{Float32, 3}, Float32, Float32, SubArray{Float32, 5, Array{Float32, 5},
Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}},
Base.Slice{Base.OneTo{Int64}}, UnitRange{Int64},
Base.Slice{Base.OneTo{Int64}}}, false}, SubArray{Float32, 5,
Array{Float32, 5}, Tuple{Base.Slice{Base.OneTo{Int64}},
Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}},
UnitRange{Int64}, Base.Slice{Base.OneTo{Int64}}}, false},
CuArray{Float32, 5, CUDA.Mem.DeviceBuffer}, DenseConvDims{3, 3, 3, 6,
3}, Int64, Int64, Int64, UnitRange{Int64}})(onethread::Bool)
# NNlib ./threadingconstructs.jl:81
[3] (::NNlib.var"#909#threadsfor_fun#504"{Array{Float32, 3}, Float32, Float32, SubArray{Float32, 5, Array{Float32, 5},
Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}},
Base.Slice{Base.OneTo{Int64}}, UnitRange{Int64},
Base.Slice{Base.OneTo{Int64}}}, false}, SubArray{Float32, 5,
Array{Float32, 5}, Tuple{Base.Slice{Base.OneTo{Int64}},
Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}},
UnitRange{Int64}, Base.Slice{Base.OneTo{Int64}}}, false},
CuArray{Float32, 5, CUDA.Mem.DeviceBuffer}, DenseConvDims{3, 3, 3, 6,
3}, Int64, Int64, Int64, UnitRange{Int64}})()
# NNlib ./threadingconstructs.jl:48 Stacktrace: [1] sync_end(c::Channel{Any})
# Base ./task.jl:369 [2] macro expansion
# ./task.jl:388 [inlined] [3] conv!(out::Array{Float32, 5}, in1::Array{Float32, 5}, in2::CuArray{Float32, 5,
CUDA.Mem.DeviceBuffer}, cdims::DenseConvDims{3, 3, 3, 6, 3};
kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(),
Tuple{}}})
# NNlib ~/.julia/packages/NNlib/0QnJJ/src/conv.jl:195 [4] conv!
# ~/.julia/packages/NNlib/0QnJJ/src/conv.jl:186 [inlined] [5] conv!(y::Array{Float32, 4}, x::Array{Float32, 4}, w::CuArray{Float32,
4, CUDA.Mem.DeviceBuffer}, cdims::DenseConvDims{2, 2, 2, 4, 2};
kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(),
Tuple{}}})
# NNlib ~/.julia/packages/NNlib/0QnJJ/src/conv.jl:145 [6] conv!
# ~/.julia/packages/NNlib/0QnJJ/src/conv.jl:145 [inlined] [7] conv(x::Array{Float32, 4}, w::CuArray{Float32, 4,
CUDA.Mem.DeviceBuffer}, cdims::DenseConvDims{2, 2, 2, 4, 2};
kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(),
Tuple{}}})
# NNlib ~/.julia/packages/NNlib/0QnJJ/src/conv.jl:88 [8] conv
# ~/.julia/packages/NNlib/0QnJJ/src/conv.jl:86 [inlined] [9] #rrule#312
# ~/.julia/packages/NNlib/0QnJJ/src/conv.jl:313 [inlined] [10] rrule
# ~/.julia/packages/NNlib/0QnJJ/src/conv.jl:304 [inlined] [11] rrule
# ~/.julia/packages/ChainRulesCore/ctmSK/src/rules.jl:134 [inlined] [12] chain_rrule
# ~/.julia/packages/Zygote/xGkZ5/src/compiler/chainrules.jl:218 [inlined] [13] macro expansion
# ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface2.jl:0 [inlined] [14] _pullback
# ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface2.jl:9 [inlined] [15] _pullback
# ~/.julia/packages/Flux/KkC79/src/layers/conv.jl:200 [inlined] [16] macro expansion
# ~/.julia/packages/Flux/KkC79/src/layers/basic.jl:53 [inlined] [17] _pullback
# ~/.julia/packages/Flux/KkC79/src/layers/basic.jl:53 [inlined] [18] _pullback(::Zygote.Context{true}, ::typeof(Flux._applychain),
::WARNING: both Losses and NNlib export "ctc_loss"; uses of it in
module Flux must be qualified Tuple{Conv{2, 2, typeof(identity),
CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1,
CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(selu), CuArray{Float32, 1,
CUDA.Mem.DeviceBuffer}, Float32, CuArray{Float32, 1,
CUDA.Mem.DeviceBuffer}}, Conv{2, 2, typeof(identity), CuArray{Float32,
4, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1,
CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(selu), CuArray{Float32, 1,
CUDA.Mem.DeviceBuffer}, Float32, CuArray{Float32, 1,
CUDA.Mem.DeviceBuffer}}}, ::Array{Float32, 4})
# Zygote ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface2.jl:0 [19]
_pullback
# ~/.julia/packages/Flux/KkC79/src/layers/basic.jl:51 [inlined] [20] _pullback(ctx::Zygote.Context{true}, f::Chain{Tuple{Conv{2, 2,
typeof(identity), CuArray{Float32, 4, CUDA.Mem.DeviceBuffer},
CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(selu),
CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CuArray{Float32,
1, CUDA.Mem.DeviceBuffer}}, Conv{2, 2, typeof(identity),
CuArray{Float32, 4, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1,
CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(selu), CuArray{Float32, 1,
CUDA.Mem.DeviceBuffer}, Float32, CuArray{Float32, 1,
CUDA.Mem.DeviceBuffer}}}}, args::Array{Float32, 4})
# Zygote ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface2.jl:0 [21]
_pullback
# ~/.julia/dev/UNet/src/model.jl:93 [inlined] [22] _pullback(ctx::Zygote.Context{true}, f::Unet, args::Array{Float32, 4})
# Zygote ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface2.jl:0 [23]
_pullback
# ~/.julia/dev/UNet/src/train.jl:37 [inlined] [24] _pullback(::Zygote.Context{true}, ::UNet.var"#loss#75"{typeof(Flux.Losses.logitcrossentropy), Unet},
::Array{Float32, 4}, ::Array{Int32, 4})
# Zygote ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface2.jl:0 [25]
_apply
# ./boot.jl:804 [inlined] [26] adjoint
# ~/.julia/packages/Zygote/xGkZ5/src/lib/lib.jl:203 [inlined] [27] _pullback
# ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined] [28] _pullback
# ~/.julia/packages/Flux/KkC79/src/optimise/train.jl:120 [inlined] [29] _pullback(::Zygote.Context{true},
::Flux.Optimise.var"#37#40"{UNet.var"#loss#75"{typeof(Flux.Losses.logitcrossentropy),
Unet}, Tuple{Array{Float32, 4}, Array{Int32, 4}}})
# Zygote ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface2.jl:0 [30]
pullback(f::Function, ps::Zygote.Params{Zygote.Buffer{Any,
Vector{Any}}})
# Zygote ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface.jl:373 [31]
gradient(f::Function, args::Zygote.Params{Zygote.Buffer{Any,
Vector{Any}}})
# Zygote ~/.julia/packages/Zygote/xGkZ5/src/compiler/interface.jl:96 [32]
macro expansion
# ~/.julia/packages/Flux/KkC79/src/optimise/train.jl:119 [inlined] [33] macro expansion
# ~/.julia/packages/ProgressLogging/6KXlp/src/ProgressLogging.jl:328
[inlined] [34] train!(loss::Function,
ps::Zygote.Params{Zygote.Buffer{Any, Vector{Any}}},
data::MLUtils.DataLoader{Tuple{Array{Float32, 4}, Array{Int32, 4}},
Random._GLOBAL_RNG}, opt::RMSProp; cb::UNet.var"#77#78"{Unet,
MLUtils.DataLoader{Tuple{Array{Float32, 4}, Array{Int32, 4}},
Random._GLOBAL_RNG}})
# Flux.Optimise ~/.julia/packages/Flux/KkC79/src/optimise/train.jl:117 [35]
train(train_dataset::ImageDataset, test_dataset::ImageDataset,
model::Unet, opt::RMSProp,
loss::UNet.var"#loss#75"{typeof(Flux.Losses.logitcrossentropy), Unet};
batch_size::Int64, num_epochs::Int64,
learning_rate_drop_rate::Float64, learning_rate_step::Int64,
device::Function, save_intermediate_models::Bool, save_model::Bool,
out_dir::String, model_file_name::String)
# UNet ~/.julia/dev/UNet/src/train.jl:123 [36] train(train_dataset::ImageDataset, test_dataset::ImageDataset,
num_channels::Int64, num_labels::Int64;
loss_function::typeof(Flux.Losses.logitcrossentropy),
learning_rate::Float64, momentum::Float64, device::Function,
kwargs::Base.Iterators.Pairs{Symbol, Any, NTuple{6, Symbol},
NamedTuple{(:num_epochs, :batch_size, :save_intermediate_models,
:save_model, :out_dir, :model_file_name), Tuple{Int64, Int64, Bool,
Bool, String, String}}})
# UNet ~/.julia/dev/UNet/src/train.jl:40 [37] top-level scope
# /train_images.jl:70 [38] include(fname::String)
# Base.MainInclude ./client.jl:444 [39] top-level scope
# REPL[2]:1 in expression starting at train_images.jl:70