BatchNorm Memory Blowup on Very Large CUDA Arrays

[BenCurran98]

Hi,

This is probably less of a bug and more a discussion/feature request, but I've noticed when transferring models from pytorch to Flux that the BatchNorm layer seems to use a lot more memory on the GPU for larger arrays than Pytorch does. This proves to be a bottleneck when doing e.g. graph-based models on very large datasets (~100k nodes). Given one could just increase the size of GPU they're using, it's obviously not a blocker, but if Pytorch can handle the same case on the same GPU it might be worth investigating why? Any help/insight would be appreciated :)

As a MWE, I ran the following on an 8GB NVIDIA RTX 2000 device with CUDA and Flux:

using Flux
using CUDA

m = BatchNorm(32) |> gpu
m(CUDA.rand(32, 800_000))

Which gives the following output:

ERROR: CUDNNError: CUDNN_STATUS_NOT_SUPPORTED (code 3000)
Stacktrace:
  [1] throw_api_error(res::cuDNN.cudnnStatus_t)
    @ cuDNN ~/.julia/packages/cuDNN/7odoD/src/libcudnn.jl:15
  [2] check
    @ ~/.julia/packages/cuDNN/7odoD/src/libcudnn.jl:26 [inlined]
  [3] cudnnBatchNormalizationForwardInference
    @ ~/.julia/packages/GPUToolbox/cZlg7/src/ccalls.jl:33 [inlined]
  [4] cudnnBNForward!(y::CuArray{…}, g::CuArray{…}, b::CuArray{…}, x::CuArray{…}, running_mean::CuArray{…}, running_var::CuArray{…}, momentum::Float32; cache::Nothing, alpha::Int64, beta::Int64, eps::Float32, training::Bool, affine::Bool, track_stats::Bool)
    @ NNlibCUDACUDNNExt ~/.julia/packages/NNlib/1TYHL/ext/NNlibCUDACUDNNExt/batchnorm.jl:88
  [5] cudnnBNForward!
    @ ~/.julia/packages/NNlib/1TYHL/ext/NNlibCUDACUDNNExt/batchnorm.jl:40 [inlined]
  [6] #batchnorm#64
    @ ~/.julia/packages/NNlib/1TYHL/ext/NNlibCUDACUDNNExt/batchnorm.jl:37 [inlined]
  [7] batchnorm
    @ ~/.julia/packages/NNlib/1TYHL/ext/NNlibCUDACUDNNExt/batchnorm.jl:35 [inlined]
  [8] #batchnorm#63
    @ ~/.julia/packages/NNlib/1TYHL/ext/NNlibCUDACUDNNExt/batchnorm.jl:31 [inlined]
  [9] (::BatchNorm{typeof(identity), CuArray{Float32, 1, CUDA.DeviceMemory}, Float32, CuArray{Float32, 1, CUDA.DeviceMemory}})(x::CuArray{Float32, 2, CUDA.DeviceMemory}, cache::Nothing)
    @ FluxCUDAcuDNNExt ~/.julia/packages/Flux/uRn8o/ext/FluxCUDAcuDNNExt/FluxCUDAcuDNNExt.jl:13
 [10] (::BatchNorm{typeof(identity), CuArray{Float32, 1, CUDA.DeviceMemory}, Float32, CuArray{Float32, 1, CUDA.DeviceMemory}})(x::CuArray{Float32, 2, CUDA.DeviceMemory})
    @ FluxCUDAcuDNNExt ~/.julia/packages/Flux/uRn8o/ext/FluxCUDAcuDNNExt/FluxCUDAcuDNNExt.jl:10
 [11] top-level scope
    @ REPL[19]:1
Some type information was truncated. Use `show(err)` to see complete types.

I can run the following using PyTorch v2.1.2 on Python v3.12.2:

import torch
device = torch.device("cuda")
m = torch.nn.BatchNorm1d(32).to(device)
input = torch.randn(800000, 32).to(device)
output = m(input)

[BenCurran98]

@CarloLucibello any thoughts on this one? This seems to also be happening using Lux.jl, probably because both packages use cuDNN.jl so it might be a problem there?

[avik-pal]

This seems to also be happening using Lux.jl

For Lux this shouldn't be an issue

julia> using Reactant, Lux, Random

julia> model = BatchNorm(32)
BatchNorm(32, affine=true, track_stats=true)  # 64 parameters, plus 65 non-trainable

julia> ps, st = Lux.setup(Random.default_rng(), model) |> reactant_device();

julia> x = rand(Float32, 32, 800_000) |> reactant_device();

julia> @code_hlo model(x, ps, Lux.testmode(st))
module @reactant_BatchNo... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<800000x32xf32>, %arg1: tensor<32xf32>, %arg2: tensor<32xf32>, %arg3: tensor<32xf32>, %arg4: tensor<32xf32>) -> tensor<800000x32xf32> {
    %0 = "stablehlo.batch_norm_inference"(%arg0, %arg1, %arg2, %arg3, %arg4) <{epsilon = 9.99999974E-6 : f32, feature_index = 1 : i64}> : (tensor<800000x32xf32>, tensor<32xf32>, tensor<32xf32>, tensor<32xf32>, tensor<32xf32>) -> tensor<800000x32xf32>
    return %0 : tensor<800000x32xf32>
  }
}

julia> @jit model(x, ps, Lux.testmode(st))
(ConcretePJRTArray{Float32, 2, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(Float32[0.6103075 0.067417756 … 0.2430162 0.6654256; 0.9981578 0.21910104 … 0.53565955 0.62876964; … ; 0.85754937 0.31250307 … 0.3459001 0.23633033; 0.09042617 0.30842292 … 0.7321898 0.7129404]), (running_mean = ConcretePJRTArray{Float32, 1, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0  …  0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), running_var = ConcretePJRTArray{Float32, 1, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(Float32[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0  …  1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]), training = Val{false}()))

[BenCurran98]

This seems to also be happening using Lux.jl

For Lux this shouldn't be an issue

julia> using Reactant, Lux, Random

julia> model = BatchNorm(32)
BatchNorm(32, affine=true, track_stats=true) # 64 parameters, plus 65 non-trainable

julia> ps, st = Lux.setup(Random.default_rng(), model) |> reactant_device();

julia> x = rand(Float32, 32, 800_000) |> reactant_device();

julia> @code_hlo model(x, ps, Lux.testmode(st))
module @reactant_BatchNo... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<800000x32xf32>, %arg1: tensor<32xf32>, %arg2: tensor<32xf32>, %arg3: tensor<32xf32>, %arg4: tensor<32xf32>) -> tensor<800000x32xf32> {
%0 = "stablehlo.batch_norm_inference"(%arg0, %arg1, %arg2, %arg3, %arg4) <{epsilon = 9.99999974E-6 : f32, feature_index = 1 : i64}> : (tensor<800000x32xf32>, tensor<32xf32>, tensor<32xf32>, tensor<32xf32>, tensor<32xf32>) -> tensor<800000x32xf32>
return %0 : tensor<800000x32xf32>
}
}

julia> @jit model(x, ps, Lux.testmode(st))
(ConcretePJRTArray{Float32, 2, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(Float32[0.6103075 0.067417756 … 0.2430162 0.6654256; 0.9981578 0.21910104 … 0.53565955 0.62876964; … ; 0.85754937 0.31250307 … 0.3459001 0.23633033; 0.09042617 0.30842292 … 0.7321898 0.7129404]), (running_mean = ConcretePJRTArray{Float32, 1, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 … 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), running_var = ConcretePJRTArray{Float32, 1, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(Float32[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0 … 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]), training = Val{false}()))

That's very helpful, thank you! This should unblock me for now, although we should probably still discuss why this seems to be an issue in the underlying cuDNN framework. I was previously just using gpu_device() (not Reactant.jl) which was causing the same errors as with Flux.jl:

julia> using Reactant, Lux, Random

julia> model = BatchNorm(32)

julia> ps, st = Lux.setup(Random.default_rng(), model) |> gpu_device();

julia> x = rand(Float32, 32, 800_000) |> gpu_device();

julia> y, st = model(x, ps, Lux.testmode(st))

ERROR: CUDNNError: CUDNN_STATUS_NOT_SUPPORTED (code 3000)
Stacktrace:
  [1] throw_api_error(res::cuDNN.cudnnStatus_t)
    @ cuDNN ~/.julia/dev/CUDA/lib/cudnn/src/libcudnn.jl:15
  [2] check
    @ ~/.julia/dev/CUDA/lib/cudnn/src/libcudnn.jl:26 [inlined]
  [3] cudnnBatchNormalizationForwardInference
    @ ~/.julia/packages/GPUToolbox/cZlg7/src/ccalls.jl:33 [inlined]
  [4] batchnorm_cudnn!(y::CuArray{…}, γ′::CuArray{…}, β′::CuArray{…}, x::CuArray{…}, rμ′::CuArray{…}, rσ²′::CuArray{…}, m::Float32, ϵ::Float32, training::Static.False)
    @ LuxLibcuDNNExt ~/.julia/dev/Lux/lib/LuxLib/ext/LuxLibcuDNNExt/batchnorm.jl:103
  [5] batchnorm_cudnn
    @ ~/.julia/dev/Lux/lib/LuxLib/ext/LuxLibcuDNNExt/batchnorm.jl:40 [inlined]
  [6] batchnorm_cudnn
    @ ~/.julia/dev/Lux/lib/LuxLib/ext/LuxLibcuDNNExt/batchnorm.jl:27 [inlined]
  [7] batchnorm
    @ ~/.julia/dev/Lux/lib/LuxLib/ext/LuxLibcuDNNExt/LuxLibcuDNNExt.jl:38 [inlined]
  [8] batchnorm
    @ ~/.julia/dev/Lux/lib/LuxLib/src/api/batchnorm.jl:43 [inlined]
  [9] (::BatchNorm{…})(x::CuArray{…}, ps::@NamedTuple{…}, st::@NamedTuple{…})
    @ Lux ~/.julia/dev/Lux/src/layers/normalize.jl:146
 [10] top-level scope
    @ REPL[11]:1
Some type information was truncated. Use `show(err)` to see complete types.
``