Closed casper2002casper closed 2 years ago
Any update on this? It seems like a major issue to me as it slows down any form of inference significantly.
I tried to reduce the example, it turns out that the problem is not GNN.jl related but due to the allocation strategy of CUDA.jl (of which I know nothing about). In your original example the comparison wasn't measuring comparable operations for GNN and NN.
Here are a few comparisons
using GraphNeuralNetworks, CUDA, Flux
N = 10000
I = 10
function test_mem(n, data)
for i in 1:I
y = n(data)
CUDA.memory_status()
end
end
GC.gc(); CUDA.reclaim();
println("GNN, memory filling")
g = GNNGraph(collect(1:N-1), collect(2:N), num_nodes = N, ndata = rand(Float32, 1, N)) |> gpu
gnnchain = GNNChain(Dense(1, 1000), Dense(1000, 1)) |> gpu
CUDA.@time test_mem(gnnchain, g)
GC.gc(); CUDA.reclaim();
println("\n\nNN equivalent to GNN, memory filling")
x = g.ndata.x
chain = Chain(gnnchain.layers...)
@assert gnnchain(g).ndata.x ≈ chain(x)
## same results with these
# x = rand(1, N) |> gpu
# chain = Chain(Dense(1, 1000), Dense(1000, 1)) |> gpu
CUDA.@time test_mem(chain, x)
GC.gc(); CUDA.reclaim();
println("\n\nNN 1, same memory")
data = rand(N, 1) |> gpu
n = Chain(Dense(N, 1000), Dense(1000, 1)) |> gpu
CUDA.@time test_mem(n, data)
println("\n\nNN 2, memory filling:")
data = rand(N, N) |> gpu
n = Chain(Dense(N, 1000), Dense(1000, 1)) |> gpu
CUDA.@time test_mem(n, data)
julia> include("bench.jl")
GNN, memory filling
Effective GPU memory usage: 21.33% (834.750 MiB/3.823 GiB)
Memory pool usage: 496.399 MiB (544.000 MiB reserved)Effective GPU memory usage: 22.96% (898.750 MiB/3.823 GiB)
Memory pool usage: 572.769 MiB (608.000 MiB reserved)Effective GPU memory usage: 25.41% (994.750 MiB/3.823 GiB)
Memory pool usage: 649.139 MiB (704.000 MiB reserved)Effective GPU memory usage: 27.05% (1.034 GiB/3.823 GiB)
Memory pool usage: 725.510 MiB (768.000 MiB reserved)Effective GPU memory usage: 29.50% (1.128 GiB/3.823 GiB)
Memory pool usage: 801.880 MiB (864.000 MiB reserved)Effective GPU memory usage: 30.32% (1.159 GiB/3.823 GiB)
Memory pool usage: 878.250 MiB (896.000 MiB reserved)Effective GPU memory usage: 31.95% (1.221 GiB/3.823 GiB)
Memory pool usage: 954.620 MiB (960.000 MiB reserved)Effective GPU memory usage: 34.41% (1.315 GiB/3.823 GiB)
Memory pool usage: 1.007 GiB (1.031 GiB reserved)Effective GPU memory usage: 36.04% (1.378 GiB/3.823 GiB)
Memory pool usage: 1.081 GiB (1.094 GiB reserved)Effective GPU memory usage: 37.68% (1.440 GiB/3.823 GiB)
Memory pool usage: 1.156 GiB (1.156 GiB reserved) 0.028500 seconds (9.53 k CPU allocations: 522.793 KiB) (40 GPU allocations: 763.702 MiB, 10.16% memmgmt time)
NN equivalent to GNN, memory filling
Effective GPU memory usage: 25.41% (994.750 MiB/3.823 GiB)
Memory pool usage: 649.025 MiB (704.000 MiB reserved)Effective GPU memory usage: 27.05% (1.034 GiB/3.823 GiB)
Memory pool usage: 725.395 MiB (768.000 MiB reserved)Effective GPU memory usage: 29.50% (1.128 GiB/3.823 GiB)
Memory pool usage: 801.765 MiB (864.000 MiB reserved)Effective GPU memory usage: 30.32% (1.159 GiB/3.823 GiB)
Memory pool usage: 878.136 MiB (896.000 MiB reserved)Effective GPU memory usage: 31.95% (1.221 GiB/3.823 GiB)
Memory pool usage: 954.506 MiB (960.000 MiB reserved)Effective GPU memory usage: 34.41% (1.315 GiB/3.823 GiB)
Memory pool usage: 1.007 GiB (1.031 GiB reserved)Effective GPU memory usage: 36.04% (1.378 GiB/3.823 GiB)
Memory pool usage: 1.081 GiB (1.094 GiB reserved)Effective GPU memory usage: 37.68% (1.440 GiB/3.823 GiB)
Memory pool usage: 1.156 GiB (1.156 GiB reserved)Effective GPU memory usage: 40.13% (1.534 GiB/3.823 GiB)
Memory pool usage: 1.230 GiB (1.250 GiB reserved)Effective GPU memory usage: 41.76% (1.596 GiB/3.823 GiB)
Memory pool usage: 1.305 GiB (1.312 GiB reserved) 0.027846 seconds (9.17 k CPU allocations: 508.121 KiB) (40 GPU allocations: 763.702 MiB, 8.10% memmgmt time)
NN 1, same memory
Effective GPU memory usage: 20.51% (802.750 MiB/3.823 GiB)
Memory pool usage: 458.027 MiB (512.000 MiB reserved)Effective GPU memory usage: 20.51% (802.750 MiB/3.823 GiB)
Memory pool usage: 458.035 MiB (512.000 MiB reserved)Effective GPU memory usage: 20.51% (802.750 MiB/3.823 GiB)
Memory pool usage: 458.042 MiB (512.000 MiB reserved)Effective GPU memory usage: 20.51% (802.750 MiB/3.823 GiB)
Memory pool usage: 458.050 MiB (512.000 MiB reserved)Effective GPU memory usage: 20.51% (802.750 MiB/3.823 GiB)
Memory pool usage: 458.057 MiB (512.000 MiB reserved)Effective GPU memory usage: 20.51% (802.750 MiB/3.823 GiB)
Memory pool usage: 458.065 MiB (512.000 MiB reserved)Effective GPU memory usage: 20.51% (802.750 MiB/3.823 GiB)
Memory pool usage: 458.073 MiB (512.000 MiB reserved)Effective GPU memory usage: 20.51% (802.750 MiB/3.823 GiB)
Memory pool usage: 458.080 MiB (512.000 MiB reserved)Effective GPU memory usage: 20.51% (802.750 MiB/3.823 GiB)
Memory pool usage: 458.088 MiB (512.000 MiB reserved)Effective GPU memory usage: 20.51% (802.750 MiB/3.823 GiB)
Memory pool usage: 458.096 MiB (512.000 MiB reserved) 0.005307 seconds (1.64 k CPU allocations: 86.688 KiB) (40 GPU allocations: 78.203 KiB, 1.58% memmgmt time)
NN 2, memory filling:
Effective GPU memory usage: 31.95% (1.221 GiB/3.823 GiB)
Memory pool usage: 954.014 MiB (960.000 MiB reserved)Effective GPU memory usage: 34.41% (1.315 GiB/3.823 GiB)
Memory pool usage: 1.006 GiB (1.031 GiB reserved)Effective GPU memory usage: 36.04% (1.378 GiB/3.823 GiB)
Memory pool usage: 1.081 GiB (1.094 GiB reserved)Effective GPU memory usage: 37.68% (1.440 GiB/3.823 GiB)
Memory pool usage: 1.155 GiB (1.156 GiB reserved)Effective GPU memory usage: 40.13% (1.534 GiB/3.823 GiB)
Memory pool usage: 1.230 GiB (1.250 GiB reserved)Effective GPU memory usage: 41.76% (1.596 GiB/3.823 GiB)
Memory pool usage: 1.305 GiB (1.312 GiB reserved)Effective GPU memory usage: 44.22% (1.690 GiB/3.823 GiB)
Memory pool usage: 1.379 GiB (1.406 GiB reserved)Effective GPU memory usage: 45.85% (1.753 GiB/3.823 GiB)
Memory pool usage: 1.454 GiB (1.469 GiB reserved)Effective GPU memory usage: 47.49% (1.815 GiB/3.823 GiB)
Memory pool usage: 1.528 GiB (1.531 GiB reserved)Effective GPU memory usage: 49.94% (1.909 GiB/3.823 GiB)
Memory pool usage: 1.603 GiB (1.625 GiB reserved) 2.422488 seconds (1.74 k CPU allocations: 88.312 KiB) (40 GPU allocations: 763.702 MiB, 0.08% memmgmt time)
I see, my bad. I wrongly assumed the equivalent of the GNN to be x = rand(N,1)
while it should be x=rand(1,N)
When repeatedly performing inference operations, the GPU memory gets filled up quite fast. This causes the GPU to have to perform garbage collection. In my implementation, this accounted for 50% of GPU time as GC was needed every couple seconds. Using normal NN the GPU memory increases much slower.