Skip to content

[MLDataDevices] CUDA backend should not require cuDNN #1728

Description

@AntonOresten

Related: #1245, LuxDL/MLDataDevices.jl#51, FluxML/Flux.jl#2440

The problem

CUDADevice does not register as loaded or functional unless cuDNN is present. Both methods live in cuDNNExt, not CUDAExt:

# cuDNNExt.jl — these are the only things making CUDADevice discoverable
MLDataDevices.loaded(::Union{CUDADevice,Type{<:CUDADevice}}) = true
MLDataDevices.functional(::Union{CUDADevice,Type{<:CUDADevice}})::Bool = ...

Meanwhile, all the actual data-transfer machinery — adapt_storage, set_device!, get_device, unsafe_free!, default_device_rng — lives in CUDAExt and never touches cuDNN. cuDNNExt provides no data-transfer functionality at all; its entire role is to flip the switch that says the backend exists.

So with CUDA loaded and fully working:

  • gpu_device() falls back to CPU (or throws with force=true);
  • even an explicit gpu_backend!("CUDA") preference cannot opt out — it routes through the same loaded/functional checks;
  • the only workaround is bypassing gpu_device() and constructing CUDADevice() by hand, which the gpu_device docstring itself recommends: "if cuDNN is not loaded you can still manually create a CUDADevice object and use it."

Reproducer:

# CUDA v6.1.0, Flux v0.16.10, no cuDNN in the environment
julia> x = rand(Float32, 4);

julia> CUDA.functional()
true

julia> gpu_device()             # ┌ Warning: No functional GPU backend found! Defaulting to CPU. ...
CPUDevice{Missing}()

julia> Flux.gpu(x) |> typeof    # silent no-op
Vector{Float32}

julia> CUDADevice()(x) |> typeof   # the "non-functional" device, moving data fine
CuArray{Float32, 1, CUDACore.DeviceMemory}

That last line and the docstring note are the tell: by the package's own admission, the device works without cuDNN. The gate is a deep-learning-stack policy enforced inside a data-movement package.

The package already does this right — for AMD

The AMDGPU extension faces the identical situation with MIOpen (AMD's cuDNN equivalent) and makes the opposite choice:

# AMDGPUExt.jl
USE_AMD_GPU[] = AMDGPU.functional()
if USE_AMD_GPU[] && !AMDGPU.functional(:MIOpen)
    @warn "MIOpen is not functional in AMDGPU.jl, some functionality will not be \
           available." maxlog = 1
end

It warns and stays functional. The same policy for CUDA would not be novel — it's already this package's behavior for the other GPU vendor. The asymmetry is a leftover from the LuxDeviceUtils era, when assuming the full Lux DL stack (LuxCUDA = CUDA + cuDNN) was reasonable; after the deliberate rebrand to a framework-agnostic package (LuxDL/MLDataDevices.jl#51), it no longer is.

Nor is it just the AMD path: sibling packages in this same monorepo don't gate on cuDNN either. WeightInitializers' CUDA extension is triggered by CUDA alone and constructs CuArrays by dispatching on CUDA RNG types — so kaiming_normal(CUDA.default_rng(), 64, 64) happily allocates on the GPU while gpu_device() insists that same GPU isn't functional.

Who this affects

Lux users are unaffected: LuxCUDA bundles cuDNN, so nothing changes for them. Affected is everyone else:

  • Flux, whose gpu(x) is defined as gpu_device()(x) — the gate decides whether Flux's GPU movement works at all, despite none of it touching cuDNN. Flux handles cuDNN at its own layer (auto-loading it when installed, warning with install instructions when not), so the gate is redundant when cuDNN is installed, and harmful when it isn't: Flux's warning promises "some Flux functionalities will not be available", but the gate downgrades gpu_device() to CPU and the whole model trains on CPU.
  • DL users who don't use cuDNN. cuDNN is a fixed menu of kernels, not a prerequisite for deep learning on NVIDIA GPUs. Anyone writing their own kernels — with plain CUDA.jl, KernelAbstractions.jl, or cuTile.jl (NVIDIA's tile-based programming model, positioned precisely as the path to custom high-performance kernels in place of fixed-function libraries like cuDNN) — has no use for cuDNN, yet must install and load it anyway, purely for data movement.
  • Non-DL users — scientific computing, anyone who loads CUDA without cuDNN and just wants arrays moved.

Prior discussion

I'm raising this again because neither discussion engaged with the layering question.

The FluxML/Flux.jl#2440 concern, at the right layer

The concern behind the gate is real: without cuDNN, a GPU convolution silently falls back to the CPU implementation and dies with a confusing scalar-indexing error. Flux#2440 was even closed citing this gate as the fix. But the gate is a shotgun: it prevents the Conv-without-cuDNN footgun by preventing every non-Conv use of CUDA through gpu_device().

The footgun belongs to the layers that own the operations: NNlib.jl#523 ("Add warnings/errors on CPU kernels when called with GPU arrays", still open) is the actual fix, as was noted in Flux#2440 itself — and the frameworks that need cuDNN already take responsibility for it (Flux auto-loads it, Lux ships LuxCUDA). Nothing is lost diagnostically either: cuDNN.jl itself already emits an @error at load time when it's installed but broken.

Proposed fix

Move loaded and functional into CUDAExt, and delete cuDNNExt:

# CUDAExt
__init__() = reset_gpu_device!()   # currently lives in cuDNNExt; must move too

MLDataDevices.loaded(::Union{CUDADevice,Type{<:CUDADevice}}) = true
MLDataDevices.functional(::Union{CUDADevice,Type{<:CUDADevice}})::Bool = CUDA.functional()

With the gate gone, nothing else in cuDNNExt has a job left: its loaded-but-broken warning duplicates the @error cuDNN.jl itself emits at load time, and its reset_gpu_device! init hook only existed because loading cuDNN used to change the outcome of device selection. If there's a reason for cuDNNExt to exist in a data-movement package beyond gating these two methods, I'd be interested to hear it — but looking at the current code, there isn't one.

Plus doc updates where cuDNN is currently listed as a required trigger package (loaded docstring, the gpu_device admonition, the "No functional GPU backend" warning text).

Lux users loading LuxCUDA are completely unaffected. The only behavioral change is that using CUDA, MLDataDevices; gpu_device() returns a CUDADevice instead of falling back to CPU — which makes CUDA behave exactly like AMDGPU, Metal, and oneAPI already do.

Drafted with Fable 5, sourced from a months-long gripe.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions